fix(engine): shrink tries after clearing (#19159)

crates/engine/tree/src/tree/payload_processor/configured_sparse_trie.rs

@@ -186,4 +186,18 @@ impl SparseTrieInterface for ConfiguredSparseTrie {
             Self::Parallel(trie) => trie.value_capacity(),
         }
     }
+
+    fn shrink_nodes_to(&mut self, size: usize) {
+        match self {
+            Self::Serial(trie) => trie.shrink_nodes_to(size),
+            Self::Parallel(trie) => trie.shrink_nodes_to(size),
+        }
+    }
+
+    fn shrink_values_to(&mut self, size: usize) {
+        match self {
+            Self::Serial(trie) => trie.shrink_values_to(size),
+            Self::Parallel(trie) => trie.shrink_values_to(size),
+        }
+    }
 }

crates/engine/tree/src/tree/payload_processor/mod.rs

@@ -66,6 +66,29 @@ use configured_sparse_trie::ConfiguredSparseTrie;
 pub const PARALLEL_SPARSE_TRIE_PARALLELISM_THRESHOLDS: ParallelismThresholds =
     ParallelismThresholds { min_revealed_nodes: 100, min_updated_nodes: 100 };
 
+/// Default node capacity for shrinking the sparse trie. This is used to limit the number of trie
+/// nodes in allocated sparse tries.
+///
+/// Node maps have a key of `Nibbles` and value of `SparseNode`.
+/// The `size_of::<Nibbles>` is 40, and `size_of::<SparseNode>` is 80.
+///
+/// If we have 1 million entries of 120 bytes each, this conservative estimate comes out at around
+/// 120MB.
+pub const SPARSE_TRIE_MAX_NODES_SHRINK_CAPACITY: usize = 1_000_000;
+
+/// Default value capacity for shrinking the sparse trie. This is used to limit the number of values
+/// in allocated sparse tries.
+///
+/// There are storage and account values, the largest of the two being account values, which are
+/// essentially `TrieAccount`s.
+///
+/// Account value maps have a key of `Nibbles` and value of `TrieAccount`.
+/// The `size_of::<Nibbles>` is 40, and `size_of::<TrieAccount>` is 104.
+///
+/// If we have 1 million entries of 144 bytes each, this conservative estimate comes out at around
+/// 144MB.
+pub const SPARSE_TRIE_MAX_VALUES_SHRINK_CAPACITY: usize = 1_000_000;
+
 /// Entrypoint for executing the payload.
 #[derive(Debug)]
 pub struct PayloadProcessor<Evm>
@@ -439,11 +462,19 @@ where
             // Send state root computation result
             let _ = state_root_tx.send(result);
 
-            // Clear the SparseStateTrie and replace it back into the mutex _after_ sending
+            // Clear the SparseStateTrie, shrink, and replace it back into the mutex _after_ sending
             // results to the next step, so that time spent clearing doesn't block the step after
             // this one.
             let _enter = debug_span!(target: "engine::tree::payload_processor", "clear").entered();
-            cleared_sparse_trie.lock().replace(ClearedSparseStateTrie::from_state_trie(trie));
+            let mut cleared_trie = ClearedSparseStateTrie::from_state_trie(trie);
+
+            // Shrink the sparse trie so that we don't have ever increasing memory.
+            cleared_trie.shrink_to(
+                SPARSE_TRIE_MAX_NODES_SHRINK_CAPACITY,
+                SPARSE_TRIE_MAX_VALUES_SHRINK_CAPACITY,
+            );
+
+            cleared_sparse_trie.lock().replace(cleared_trie);
         });
     }
 }

crates/trie/sparse-parallel/src/lower.rs

@@ -122,4 +122,26 @@ impl LowerSparseSubtrie {
             Self::Blind(None) => 0,
         }
     }
+
+    /// Shrinks the capacity of the subtrie's node storage.
+    /// Works for both revealed and blind tries with allocated storage.
+    pub(crate) fn shrink_nodes_to(&mut self, size: usize) {
+        match self {
+            Self::Revealed(trie) | Self::Blind(Some(trie)) => {
+                trie.shrink_nodes_to(size);
+            }
+            Self::Blind(None) => {}
+        }
+    }
+
+    /// Shrinks the capacity of the subtrie's value storage.
+    /// Works for both revealed and blind tries with allocated storage.
+    pub(crate) fn shrink_values_to(&mut self, size: usize) {
+        match self {
+            Self::Revealed(trie) | Self::Blind(Some(trie)) => {
+                trie.shrink_values_to(size);
+            }
+            Self::Blind(None) => {}
+        }
+    }
 }

crates/trie/sparse-parallel/src/trie.rs

@@ -883,6 +883,42 @@ impl SparseTrieInterface for ParallelSparseTrie {
         self.upper_subtrie.value_capacity() +
             self.lower_subtries.iter().map(|trie| trie.value_capacity()).sum::<usize>()
     }
+
+    fn shrink_nodes_to(&mut self, size: usize) {
+        // Distribute the capacity across upper and lower subtries
+        //
+        // Always include upper subtrie, plus any lower subtries
+        let total_subtries = 1 + NUM_LOWER_SUBTRIES;
+        let size_per_subtrie = size / total_subtries;
+
+        // Shrink the upper subtrie
+        self.upper_subtrie.shrink_nodes_to(size_per_subtrie);
+
+        // Shrink lower subtries (works for both revealed and blind with allocation)
+        for subtrie in &mut self.lower_subtries {
+            subtrie.shrink_nodes_to(size_per_subtrie);
+        }
+
+        // shrink masks maps
+        self.branch_node_hash_masks.shrink_to(size);
+        self.branch_node_tree_masks.shrink_to(size);
+    }
+
+    fn shrink_values_to(&mut self, size: usize) {
+        // Distribute the capacity across upper and lower subtries
+        //
+        // Always include upper subtrie, plus any lower subtries
+        let total_subtries = 1 + NUM_LOWER_SUBTRIES;
+        let size_per_subtrie = size / total_subtries;
+
+        // Shrink the upper subtrie
+        self.upper_subtrie.shrink_values_to(size_per_subtrie);
+
+        // Shrink lower subtries (works for both revealed and blind with allocation)
+        for subtrie in &mut self.lower_subtries {
+            subtrie.shrink_values_to(size_per_subtrie);
+        }
+    }
 }
 
 impl ParallelSparseTrie {
@@ -2111,6 +2147,16 @@ impl SparseSubtrie {
     pub(crate) fn value_capacity(&self) -> usize {
         self.inner.value_capacity()
     }
+
+    /// Shrinks the capacity of the subtrie's node storage.
+    pub(crate) fn shrink_nodes_to(&mut self, size: usize) {
+        self.nodes.shrink_to(size);
+    }
+
+    /// Shrinks the capacity of the subtrie's value storage.
+    pub(crate) fn shrink_values_to(&mut self, size: usize) {
+        self.inner.values.shrink_to(size);
+    }
 }
 
 /// Helper type for [`SparseSubtrie`] to mutably access only a subset of fields from the original
@@ -2571,10 +2617,19 @@ impl SparseSubtrieBuffers {
     /// Clears all buffers.
     fn clear(&mut self) {
         self.path_stack.clear();
+        self.path_stack.shrink_to_fit();
+
         self.rlp_node_stack.clear();
+        self.rlp_node_stack.shrink_to_fit();
+
         self.branch_child_buf.clear();
+        self.branch_child_buf.shrink_to_fit();
+
         self.branch_value_stack_buf.clear();
+        self.branch_value_stack_buf.shrink_to_fit();
+
         self.rlp_buf.clear();
+        self.rlp_buf.shrink_to_fit();
     }
 }
 

crates/trie/sparse/src/state.rs

@@ -43,6 +43,32 @@ where
         Self(trie)
     }
 
+    /// Shrink the cleared sparse trie's capacity to the given node and value size.
+    /// This helps reduce memory usage when the trie has excess capacity.
+    /// The capacity is distributed equally across the account trie and all storage tries.
+    pub fn shrink_to(&mut self, node_size: usize, value_size: usize) {
+        // Count total number of storage tries (active + cleared + default)
+        let storage_tries_count = self.0.storage.tries.len() + self.0.storage.cleared_tries.len();
+
+        // Total tries = 1 account trie + all storage tries
+        let total_tries = 1 + storage_tries_count;
+
+        // Distribute capacity equally among all tries
+        let node_size_per_trie = node_size / total_tries;
+        let value_size_per_trie = value_size / total_tries;
+
+        // Shrink the account trie
+        self.0.state.shrink_nodes_to(node_size_per_trie);
+        self.0.state.shrink_values_to(value_size_per_trie);
+
+        // Give storage tries the remaining capacity after account trie allocation
+        let storage_node_size = node_size.saturating_sub(node_size_per_trie);
+        let storage_value_size = value_size.saturating_sub(value_size_per_trie);
+
+        // Shrink all storage tries (they will redistribute internally)
+        self.0.storage.shrink_to(storage_node_size, storage_value_size);
+    }
+
     /// Returns the cleared [`SparseStateTrie`], consuming this instance.
     pub fn into_inner(self) -> SparseStateTrie<A, S> {
         self.0
@@ -860,6 +886,31 @@ impl<S: SparseTrieInterface> StorageTries<S> {
             set
         }));
     }
+
+    /// Shrinks the capacity of all storage tries (active, cleared, and default) to the given sizes.
+    /// The capacity is distributed equally among all tries that have allocations.
+    fn shrink_to(&mut self, node_size: usize, value_size: usize) {
+        // Count total number of tries with capacity (active + cleared + default)
+        let active_count = self.tries.len();
+        let cleared_count = self.cleared_tries.len();
+        let total_tries = 1 + active_count + cleared_count;
+
+        // Distribute capacity equally among all tries
+        let node_size_per_trie = node_size / total_tries;
+        let value_size_per_trie = value_size / total_tries;
+
+        // Shrink active storage tries
+        for trie in self.tries.values_mut() {
+            trie.shrink_nodes_to(node_size_per_trie);
+            trie.shrink_values_to(value_size_per_trie);
+        }
+
+        // Shrink cleared storage tries
+        for trie in &mut self.cleared_tries {
+            trie.shrink_nodes_to(node_size_per_trie);
+            trie.shrink_values_to(value_size_per_trie);
+        }
+    }
 }
 
 impl<S: SparseTrieInterface + Clone> StorageTries<S> {

crates/trie/sparse/src/traits.rs

@@ -228,6 +228,14 @@ pub trait SparseTrieInterface: Sized + Debug + Send + Sync {
 
     /// This returns the capacity of any inner data structures which store leaf values.
     fn value_capacity(&self) -> usize;
+
+    /// Shrink the capacity of the sparse trie's node storage to the given size.
+    /// This will reduce memory usage if the current capacity is higher than the given size.
+    fn shrink_nodes_to(&mut self, size: usize);
+
+    /// Shrink the capacity of the sparse trie's value storage to the given size.
+    /// This will reduce memory usage if the current capacity is higher than the given size.
+    fn shrink_values_to(&mut self, size: usize);
 }
 
 /// Struct for passing around branch node mask information.

crates/trie/sparse/src/trie.rs

@@ -275,6 +275,28 @@ impl<T: SparseTrieInterface> SparseTrie<T> {
             _ => 0,
         }
     }
+
+    /// Shrinks the capacity of the sparse trie's node storage.
+    /// Works for both revealed and blind tries with allocated storage.
+    pub fn shrink_nodes_to(&mut self, size: usize) {
+        match self {
+            Self::Blind(Some(trie)) | Self::Revealed(trie) => {
+                trie.shrink_nodes_to(size);
+            }
+            _ => {}
+        }
+    }
+
+    /// Shrinks the capacity of the sparse trie's value storage.
+    /// Works for both revealed and blind tries with allocated storage.
+    pub fn shrink_values_to(&mut self, size: usize) {
+        match self {
+            Self::Blind(Some(trie)) | Self::Revealed(trie) => {
+                trie.shrink_values_to(size);
+            }
+            _ => {}
+        }
+    }
 }
 
 /// The representation of revealed sparse trie.
@@ -1088,6 +1110,16 @@ impl SparseTrieInterface for SerialSparseTrie {
     fn value_capacity(&self) -> usize {
         self.values.capacity()
     }
+
+    fn shrink_nodes_to(&mut self, size: usize) {
+        self.nodes.shrink_to(size);
+        self.branch_node_tree_masks.shrink_to(size);
+        self.branch_node_hash_masks.shrink_to(size);
+    }
+
+    fn shrink_values_to(&mut self, size: usize) {
+        self.values.shrink_to(size);
+    }
 }
 
 impl SerialSparseTrie {