feat(trie): TrieInput (paradigmxyz#10918)

NBnet · Sep 16, 2024 · b0eb78c · b0eb78c
1 parent f37bd47
commit b0eb78c
Show file tree

Hide file tree

Showing 14 changed files with 131 additions and 203 deletions.
diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs
@@ -839,7 +839,7 @@ mod tests {
         AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
         StorageRootProvider,
     };
-    use reth_trie::{prefix_set::TriePrefixSetsMut, AccountProof, HashedStorage};
+    use reth_trie::{AccountProof, HashedStorage, TrieInput};
 
     fn create_mock_state(
         test_block_builder: &mut TestBlockBuilder,
@@ -913,12 +913,7 @@ mod tests {
             Ok(B256::random())
         }
 
-        fn state_root_from_nodes(
-            &self,
-            _nodes: TrieUpdates,
-            _post_state: HashedPostState,
-            _prefix_sets: TriePrefixSetsMut,
-        ) -> ProviderResult<B256> {
+        fn state_root_from_nodes(&self, _input: TrieInput) -> ProviderResult<B256> {
             Ok(B256::random())
         }
 
@@ -931,9 +926,7 @@ mod tests {
 
         fn state_root_from_nodes_with_updates(
             &self,
-            _nodes: TrieUpdates,
-            _post_state: HashedPostState,
-            _prefix_sets: TriePrefixSetsMut,
+            _input: TrieInput,
         ) -> ProviderResult<(B256, TrieUpdates)> {
             Ok((B256::random(), TrieUpdates::default()))
         }

diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs
@@ -7,10 +7,7 @@ use reth_storage_api::{
     AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateProviderBox,
     StateRootProvider, StorageRootProvider,
 };
-use reth_trie::{
-    prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
-    HashedStorage,
-};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
 use std::{collections::HashMap, sync::OnceLock};
 
 /// A state provider that stores references to in-memory blocks along with their state as well as
@@ -45,13 +42,12 @@ impl MemoryOverlayStateProvider {
     /// Return lazy-loaded trie state aggregated from in-memory blocks.
     fn trie_state(&self) -> &MemoryOverlayTrieState {
         self.trie_state.get_or_init(|| {
-            let mut hashed_state = HashedPostState::default();
-            let mut trie_nodes = TrieUpdates::default();
+            let mut trie_state = MemoryOverlayTrieState::default();
             for block in self.in_memory.iter().rev() {
-                hashed_state.extend_ref(block.hashed_state.as_ref());
-                trie_nodes.extend_ref(block.trie.as_ref());
+                trie_state.state.extend_ref(block.hashed_state.as_ref());
+                trie_state.nodes.extend_ref(block.trie.as_ref());
             }
-            MemoryOverlayTrieState { trie_nodes, hashed_state }
+            trie_state
         })
     }
 }
@@ -102,54 +98,38 @@ impl AccountReader for MemoryOverlayStateProvider {
 }
 
 impl StateRootProvider for MemoryOverlayStateProvider {
-    fn state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
-        let prefix_sets = hashed_state.construct_prefix_sets();
-        self.state_root_from_nodes(TrieUpdates::default(), hashed_state, prefix_sets)
+    fn state_root(&self, state: HashedPostState) -> ProviderResult<B256> {
+        self.state_root_from_nodes(TrieInput::from_state(state))
     }
 
-    fn state_root_from_nodes(
-        &self,
-        nodes: TrieUpdates,
-        state: HashedPostState,
-        prefix_sets: TriePrefixSetsMut,
-    ) -> ProviderResult<B256> {
-        let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state().clone();
-        trie_nodes.extend(nodes);
-        hashed_state.extend(state);
-        self.historical.state_root_from_nodes(trie_nodes, hashed_state, prefix_sets)
+    fn state_root_from_nodes(&self, mut input: TrieInput) -> ProviderResult<B256> {
+        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+        input.prepend_cached(nodes, state);
+        self.historical.state_root_from_nodes(input)
     }
 
     fn state_root_with_updates(
         &self,
-        hashed_state: HashedPostState,
+        state: HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        let prefix_sets = hashed_state.construct_prefix_sets();
-        self.state_root_from_nodes_with_updates(TrieUpdates::default(), hashed_state, prefix_sets)
+        self.state_root_from_nodes_with_updates(TrieInput::from_state(state))
     }
 
     fn state_root_from_nodes_with_updates(
         &self,
-        nodes: TrieUpdates,
-        state: HashedPostState,
-        prefix_sets: TriePrefixSetsMut,
+        mut input: TrieInput,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state().clone();
-        trie_nodes.extend(nodes);
-        hashed_state.extend(state);
-        self.historical.state_root_from_nodes_with_updates(trie_nodes, hashed_state, prefix_sets)
+        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+        input.prepend_cached(nodes, state);
+        self.historical.state_root_from_nodes_with_updates(input)
     }
 }
 
 impl StorageRootProvider for MemoryOverlayStateProvider {
     // TODO: Currently this does not reuse available in-memory trie nodes.
     fn storage_root(&self, address: Address, storage: HashedStorage) -> ProviderResult<B256> {
-        let mut hashed_storage = self
-            .trie_state()
-            .hashed_state
-            .storages
-            .get(&keccak256(address))
-            .cloned()
-            .unwrap_or_default();
+        let mut hashed_storage =
+            self.trie_state().state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
         hashed_storage.extend(&storage);
         self.historical.storage_root(address, hashed_storage)
     }
@@ -163,7 +143,7 @@ impl StateProofProvider for MemoryOverlayStateProvider {
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
-        let mut hashed_state = self.trie_state().hashed_state.clone();
+        let mut hashed_state = self.trie_state().state.clone();
         hashed_state.extend(state);
         self.historical.proof(hashed_state, address, slots)
     }
@@ -174,7 +154,7 @@ impl StateProofProvider for MemoryOverlayStateProvider {
         overlay: HashedPostState,
         target: HashedPostState,
     ) -> ProviderResult<HashMap<B256, Bytes>> {
-        let mut hashed_state = self.trie_state().hashed_state.clone();
+        let mut hashed_state = self.trie_state().state.clone();
         hashed_state.extend(overlay);
         self.historical.witness(hashed_state, target)
     }
@@ -207,10 +187,10 @@ impl StateProvider for MemoryOverlayStateProvider {
 }
 
 /// The collection of data necessary for trie-related operations for [`MemoryOverlayStateProvider`].
-#[derive(Clone, Debug)]
+#[derive(Clone, Default, Debug)]
 pub(crate) struct MemoryOverlayTrieState {
     /// The collection of aggregated in-memory trie updates.
-    pub(crate) trie_nodes: TrieUpdates,
+    pub(crate) nodes: TrieUpdates,
     /// The collection of hashed state from in-memory blocks.
-    pub(crate) hashed_state: HashedPostState,
+    pub(crate) state: HashedPostState,
 }
diff --git a/crates/revm/src/test_utils.rs b/crates/revm/src/test_utils.rs
@@ -8,10 +8,7 @@ use reth_storage_api::{
     StorageRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{
-    prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
-    HashedStorage,
-};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
 
 /// Mock state for testing
 #[derive(Debug, Default, Clone, Eq, PartialEq)]
@@ -74,12 +71,7 @@ impl StateRootProvider for StateProviderTest {
         unimplemented!("state root computation is not supported")
     }
 
-    fn state_root_from_nodes(
-        &self,
-        _nodes: TrieUpdates,
-        _hashed_state: HashedPostState,
-        _prefix_sets: TriePrefixSetsMut,
-    ) -> ProviderResult<B256> {
+    fn state_root_from_nodes(&self, _input: TrieInput) -> ProviderResult<B256> {
         unimplemented!("state root computation is not supported")
     }
 
@@ -92,9 +84,7 @@ impl StateRootProvider for StateProviderTest {
 
     fn state_root_from_nodes_with_updates(
         &self,
-        _nodes: TrieUpdates,
-        _hashed_state: HashedPostState,
-        _prefix_sets: TriePrefixSetsMut,
+        _input: TrieInput,
     ) -> ProviderResult<(B256, TrieUpdates)> {
         unimplemented!("state root computation is not supported")
     }

diff --git a/crates/rpc/rpc-eth-types/src/cache/db.rs b/crates/rpc/rpc-eth-types/src/cache/db.rs
@@ -27,11 +27,9 @@ impl<'a> reth_storage_api::StateRootProvider for StateProviderTraitObjWrapper<'a
 
     fn state_root_from_nodes(
         &self,
-        nodes: reth_trie::updates::TrieUpdates,
-        hashed_state: reth_trie::HashedPostState,
-        prefix_sets: reth_trie::prefix_set::TriePrefixSetsMut,
+        input: reth_trie::TrieInput,
     ) -> reth_errors::ProviderResult<B256> {
-        self.0.state_root_from_nodes(nodes, hashed_state, prefix_sets)
+        self.0.state_root_from_nodes(input)
     }
 
     fn state_root_with_updates(
@@ -43,11 +41,9 @@ impl<'a> reth_storage_api::StateRootProvider for StateProviderTraitObjWrapper<'a
 
     fn state_root_from_nodes_with_updates(
         &self,
-        nodes: reth_trie::updates::TrieUpdates,
-        hashed_state: reth_trie::HashedPostState,
-        prefix_sets: reth_trie::prefix_set::TriePrefixSetsMut,
+        input: reth_trie::TrieInput,
     ) -> reth_errors::ProviderResult<(B256, reth_trie::updates::TrieUpdates)> {
-        self.0.state_root_from_nodes_with_updates(nodes, hashed_state, prefix_sets)
+        self.0.state_root_from_nodes_with_updates(input)
     }
 }
 

diff --git a/crates/storage/provider/src/providers/bundle_state_provider.rs b/crates/storage/provider/src/providers/bundle_state_provider.rs
@@ -4,10 +4,7 @@ use crate::{
 use reth_primitives::{Account, Address, BlockNumber, Bytecode, Bytes, B256};
 use reth_storage_api::{StateProofProvider, StorageRootProvider};
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{
-    prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
-    HashedStorage,
-};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
 use std::collections::HashMap;
 
 /// A state provider that resolves to data from either a wrapped [`crate::ExecutionOutcome`]
@@ -75,12 +72,7 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateRootProvider
         self.state_provider.state_root(state)
     }
 
-    fn state_root_from_nodes(
-        &self,
-        _nodes: TrieUpdates,
-        _hashed_state: HashedPostState,
-        _prefix_sets: TriePrefixSetsMut,
-    ) -> ProviderResult<B256> {
+    fn state_root_from_nodes(&self, _input: TrieInput) -> ProviderResult<B256> {
         unimplemented!()
     }
 
@@ -96,16 +88,11 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateRootProvider
 
     fn state_root_from_nodes_with_updates(
         &self,
-        nodes: TrieUpdates,
-        hashed_state: HashedPostState,
-        prefix_sets: TriePrefixSetsMut,
+        mut input: TrieInput,
     ) -> ProviderResult<(B256, TrieUpdates)> {
         let bundle_state = self.block_execution_data_provider.execution_outcome().state();
-        let mut state = HashedPostState::from_bundle_state(&bundle_state.state);
-        let mut state_prefix_sets = state.construct_prefix_sets();
-        state.extend(hashed_state);
-        state_prefix_sets.extend(prefix_sets);
-        self.state_provider.state_root_from_nodes_with_updates(nodes, state, state_prefix_sets)
+        input.prepend(HashedPostState::from_bundle_state(&bundle_state.state));
+        self.state_provider.state_root_from_nodes_with_updates(input)
     }
 }
 

diff --git a/crates/storage/provider/src/providers/state/historical.rs b/crates/storage/provider/src/providers/state/historical.rs
@@ -16,8 +16,8 @@ use reth_primitives::{
 use reth_storage_api::{StateProofProvider, StorageRootProvider};
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
-    prefix_set::TriePrefixSetsMut, proof::Proof, updates::TrieUpdates, witness::TrieWitness,
-    AccountProof, HashedPostState, HashedStorage, StateRoot, StorageRoot,
+    proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
+    HashedStorage, StateRoot, StorageRoot, TrieInput,
 };
 use reth_trie_db::{
     DatabaseHashedPostState, DatabaseHashedStorage, DatabaseProof, DatabaseStateRoot,
@@ -292,17 +292,9 @@ impl<'b, TX: DbTx> StateRootProvider for HistoricalStateProviderRef<'b, TX> {
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
-    fn state_root_from_nodes(
-        &self,
-        nodes: TrieUpdates,
-        hashed_state: HashedPostState,
-        prefix_sets: TriePrefixSetsMut,
-    ) -> ProviderResult<B256> {
-        let mut revert_state = self.revert_state()?;
-        let mut revert_prefix_sets = revert_state.construct_prefix_sets();
-        revert_state.extend(hashed_state);
-        revert_prefix_sets.extend(prefix_sets);
-        StateRoot::overlay_root_from_nodes(self.tx, nodes, revert_state, revert_prefix_sets)
+    fn state_root_from_nodes(&self, mut input: TrieInput) -> ProviderResult<B256> {
+        input.prepend(self.revert_state()?);
+        StateRoot::overlay_root_from_nodes(self.tx, input)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
@@ -318,21 +310,11 @@ impl<'b, TX: DbTx> StateRootProvider for HistoricalStateProviderRef<'b, TX> {
 
     fn state_root_from_nodes_with_updates(
         &self,
-        nodes: TrieUpdates,
-        hashed_state: HashedPostState,
-        prefix_sets: TriePrefixSetsMut,
+        mut input: TrieInput,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        let mut revert_state = self.revert_state()?;
-        let mut revert_prefix_sets = revert_state.construct_prefix_sets();
-        revert_state.extend(hashed_state);
-        revert_prefix_sets.extend(prefix_sets);
-        StateRoot::overlay_root_from_nodes_with_updates(
-            self.tx,
-            nodes,
-            revert_state,
-            revert_prefix_sets,
-        )
-        .map_err(|err| ProviderError::Database(err.into()))
+        input.prepend(self.revert_state()?);
+        StateRoot::overlay_root_from_nodes_with_updates(self.tx, input)
+            .map_err(|err| ProviderError::Database(err.into()))
     }
 }
 

diff --git a/crates/storage/provider/src/providers/state/latest.rs b/crates/storage/provider/src/providers/state/latest.rs
@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use crate::{
     providers::{state::macros::delegate_provider_impls, StaticFileProvider},
     AccountReader, BlockHashReader, StateProvider, StateRootProvider,
@@ -16,10 +14,11 @@ use reth_primitives::{
 use reth_storage_api::{StateProofProvider, StorageRootProvider};
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{
-    prefix_set::TriePrefixSetsMut, proof::Proof, updates::TrieUpdates, witness::TrieWitness,
-    AccountProof, HashedPostState, HashedStorage, StateRoot, StorageRoot,
+    proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
+    HashedStorage, StateRoot, StorageRoot, TrieInput,
 };
 use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieWitness};
+use std::collections::HashMap;
 
 /// State provider over latest state that takes tx reference.
 #[derive(Debug)]
@@ -86,13 +85,8 @@ impl<'b, TX: DbTx> StateRootProvider for LatestStateProviderRef<'b, TX> {
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
-    fn state_root_from_nodes(
-        &self,
-        nodes: TrieUpdates,
-        hashed_state: HashedPostState,
-        prefix_sets: TriePrefixSetsMut,
-    ) -> ProviderResult<B256> {
-        StateRoot::overlay_root_from_nodes(self.tx, nodes, hashed_state, prefix_sets)
+    fn state_root_from_nodes(&self, input: TrieInput) -> ProviderResult<B256> {
+        StateRoot::overlay_root_from_nodes(self.tx, input)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
@@ -106,11 +100,9 @@ impl<'b, TX: DbTx> StateRootProvider for LatestStateProviderRef<'b, TX> {
 
     fn state_root_from_nodes_with_updates(
         &self,
-        nodes: TrieUpdates,
-        hashed_state: HashedPostState,
-        prefix_sets: TriePrefixSetsMut,
+        input: TrieInput,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        StateRoot::overlay_root_from_nodes_with_updates(self.tx, nodes, hashed_state, prefix_sets)
+        StateRoot::overlay_root_from_nodes_with_updates(self.tx, input)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 }

diff --git a/crates/storage/provider/src/providers/state/macros.rs b/crates/storage/provider/src/providers/state/macros.rs
@@ -43,9 +43,9 @@ macro_rules! delegate_provider_impls {
             }
             StateRootProvider $(where [$($generics)*])? {
                 fn state_root(&self, state: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;
-                fn state_root_from_nodes(&self,  nodes: reth_trie::updates::TrieUpdates, state: reth_trie::HashedPostState, prefix_sets: reth_trie::prefix_set::TriePrefixSetsMut) -> reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;
+                fn state_root_from_nodes(&self, input: reth_trie::TrieInput) -> reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;
                 fn state_root_with_updates(&self, state: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult<(reth_primitives::B256, reth_trie::updates::TrieUpdates)>;
-                fn state_root_from_nodes_with_updates(&self,  nodes: reth_trie::updates::TrieUpdates, state: reth_trie::HashedPostState, prefix_sets: reth_trie::prefix_set::TriePrefixSetsMut) -> reth_storage_errors::provider::ProviderResult<(reth_primitives::B256, reth_trie::updates::TrieUpdates)>;
+                fn state_root_from_nodes_with_updates(&self, input: reth_trie::TrieInput) -> reth_storage_errors::provider::ProviderResult<(reth_primitives::B256, reth_trie::updates::TrieUpdates)>;
             }
             StorageRootProvider $(where [$($generics)*])? {
                 fn storage_root(&self, address: reth_primitives::Address, storage: reth_trie::HashedStorage) ->  reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;