chore: refactor to use dedicated SparseTrieState struct and new variant

Author: Rjected

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

@@ -28,7 +28,7 @@ use reth_trie_parallel::{
     proof_task::{ProofTaskCtx, ProofTaskManager},
     root::ParallelStateRootError,
 };
-use reth_trie_sparse::{blinded::DefaultBlindedProvider, SparseTrie};
+use reth_trie_sparse::SparseTrieState;
 use std::{
     collections::VecDeque,
     sync::{
@@ -70,7 +70,7 @@ where
     precompile_cache_map: PrecompileCacheMap<SpecFor<Evm>>,
     /// A sparse trie, kept around to be used for the state root computation so that allocations
     /// can be minimized.
-    sparse_trie: Option<SparseTrie<DefaultBlindedProvider>>,
+    sparse_trie: Option<SparseTrieState>,
     _marker: std::marker::PhantomData<N>,
 }
 
@@ -251,7 +251,7 @@ where
     }
 
     /// Sets the sparse trie to be kept around for the state root computation.
-    pub(super) fn set_sparse_trie(&mut self, sparse_trie: SparseTrie<DefaultBlindedProvider>) {
+    pub(super) fn set_sparse_trie(&mut self, sparse_trie: SparseTrieState) {
         self.sparse_trie = Some(sparse_trie);
     }
 

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

@@ -9,9 +9,9 @@ use rayon::iter::{ParallelBridge, ParallelIterator};
 use reth_trie::{updates::TrieUpdates, Nibbles};
 use reth_trie_parallel::root::ParallelStateRootError;
 use reth_trie_sparse::{
-    blinded::{BlindedProvider, BlindedProviderFactory, DefaultBlindedProvider},
+    blinded::{BlindedProvider, BlindedProviderFactory},
     errors::{SparseStateTrieResult, SparseTrieErrorKind},
-    SparseStateTrie, SparseTrie,
+    SparseStateTrie, SparseTrieState,
 };
 use std::{
     sync::mpsc,
@@ -63,41 +63,39 @@ where
         }
     }
 
-    /// Creates a new sparse trie, populating the accounts trie with the given cleared `SparseTrie`
-    /// if it exists.
+    /// Creates a new sparse trie, populating the accounts trie with the given cleared
+    /// `SparseTrieState` if it exists.
     pub(super) fn new_with_stored_trie(
         executor: WorkloadExecutor,
         updates: mpsc::Receiver<SparseTrieUpdate>,
         blinded_provider_factory: BPF,
         trie_metrics: MultiProofTaskMetrics,
-        sparse_trie: Option<SparseTrie<DefaultBlindedProvider>>,
+        sparse_trie_state: Option<SparseTrieState>,
     ) -> Self {
-        if let Some(sparse_trie) = sparse_trie {
+        if let Some(sparse_trie_state) = sparse_trie_state {
             Self::with_accounts_trie(
                 executor,
                 updates,
                 blinded_provider_factory,
                 trie_metrics,
-                sparse_trie,
+                sparse_trie_state,
             )
         } else {
             Self::new(executor, updates, blinded_provider_factory, trie_metrics)
         }
     }
 
-    /// Creates a new sparse trie task, using the given cleared `SparseTrie` for the accounts trie.
+    /// Creates a new sparse trie task, using the given cleared `SparseTrieState` for the accounts
+    /// trie.
     pub(super) fn with_accounts_trie(
         executor: WorkloadExecutor,
         updates: mpsc::Receiver<SparseTrieUpdate>,
         blinded_provider_factory: BPF,
         metrics: MultiProofTaskMetrics,
-        cleared_accounts_trie: SparseTrie<DefaultBlindedProvider>,
+        sparse_trie_state: SparseTrieState,
     ) -> Self {
         let mut trie = SparseStateTrie::new(blinded_provider_factory).with_updates(true);
-
-        if let SparseTrie::Revealed(cleared) = cleared_accounts_trie {
-            trie.populate_from(cleared);
-        }
+        trie.populate_from(sparse_trie_state);
 
         Self { executor, updates, metrics, trie }
     }
@@ -149,7 +147,7 @@ where
         self.metrics.sparse_trie_total_duration_histogram.record(now.elapsed());
 
         // take the account trie
-        let trie = self.trie.take_cleared_account_trie();
+        let trie = self.trie.take_cleared_account_trie_state();
 
         Ok(StateRootComputeOutcome { state_root, trie_updates, trie })
     }
@@ -164,7 +162,7 @@ pub struct StateRootComputeOutcome {
     /// The trie updates.
     pub trie_updates: TrieUpdates,
     /// The account state trie.
-    pub trie: SparseTrie<DefaultBlindedProvider>,
+    pub trie: SparseTrieState,
 }
 
 /// Updates the sparse trie with the given proofs and state, and returns the elapsed time.

crates/trie/sparse/src/state.rs

@@ -1,11 +1,8 @@
 use crate::{
-    blinded::{
-        BlindedProvider, BlindedProviderFactory, DefaultBlindedProvider,
-        DefaultBlindedProviderFactory,
-    },
-    LeafLookup, RevealedSparseTrie, SparseTrie, TrieMasks,
+    blinded::{BlindedProvider, BlindedProviderFactory, DefaultBlindedProviderFactory},
+    LeafLookup, RevealedSparseTrie, SparseTrie, SparseTrieState, TrieMasks,
 };
-use alloc::{boxed::Box, collections::VecDeque, vec::Vec};
+use alloc::{collections::VecDeque, vec::Vec};
 use alloy_primitives::{
     hex,
     map::{B256Map, HashMap, HashSet},
@@ -110,15 +107,16 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         self.revealed_account_paths.contains(&Nibbles::unpack(account))
     }
 
-    /// Uses the input `SparseTrie` to populate the backing data structures in the `state` trie.
-    pub fn populate_from<P>(&mut self, trie: Box<RevealedSparseTrie<P>>) {
+    /// Uses the input `SparseTrieState` to populate the backing data structures in the `state`
+    /// trie.
+    pub fn populate_from(&mut self, trie: SparseTrieState) {
         if let Some(new_trie) = self.state.as_revealed_mut() {
-            new_trie.populate_from_trie(trie);
+            new_trie.use_allocated_state(trie);
         } else {
             self.state = SparseTrie::revealed_with_provider(
                 self.provider_factory.account_node_provider(),
                 trie,
-            );
+            )
         }
     }
 
@@ -662,7 +660,7 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         &mut self,
     ) -> SparseStateTrieResult<&mut RevealedSparseTrie<F::AccountNodeProvider>> {
         match self.state {
-            SparseTrie::Blind { allocated: _ } => {
+            SparseTrie::Blind | SparseTrie::AllocatedEmpty { .. } => {
                 let (root_node, hash_mask, tree_mask) = self
                     .provider_factory
                     .account_node_provider()
@@ -882,7 +880,7 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
     }
 
     /// Clears and takes the account trie.
-    pub fn take_cleared_account_trie(&mut self) -> SparseTrie<DefaultBlindedProvider> {
+    pub fn take_cleared_account_trie_state(&mut self) -> SparseTrieState {
         let trie = core::mem::take(&mut self.state);
         trie.cleared()
     }
@@ -988,7 +986,7 @@ mod tests {
         assert_eq!(proofs.len(), 1);
 
         let mut sparse = SparseStateTrie::default();
-        assert_eq!(sparse.state, SparseTrie::Blind { allocated: None });
+        assert_eq!(sparse.state, SparseTrie::Blind);
 
         sparse.reveal_account(Default::default(), proofs.into_inner()).unwrap();
         assert_eq!(sparse.state, SparseTrie::revealed_empty());