fix(replay): fix execution with eth_getProof

cmd/ethrex_replay/src/cli.rs

@@ -23,26 +23,28 @@ use ethrex_storage::{
 };
 #[cfg(feature = "l2")]
 use ethrex_storage_rollup::EngineTypeRollup;
-use ethrex_trie::{InMemoryTrieDB, Node, NodeHash, NodeRef, node::LeafNode};
+use ethrex_trie::{
+    InMemoryTrieDB, Node,
+    node::{BranchNode, LeafNode},
+};
 use reqwest::Url;
 #[cfg(feature = "l2")]
 use std::path::Path;
 use std::{
     cmp::max,
-    collections::HashSet,
     io::Write,
     sync::{Arc, RwLock},
     time::{Instant, SystemTime},
 };
 use tracing::info;
 
-use crate::bench::run_and_measure;
 use crate::fetcher::get_blockdata;
 #[cfg(not(feature = "l2"))]
 use crate::fetcher::get_rangedata;
 #[cfg(not(feature = "l2"))]
 use crate::plot_composition::plot;
 use crate::run::{exec, prove, run_tx};
+use crate::{bench::run_and_measure, helpers::get_referenced_hashes};
 use crate::{
     block_run_report::{BlockRunReport, ReplayerMode},
     cache::Cache,
@@ -463,7 +465,7 @@ async fn replay_no_zkvm(cache: Cache, opts: &EthrexReplayOptions) -> eyre::Resul
     let guest_program = GuestProgramState::try_from(witness.clone())?;
 
     // This will contain all code hashes with the corresponding bytecode
-    // For the code hashes that we don't have we'll will it with <CodeHash, Bytes::new()>
+    // For the code hashes that we don't have we'll fill it with <CodeHash, Bytes::new()>
     let mut all_codes_hashed = guest_program.codes_hashed.clone();
 
     let in_memory_store = InMemoryStore::new();
@@ -479,42 +481,16 @@ async fn replay_no_zkvm(cache: Cache, opts: &EthrexReplayOptions) -> eyre::Resul
         //   2. Get all code hashes that exist in the accounts that we have so that if we don't have the code we set it to empty bytes.
         // We do these things because sometimes the witness may be incomplete and in those cases we don't want failures for missing data.
         // This only applies when we use the InMemoryDatabase and not when we use the ExecutionWitness as database, that's because in the latter failures are dismissed and we fall back to default values.
-        let mut nodes = state_trie_nodes.lock().unwrap();
-        let mut referenced_node_hashes: HashSet<NodeHash> = HashSet::new(); // All hashes referenced in the trie (by Branch or Ext nodes).
-
-        for (_node_hash, node_rlp) in nodes.iter() {
-            let node = Node::decode(node_rlp)?;
-            match node {
-                Node::Branch(node) => {
-                    for choice in &node.choices {
-                        let NodeRef::Hash(hash) = *choice else {
-                            unreachable!()
-                        };
-
-                        referenced_node_hashes.insert(hash);
-                    }
-                }
-                Node::Extension(node) => {
-                    let NodeRef::Hash(hash) = node.child else {
-                        unreachable!()
-                    };
-
-                    referenced_node_hashes.insert(hash);
-                }
-                Node::Leaf(node) => {
-                    let info = AccountState::decode(&node.value)?;
-                    all_codes_hashed.entry(info.code_hash).or_insert(vec![]);
-                }
-            }
-        }
+        let mut state_nodes = state_trie_nodes.lock().unwrap();
+        let referenced_node_hashes = get_referenced_hashes(&state_nodes)?;
 
+        let dummy_leaf = Node::from(LeafNode::default()).encode_to_vec();
         // Insert arbitrary leaf nodes to state trie.
         for hash in referenced_node_hashes {
-            let dummy_leaf: Node = LeafNode::default().into();
-            nodes.entry(hash).or_insert(dummy_leaf.encode_to_vec());
+            state_nodes.entry(hash).or_insert(dummy_leaf.clone());
         }
 
-        drop(nodes);
+        drop(state_nodes);
 
         let mut inner_store = in_memory_store.inner()?;
 
@@ -541,13 +517,33 @@ async fn replay_no_zkvm(cache: Cache, opts: &EthrexReplayOptions) -> eyre::Resul
                 continue;
             };
 
-            let storage_root = AccountState::decode(&account_state_rlp)?.storage_root;
+            let account_state = AccountState::decode(&account_state_rlp)?;
 
+            // If code hash of account isn't present insert empty code so that if not found the execution doesn't break.
+            let code_hash = account_state.code_hash;
+            all_codes_hashed.entry(code_hash).or_insert(vec![]);
+
+            let storage_root = account_state.storage_root;
             let storage_trie = match InMemoryTrieDB::from_nodes(storage_root, all_nodes) {
                 Ok(trie) => trie.inner,
                 Err(_) => continue,
             };
 
+            // Fill storage trie with dummy branch nodes that have the hash of the missing nodes
+            // This is useful for eth_getProofs when we want to restructure the trie after removing a node whose sibling isn't known
+            // We assume the sibling is a branch node because we already covered the cases in which it's a Leaf or Extension node by injecting nodes in the witness.
+            // For more info read: https://github.com/kkrt-labs/zk-pig/blob/v0.8.0/docs/modified-mpt.md
+            {
+                let mut storage_nodes = storage_trie.lock().unwrap();
+                let dummy_branch = Node::from(BranchNode::default()).encode_to_vec();
+
+                let referenced_storage_node_hashes = get_referenced_hashes(&storage_nodes)?;
+
+                for hash in referenced_storage_node_hashes {
+                    storage_nodes.entry(hash).or_insert(dummy_branch.clone());
+                }
+            }
+
             inner_store
                 .storage_trie_nodes
                 .insert(H256::from_slice(&hashed_address), storage_trie);

cmd/ethrex_replay/src/helpers.rs

@@ -0,0 +1,36 @@
+use std::collections::{BTreeMap, HashSet};
+
+use ethrex_rlp::decode::RLPDecode;
+use ethrex_trie::{Node, NodeHash, NodeRef};
+
+/// Given a mapping of nodes with their corresponding hash get all hashes referenced by branch and extension nodes.
+pub fn get_referenced_hashes(
+    nodes: &BTreeMap<NodeHash, Vec<u8>>,
+) -> eyre::Result<HashSet<NodeHash>> {
+    let mut referenced_hashes: HashSet<NodeHash> = HashSet::new();
+
+    for (_node_hash, node_rlp) in nodes.iter() {
+        let node = Node::decode(node_rlp)?;
+        match node {
+            Node::Branch(node) => {
+                for choice in &node.choices {
+                    if let NodeRef::Hash(hash) = *choice {
+                        referenced_hashes.insert(hash);
+                    } else {
+                        return Err(eyre::eyre!("Branch node contains non-hash reference"));
+                    }
+                }
+            }
+            Node::Extension(node) => {
+                if let NodeRef::Hash(hash) = node.child {
+                    referenced_hashes.insert(hash);
+                } else {
+                    return Err(eyre::eyre!("Extension node contains non-hash reference"));
+                }
+            }
+            Node::Leaf(_node) => {}
+        }
+    }
+
+    Ok(referenced_hashes)
+}

cmd/ethrex_replay/src/lib.rs

@@ -3,6 +3,7 @@ pub mod block_run_report;
 mod cache;
 pub mod cli;
 mod fetcher;
+mod helpers;
 #[cfg(not(feature = "l2"))]
 mod plot_composition;
 pub mod rpc;

cmd/ethrex_replay/src/rpc/db.rs

@@ -581,7 +581,9 @@ pub fn get_potential_child_nodes(proof: &[NodeRLP], key: &PathRLP) -> Option<Vec
     }
 
     let hash = if let Some(root) = proof.first() {
-        H256::from_slice(&Keccak256::digest(root))
+        let hash = H256::from_slice(&Keccak256::digest(root));
+        state_nodes.insert(hash, root.clone());
+        hash
     } else {
         *EMPTY_KECCACK_HASH
     };

crates/common/trie/Cargo.toml

@@ -45,4 +45,3 @@ path = "./trie.rs"
 [[bench]]
 name = "trie_bench"
 harness = false
-

crates/common/trie/node/branch.rs

@@ -6,7 +6,7 @@ use super::{ExtensionNode, LeafNode, Node, NodeRef, ValueOrHash};
 
 /// Branch Node of an an Ethereum Compatible Patricia Merkle Trie
 /// Contains the node's value and the hash of its children nodes
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, Default)]
 pub struct BranchNode {
     pub choices: [NodeRef; 16],
     pub value: ValueRLP,