feat(trie): StorageRootProvider::storage_multiproof (#12879)

paradigmxyz · Nov 26, 2024 · a0f99df · a0f99df
1 parent b6d6bf5
commit a0f99df
Show file tree

Hide file tree

Showing 12 changed files with 149 additions and 11 deletions.
diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs
@@ -944,7 +944,9 @@ mod tests {
         AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
         StorageRootProvider,
     };
-    use reth_trie::{AccountProof, HashedStorage, MultiProof, StorageProof, TrieInput};
+    use reth_trie::{
+        AccountProof, HashedStorage, MultiProof, StorageMultiProof, StorageProof, TrieInput,
+    };
 
     fn create_mock_state(
         test_block_builder: &mut TestBlockBuilder<EthPrimitives>,
@@ -1054,6 +1056,15 @@ mod tests {
         ) -> ProviderResult<StorageProof> {
             Ok(StorageProof::new(slot))
         }
+
+        fn storage_multiproof(
+            &self,
+            _address: Address,
+            _slots: &[B256],
+            _hashed_storage: HashedStorage,
+        ) -> ProviderResult<StorageMultiProof> {
+            Ok(StorageMultiProof::empty())
+        }
     }
 
     impl StateProofProvider for MockStateProvider {

diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs
@@ -11,7 +11,8 @@ use reth_storage_api::{
     StorageRootProvider,
 };
 use reth_trie::{
-    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
+    StorageMultiProof, TrieInput,
 };
 use std::sync::OnceLock;
 
@@ -167,6 +168,20 @@ macro_rules! impl_state_provider {
                 hashed_storage.extend(&storage);
                 self.historical.storage_proof(address, slot, hashed_storage)
             }
+
+             // TODO: Currently this does not reuse available in-memory trie nodes.
+             fn storage_multiproof(
+                &self,
+                address: Address,
+                slots: &[B256],
+                storage: HashedStorage,
+            ) -> ProviderResult<StorageMultiProof> {
+                let state = &self.trie_state().state;
+                let mut hashed_storage =
+                    state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
+                hashed_storage.extend(&storage);
+                self.historical.storage_multiproof(address, slots, hashed_storage)
+            }
         }
 
         impl $($tokens)* StateProofProvider for $type {

diff --git a/crates/revm/src/test_utils.rs b/crates/revm/src/test_utils.rs
@@ -11,8 +11,8 @@ use reth_storage_api::{
 };
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
-    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, StorageProof,
-    TrieInput,
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
+    StorageMultiProof, StorageProof, TrieInput,
 };
 
 /// Mock state for testing
@@ -112,6 +112,15 @@ impl StorageRootProvider for StateProviderTest {
     ) -> ProviderResult<StorageProof> {
         unimplemented!("proof generation is not supported")
     }
+
+    fn storage_multiproof(
+        &self,
+        _address: Address,
+        _slots: &[B256],
+        _hashed_storage: HashedStorage,
+    ) -> ProviderResult<StorageMultiProof> {
+        unimplemented!("proof generation is not supported")
+    }
 }
 
 impl StateProofProvider for StateProviderTest {

diff --git a/crates/rpc/rpc-eth-types/src/cache/db.rs b/crates/rpc/rpc-eth-types/src/cache/db.rs
@@ -67,6 +67,15 @@ impl reth_storage_api::StorageRootProvider for StateProviderTraitObjWrapper<'_>
     ) -> ProviderResult<reth_trie::StorageProof> {
         self.0.storage_proof(address, slot, hashed_storage)
     }
+
+    fn storage_multiproof(
+        &self,
+        address: Address,
+        slots: &[B256],
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<reth_trie::StorageMultiProof> {
+        self.0.storage_multiproof(address, slots, hashed_storage)
+    }
 }
 
 impl reth_storage_api::StateProofProvider for StateProviderTraitObjWrapper<'_> {

diff --git a/crates/storage/provider/src/providers/bundle_state_provider.rs b/crates/storage/provider/src/providers/bundle_state_provider.rs
@@ -9,7 +9,8 @@ use reth_primitives::{Account, Bytecode};
 use reth_storage_api::{StateProofProvider, StorageRootProvider};
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
-    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
+    StorageMultiProof, TrieInput,
 };
 
 /// A state provider that resolves to data from either a wrapped [`crate::ExecutionOutcome`]
@@ -138,6 +139,17 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StorageRootProvider
         storage.extend(&hashed_storage);
         self.state_provider.storage_proof(address, slot, storage)
     }
+
+    fn storage_multiproof(
+        &self,
+        address: Address,
+        slots: &[B256],
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<StorageMultiProof> {
+        let mut storage = self.get_hashed_storage(address);
+        storage.extend(&hashed_storage);
+        self.state_provider.storage_multiproof(address, slots, storage)
+    }
 }
 
 impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProofProvider

diff --git a/crates/storage/provider/src/providers/state/historical.rs b/crates/storage/provider/src/providers/state/historical.rs
@@ -21,7 +21,8 @@ use reth_trie::{
     proof::{Proof, StorageProof},
     updates::TrieUpdates,
     witness::TrieWitness,
-    AccountProof, HashedPostState, HashedStorage, MultiProof, StateRoot, StorageRoot, TrieInput,
+    AccountProof, HashedPostState, HashedStorage, MultiProof, StateRoot, StorageMultiProof,
+    StorageRoot, TrieInput,
 };
 use reth_trie_db::{
     DatabaseHashedPostState, DatabaseHashedStorage, DatabaseProof, DatabaseStateRoot,
@@ -341,6 +342,18 @@ impl<Provider: DBProvider + BlockNumReader> StorageRootProvider
         StorageProof::overlay_storage_proof(self.tx(), address, slot, revert_storage)
             .map_err(Into::<ProviderError>::into)
     }
+
+    fn storage_multiproof(
+        &self,
+        address: Address,
+        slots: &[B256],
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<StorageMultiProof> {
+        let mut revert_storage = self.revert_storage(address)?;
+        revert_storage.extend(&hashed_storage);
+        StorageProof::overlay_storage_multiproof(self.tx(), address, slots, revert_storage)
+            .map_err(Into::<ProviderError>::into)
+    }
 }
 
 impl<Provider: DBProvider + BlockNumReader> StateProofProvider

diff --git a/crates/storage/provider/src/providers/state/latest.rs b/crates/storage/provider/src/providers/state/latest.rs
@@ -15,7 +15,8 @@ use reth_trie::{
     proof::{Proof, StorageProof},
     updates::TrieUpdates,
     witness::TrieWitness,
-    AccountProof, HashedPostState, HashedStorage, MultiProof, StateRoot, StorageRoot, TrieInput,
+    AccountProof, HashedPostState, HashedStorage, MultiProof, StateRoot, StorageMultiProof,
+    StorageRoot, TrieInput,
 };
 use reth_trie_db::{
     DatabaseProof, DatabaseStateRoot, DatabaseStorageProof, DatabaseStorageRoot,
@@ -108,6 +109,16 @@ impl<Provider: DBProvider> StorageRootProvider for LatestStateProviderRef<'_, Pr
         StorageProof::overlay_storage_proof(self.tx(), address, slot, hashed_storage)
             .map_err(Into::<ProviderError>::into)
     }
+
+    fn storage_multiproof(
+        &self,
+        address: Address,
+        slots: &[B256],
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<StorageMultiProof> {
+        StorageProof::overlay_storage_multiproof(self.tx(), address, slots, hashed_storage)
+            .map_err(Into::<ProviderError>::into)
+    }
 }
 
 impl<Provider: DBProvider> StateProofProvider for LatestStateProviderRef<'_, Provider> {

diff --git a/crates/storage/provider/src/providers/state/macros.rs b/crates/storage/provider/src/providers/state/macros.rs
@@ -50,6 +50,7 @@ macro_rules! delegate_provider_impls {
             StorageRootProvider $(where [$($generics)*])? {
                 fn storage_root(&self, address: alloy_primitives::Address, storage: reth_trie::HashedStorage) -> reth_storage_errors::provider::ProviderResult<alloy_primitives::B256>;
                 fn storage_proof(&self, address: alloy_primitives::Address, slot: alloy_primitives::B256, storage: reth_trie::HashedStorage) -> reth_storage_errors::provider::ProviderResult<reth_trie::StorageProof>;
+                fn storage_multiproof(&self, address: alloy_primitives::Address, slots: &[alloy_primitives::B256], storage: reth_trie::HashedStorage) -> reth_storage_errors::provider::ProviderResult<reth_trie::StorageMultiProof>;
             }
             StateProofProvider $(where [$($generics)*])? {
                 fn proof(&self, input: reth_trie::TrieInput, address: alloy_primitives::Address, slots: &[alloy_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;

diff --git a/crates/storage/provider/src/test_utils/mock.rs b/crates/storage/provider/src/test_utils/mock.rs
@@ -32,8 +32,8 @@ use reth_storage_api::{
 };
 use reth_storage_errors::provider::{ConsistentViewError, ProviderError, ProviderResult};
 use reth_trie::{
-    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, StorageProof,
-    TrieInput,
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
+    StorageMultiProof, StorageProof, TrieInput,
 };
 use reth_trie_db::MerklePatriciaTrie;
 use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg};
@@ -639,6 +639,15 @@ impl StorageRootProvider for MockEthProvider {
     ) -> ProviderResult<reth_trie::StorageProof> {
         Ok(StorageProof::new(slot))
     }
+
+    fn storage_multiproof(
+        &self,
+        _address: Address,
+        _slots: &[B256],
+        _hashed_storage: HashedStorage,
+    ) -> ProviderResult<StorageMultiProof> {
+        Ok(StorageMultiProof::empty())
+    }
 }
 
 impl StateProofProvider for MockEthProvider {

diff --git a/crates/storage/provider/src/test_utils/noop.rs b/crates/storage/provider/src/test_utils/noop.rs
@@ -372,6 +372,15 @@ impl StorageRootProvider for NoopProvider {
     ) -> ProviderResult<reth_trie::StorageProof> {
         Ok(reth_trie::StorageProof::new(slot))
     }
+
+    fn storage_multiproof(
+        &self,
+        _address: Address,
+        _slots: &[B256],
+        _hashed_storage: HashedStorage,
+    ) -> ProviderResult<reth_trie::StorageMultiProof> {
+        Ok(reth_trie::StorageMultiProof::empty())
+    }
 }
 
 impl StateProofProvider for NoopProvider {

diff --git a/crates/storage/storage-api/src/trie.rs b/crates/storage/storage-api/src/trie.rs
@@ -5,7 +5,8 @@ use alloy_primitives::{
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
     updates::{StorageTrieUpdates, TrieUpdates},
-    AccountProof, HashedPostState, HashedStorage, MultiProof, StorageProof, TrieInput,
+    AccountProof, HashedPostState, HashedStorage, MultiProof, StorageMultiProof, StorageProof,
+    TrieInput,
 };
 
 /// A type that can compute the state root of a given post state.
@@ -56,6 +57,14 @@ pub trait StorageRootProvider: Send + Sync {
         slot: B256,
         hashed_storage: HashedStorage,
     ) -> ProviderResult<StorageProof>;
+
+    /// Returns the storage multiproof for target slots.
+    fn storage_multiproof(
+        &self,
+        address: Address,
+        slots: &[B256],
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<StorageMultiProof>;
 }
 
 /// A type that can generate state proof on top of a given post state.

diff --git a/crates/trie/db/src/proof.rs b/crates/trie/db/src/proof.rs
@@ -10,7 +10,7 @@ use reth_trie::{
     hashed_cursor::HashedPostStateCursorFactory,
     proof::{Proof, StorageProof},
     trie_cursor::InMemoryTrieCursorFactory,
-    AccountProof, HashedPostStateSorted, HashedStorage, MultiProof, TrieInput,
+    AccountProof, HashedPostStateSorted, HashedStorage, MultiProof, StorageMultiProof, TrieInput,
 };
 
 /// Extends [`Proof`] with operations specific for working with a database transaction.
@@ -96,6 +96,14 @@ pub trait DatabaseStorageProof<'a, TX> {
         slot: B256,
         storage: HashedStorage,
     ) -> Result<reth_trie::StorageProof, StateProofError>;
+
+    /// Generates the storage multiproof for target slots based on [`TrieInput`].
+    fn overlay_storage_multiproof(
+        tx: &'a TX,
+        address: Address,
+        slots: &[B256],
+        storage: HashedStorage,
+    ) -> Result<StorageMultiProof, StateProofError>;
 }
 
 impl<'a, TX: DbTx> DatabaseStorageProof<'a, TX>
@@ -125,4 +133,26 @@ impl<'a, TX: DbTx> DatabaseStorageProof<'a, TX>
             .with_prefix_set_mut(prefix_set)
             .storage_proof(slot)
     }
+
+    fn overlay_storage_multiproof(
+        tx: &'a TX,
+        address: Address,
+        slots: &[B256],
+        storage: HashedStorage,
+    ) -> Result<StorageMultiProof, StateProofError> {
+        let hashed_address = keccak256(address);
+        let targets = slots.iter().map(keccak256).collect();
+        let prefix_set = storage.construct_prefix_set();
+        let state_sorted = HashedPostStateSorted::new(
+            Default::default(),
+            HashMap::from([(hashed_address, storage.into_sorted())]),
+        );
+        Self::from_tx(tx, address)
+            .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
+                DatabaseHashedCursorFactory::new(tx),
+                &state_sorted,
+            ))
+            .with_prefix_set_mut(prefix_set)
+            .storage_multiproof(targets)
+    }
 }