feat: implement new trie_walker which can iterate full+partial tries

portal-bridge/src/bridge/state.rs

@@ -25,9 +25,9 @@ use trin_execution::{
         create_contract_content_value, create_storage_content_key, create_storage_content_value,
     },
     execution::TrinExecution,
-    trie_walker::TrieWalker,
     types::{block_to_trace::BlockToTrace, trie_proof::TrieProof},
     utils::full_nibble_path_to_address_hash,
+    walkers::{memory_db::ReadOnlyMemoryDB, trie_walker::TrieWalker},
 };
 use trin_metrics::bridge::BridgeMetricsReporter;
 use trin_utils::dir::create_temp_dir;
@@ -159,13 +159,14 @@ impl StateBridge {
             .header
             .hash();
 
-        let walk_diff = TrieWalker::new(root_with_trie_diff.root, root_with_trie_diff.trie_diff);
+        let walk_diff = TrieWalker::new_partial_trie(
+            root_with_trie_diff.root,
+            ReadOnlyMemoryDB::new(root_with_trie_diff.trie_diff),
+        )?;
 
         // gossip block's new state transitions
         let mut content_idx = 0;
-        for node in walk_diff.nodes.keys() {
-            let account_proof = walk_diff.get_proof(*node);
-
+        for account_proof in walk_diff {
             // gossip the account
             self.gossip_account(&account_proof, block_hash, content_idx)
                 .await?;
@@ -213,10 +214,12 @@ impl StateBridge {
             // gossip contract storage
             let storage_changed_nodes = trin_execution.database.get_storage_trie_diff(address_hash);
 
-            let storage_walk_diff = TrieWalker::new(account.storage_root, storage_changed_nodes);
+            let storage_walk_diff = TrieWalker::new_partial_trie(
+                account.storage_root,
+                ReadOnlyMemoryDB::new(storage_changed_nodes),
+            )?;
 
-            for storage_node in storage_walk_diff.nodes.keys() {
-                let storage_proof = storage_walk_diff.get_proof(*storage_node);
+            for storage_proof in storage_walk_diff {
                 self.gossip_storage(
                     &account_proof,
                     &storage_proof,

trin-execution/src/lib.rs

@@ -8,6 +8,6 @@ pub mod execution;
 pub mod metrics;
 pub mod storage;
 pub mod subcommands;
-pub mod trie_walker;
 pub mod types;
 pub mod utils;
+pub mod walkers;

trin-execution/src/walkers/memory_db.rs

@@ -0,0 +1,47 @@
+use std::sync::Arc;
+
+use anyhow::anyhow;
+use eth_trie::DB;
+use hashbrown::HashMap;
+use parking_lot::Mutex;
+use revm_primitives::B256;
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum ReadOnlyMemoryDBError {
+    #[error("read only memory db error {0}")]
+    ANYHOW(#[from] anyhow::Error),
+}
+
+#[derive(Debug)]
+pub struct ReadOnlyMemoryDB {
+    storage: Arc<Mutex<HashMap<B256, Vec<u8>>>>,
+}
+
+impl ReadOnlyMemoryDB {
+    pub fn new(storage: HashMap<B256, Vec<u8>>) -> Self {
+        ReadOnlyMemoryDB {
+            storage: Arc::new(Mutex::new(storage)),
+        }
+    }
+}
+
+impl DB for ReadOnlyMemoryDB {
+    type Error = ReadOnlyMemoryDBError;
+
+    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+        Ok(self.storage.lock().get(key).cloned())
+    }
+
+    fn insert(&self, _key: &[u8], _value: Vec<u8>) -> Result<(), Self::Error> {
+        Err(anyhow!("Cannot insert into read only memory db").into())
+    }
+
+    fn remove(&self, _key: &[u8]) -> Result<(), Self::Error> {
+        Err(anyhow!("Cannot remove from read only memory db").into())
+    }
+
+    fn flush(&self) -> Result<(), Self::Error> {
+        Err(anyhow!("Cannot flush read only memory db").into())
+    }
+}

trin-execution/src/walkers/mod.rs

@@ -0,0 +1,2 @@
+pub mod memory_db;
+pub mod trie_walker;