fix(l1): ethrex_replay reduce cases were OOM happen with SP1 (#4076)

**Motivation**

Some of the blocks execution with `ethrex_replay` ans SP1 where failing
due to Out Of Memory. We fix some of them here

**Description**

- Remove unnecessary clones and variables from the `from_nodes(..)`
function
- Change the behavior of the `storage_tries`, we build them on demand
and store them in a hash map indexed by the address

**How to Test**

Running this with an url for a Hoodi client (using a cache for the block
is faster):
```
cargo run --release --features sp1 -- execute block 989143 --rpc-url <RPC_URL> --network hoodi
```

This block now succeeds when in the base branch it doesn't.

**What's next**

This could be more optimized following this
[issue](#4094)

Advances #4089

---------

Co-authored-by: ilitteri <ilitteri@fi.uba.ar>
Co-authored-by: Jeremías Salomón <48994069+JereSalo@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: JereSalo <jeresalo17@gmail.com>
Co-authored-by: Ivan Litteri <67517699+ilitteri@users.noreply.github.com>
Co-authored-by: LeanSerra <46695152+LeanSerra@users.noreply.github.com>

Author: 5 people

cmd/ethrex_replay/src/run.rs

@@ -33,7 +33,7 @@ pub async fn run_tx(
         .ok_or(eyre::Error::msg("missing block data"))?;
     let mut remaining_gas = block.header.gas_limit;
     let mut prover_db = cache.witness;
-    prover_db.rebuild_tries()?;
+    prover_db.rebuild_state_trie()?;
     let mut wrapped_db = ExecutionWitnessWrapper::new(prover_db);
 
     let vm_type = if l2 { VMType::L2 } else { VMType::L1 };

crates/blockchain/blockchain.rs

@@ -350,10 +350,10 @@ impl Blockchain {
             codes,
             //TODO: See if we should call rebuild_tries() here for initializing these fields so that we don't have an inconsistent struct. (#4056)
             state_trie: None,
-            storage_tries: None,
             block_headers,
             chain_config,
             state_trie_nodes: used_trie_nodes.into_iter().map(Bytes::from).collect(),
+            storage_tries: HashMap::new(),
             parent_block_header: self
                 .storage
                 .get_block_header_by_hash(first_block_header.parent_hash)?

crates/common/trie/trie.rs

@@ -262,21 +262,17 @@ impl Trie {
             .map(|node| {
                 (
                     NodeHash::from_slice(&Keccak256::new_with_prefix(node).finalize()),
-                    node,
+                    node.clone(),
                 )
             })
             .collect::<HashMap<_, _>>();
-        let nodes = storage
-            .iter()
-            .map(|(node_hash, nodes)| (*node_hash, (*nodes).clone()))
-            .collect::<HashMap<_, _>>();
         let Some(root) = root else {
-            let in_memory_trie = Box::new(InMemoryTrieDB::new(Arc::new(Mutex::new(nodes))));
+            let in_memory_trie = Box::new(InMemoryTrieDB::new(Arc::new(Mutex::new(storage))));
             return Ok(Trie::new(in_memory_trie));
         };
 
         fn inner(
-            storage: &mut HashMap<NodeHash, &Vec<u8>>,
+            storage: &mut HashMap<NodeHash, Vec<u8>>,
             node: &NodeRLP,
         ) -> Result<Node, TrieError> {
             Ok(match Node::decode_raw(node)? {
@@ -288,7 +284,7 @@ impl Trie {
 
                         if hash.is_valid() {
                             *choice = match storage.remove(&hash) {
-                                Some(rlp) => inner(storage, rlp)?.into(),
+                                Some(rlp) => inner(storage, &rlp)?.into(),
                                 None => hash.into(),
                             };
                         }
@@ -302,7 +298,7 @@ impl Trie {
                     };
 
                     node.child = match storage.remove(&hash) {
-                        Some(rlp) => inner(storage, rlp)?.into(),
+                        Some(rlp) => inner(storage, &rlp)?.into(),
                         None => hash.into(),
                     };
 
@@ -313,11 +309,7 @@ impl Trie {
         }
 
         let root = inner(&mut storage, root)?.into();
-        let nodes = storage
-            .into_iter()
-            .map(|(node_hash, nodes)| (node_hash, nodes.clone()))
-            .collect::<HashMap<_, _>>();
-        let in_memory_trie = Box::new(InMemoryTrieDB::new(Arc::new(Mutex::new(nodes))));
+        let in_memory_trie = Box::new(InMemoryTrieDB::new(Arc::new(Mutex::new(storage))));
 
         let mut trie = Trie::new(in_memory_trie);
         trie.root = root;

crates/common/types/block_execution_witness.rs

@@ -51,11 +51,10 @@ pub struct ExecutionWitnessResult {
     #[serde(skip)]
     #[rkyv(with = rkyv::with::Skip)]
     pub state_trie: Option<Trie>,
-    // Indexed by account
-    // Pruned storage MPT
+    // Storage tries accessed by account address
     #[serde(skip)]
     #[rkyv(with = rkyv::with::Skip)]
-    pub storage_tries: Option<HashMap<Address, Trie>>,
+    pub storage_tries: HashMap<Address, Trie>,
     // Block headers needed for BLOCKHASH opcode
     pub block_headers: HashMap<u64, BlockHeader>,
     // Parent block header to get the initial state root
@@ -85,36 +84,17 @@ pub enum ExecutionWitnessError {
 }
 
 impl ExecutionWitnessResult {
-    pub fn rebuild_tries(&mut self) -> Result<(), ExecutionWitnessError> {
+    /// Use the state nodes to build the state trie and store them in `self.state_trie`
+    /// This function will fail if the state trie cannot be rebuilt.
+    pub fn rebuild_state_trie(&mut self) -> Result<(), ExecutionWitnessError> {
         let state_trie = rebuild_trie(self.parent_block_header.state_root, &self.state_trie_nodes)?;
-
-        // Keys can either be account addresses or storage slots. They have different sizes,
-        // so we filter them by size. The from_slice method panics if the input has the wrong size.
-        let addresses: Vec<Address> = self
-            .keys
-            .iter()
-            .filter(|k| k.len() == Address::len_bytes())
-            .map(|k| Address::from_slice(k))
-            .collect();
-
-        let storage_tries: HashMap<Address, Trie> = HashMap::from_iter(
-            addresses
-                .iter()
-                .filter_map(|addr| {
-                    Some((
-                        *addr,
-                        Self::rebuild_storage_trie(addr, &state_trie, &self.state_trie_nodes)?,
-                    ))
-                })
-                .collect::<Vec<(Address, Trie)>>(),
-        );
-
         self.state_trie = Some(state_trie);
-        self.storage_tries = Some(storage_tries);
 
         Ok(())
     }
 
+    /// Helper function to rebuild the storage trie for a given account address
+    /// Returns if root is not empty, an Option with the rebuilt trie
     // This function is an option because we expect it to fail sometimes, and we just want to filter it
     pub fn rebuild_storage_trie(address: &H160, trie: &Trie, state: &[Bytes]) -> Option<Trie> {
         let account_state_rlp = trie.get(&hash_address(address)).ok()??;
@@ -128,12 +108,14 @@ impl ExecutionWitnessResult {
         rebuild_trie(account_state.storage_root, state).ok()
     }
 
+    /// Helper function to apply account updates to the execution witness
+    /// It updates the state trie and storage tries with the given account updates
+    /// Returns an error if the updates cannot be applied
     pub fn apply_account_updates(
         &mut self,
         account_updates: &[AccountUpdate],
     ) -> Result<(), ExecutionWitnessError> {
-        let (Some(state_trie), Some(storage_tries_map)) =
-            (self.state_trie.as_mut(), self.storage_tries.as_mut())
+        let (Some(state_trie), storage_tries) = (self.state_trie.as_mut(), &mut self.storage_tries)
         else {
             return Err(ExecutionWitnessError::ApplyAccountUpdates(
                 "Tried to apply account updates before rebuilding the tries".to_string(),
@@ -169,10 +151,9 @@ impl ExecutionWitnessResult {
                 }
                 // Store the added storage in the account's storage trie and compute its new root
                 if !update.added_storage.is_empty() {
-                    let storage_trie =
-                        storage_tries_map.entry(update.address).or_insert_with(|| {
-                            Trie::from_nodes(None, &[]).expect("failed to create empty trie")
-                        });
+                    let storage_trie = storage_tries.entry(update.address).or_insert_with(|| {
+                        Trie::from_nodes(None, &[]).expect("failed to create empty trie")
+                    });
 
                     // Inserts must come before deletes, otherwise deletes might require extra nodes
                     // Example:
@@ -207,6 +188,8 @@ impl ExecutionWitnessResult {
         Ok(())
     }
 
+    /// Returns the root hash of the state trie
+    /// Returns an error if the state trie is not built yet
     pub fn state_trie_root(&self) -> Result<H256, ExecutionWitnessError> {
         let state_trie = self
             .state_trie
@@ -254,6 +237,8 @@ impl ExecutionWitnessResult {
         Ok(None)
     }
 
+    /// Retrieves the parent block header for the specified block number
+    /// Searches within `self.block_headers`
     pub fn get_block_parent_header(
         &self,
         block_number: u64,
@@ -263,6 +248,8 @@ impl ExecutionWitnessResult {
             .ok_or(ExecutionWitnessError::MissingParentHeaderOf(block_number))
     }
 
+    /// Retrieves the account info based on what is stored in the state trie.
+    /// Returns an error if the state trie is not rebuilt or if decoding the account state fails.
     pub fn get_account_info(
         &self,
         address: Address,
@@ -289,6 +276,8 @@ impl ExecutionWitnessResult {
         }))
     }
 
+    /// Fetches the block hash for a specific block number.
+    /// Looks up `self.block_headers` and computes the hash if it is not already computed.
     pub fn get_block_hash(&self, block_number: u64) -> Result<H256, ExecutionWitnessError> {
         self.block_headers
             .get(&block_number)
@@ -300,21 +289,31 @@ impl ExecutionWitnessResult {
             })
     }
 
+    /// Retrieves a storage slot value for an account in its storage trie.
+    ///
+    /// Lazily builds the storage trie for the address if not already available.
+    /// This lazy loading approach minimizes memory usage by only building tries when needed.
     pub fn get_storage_slot(
-        &self,
+        &mut self,
         address: Address,
         key: H256,
     ) -> Result<Option<U256>, ExecutionWitnessError> {
-        let storage_tries_map =
-            self.storage_tries
-                .as_ref()
-                .ok_or(ExecutionWitnessError::Database(
-                    "ExecutionWitness: Tried to get storage slot before rebuilding tries"
+        let storage_trie = if let Some(storage_trie) = self.storage_tries.get(&address) {
+            storage_trie
+        } else {
+            let Some(state_trie) = self.state_trie.as_ref() else {
+                return Err(ExecutionWitnessError::Database(
+                    "ExecutionWitness: Tried to get storage slot before rebuilding state trie."
                         .to_string(),
-                ))?;
+                ));
+            };
+            let Some(storage_trie) =
+                Self::rebuild_storage_trie(&address, state_trie, &self.state_trie_nodes)
+            else {
+                return Ok(None);
+            };
 
-        let Some(storage_trie) = storage_tries_map.get(&address) else {
-            return Ok(None);
+            self.storage_tries.entry(address).or_insert(storage_trie)
         };
         let hashed_key = hash_key(&key);
         if let Some(encoded_key) = storage_trie
@@ -331,10 +330,13 @@ impl ExecutionWitnessResult {
         }
     }
 
+    /// Retrieves the chain configuration for the execution witness.
     pub fn get_chain_config(&self) -> Result<ChainConfig, ExecutionWitnessError> {
         Ok(self.chain_config)
     }
 
+    /// Retrieves the account code for a specific account.
+    /// Returns an Err if the code is not found.
     pub fn get_account_code(&self, code_hash: H256) -> Result<bytes::Bytes, ExecutionWitnessError> {
         if code_hash == *EMPTY_KECCACK_HASH {
             return Ok(Bytes::new());
@@ -447,16 +449,3 @@ pub fn rebuild_trie(initial_state: H256, state: &[Bytes]) -> Result<Trie, Execut
     )
     .map_err(|e| ExecutionWitnessError::RebuildTrie(format!("Failed to build state trie {e}")))
 }
-
-// This function is an option because we expect it to fail sometimes, and we just want to filter it
-pub fn rebuild_storage_trie(address: &H160, trie: &Trie, state: &[Bytes]) -> Option<Trie> {
-    let account_state_rlp = trie.get(&hash_address(address)).ok()??;
-
-    let account_state = AccountState::decode(&account_state_rlp).ok()?;
-
-    if account_state.storage_root == *EMPTY_TRIE_HASH {
-        return None;
-    }
-
-    rebuild_trie(account_state.storage_root, state).ok()
-}

crates/l2/prover/zkvm/interface/sp1/Cargo.lock

@@ -163,7 +163,7 @@ pub fn stateless_validation_l2(
         codes: db.codes.clone(),
         keys: db.keys.clone(),
         state_trie: None,
-        storage_tries: None,
+        storage_tries: HashMap::new(),
         state_trie_nodes: db.state_trie_nodes.clone(),
         parent_block_header: db.parent_block_header.clone(),
     };
@@ -192,7 +192,7 @@ pub fn stateless_validation_l2(
     // Check state diffs are valid
     let blob_versioned_hash = if !validium {
         initial_db
-            .rebuild_tries()
+            .rebuild_state_trie()
             .map_err(|_| StatelessExecutionError::InvalidInitialStateTrie)?;
         let wrapped_db = ExecutionWitnessWrapper::new(initial_db);
         let state_diff = prepare_state_diff(
@@ -234,7 +234,7 @@ fn execute_stateless(
     mut db: ExecutionWitnessResult,
     elasticity_multiplier: u64,
 ) -> Result<StatelessResult, StatelessExecutionError> {
-    db.rebuild_tries()
+    db.rebuild_state_trie()
         .map_err(|_| StatelessExecutionError::InvalidInitialStateTrie)?;
 
     let mut wrapped_db = ExecutionWitnessWrapper::new(db);
@@ -271,7 +271,6 @@ fn execute_stateless(
     let initial_state_hash = wrapped_db
         .state_trie_root()
         .map_err(StatelessExecutionError::ExecutionWitness)?;
-
     if initial_state_hash != parent_block_header.state_root {
         return Err(StatelessExecutionError::InvalidInitialStateTrie);
     }

crates/l2/prover/zkvm/interface/src/execution.rs

@@ -92,14 +92,14 @@ pub fn execution_witness_from_rpc_chain_config(
         state_trie_nodes: rpc_witness.state,
         keys: rpc_witness.keys,
         codes,
-        state_trie: None,    // `None` because we'll rebuild the tries afterwards
-        storage_tries: None, // `None` because we'll rebuild the tries afterwards
+        state_trie: None, // `None` because we'll rebuild the tries afterwards
+        storage_tries: HashMap::new(), // empty map because we'll rebuild the tries afterwards
         block_headers,
         chain_config,
         parent_block_header: parent_header,
     };
 
-    witness.rebuild_tries()?;
+    witness.rebuild_state_trie()?;
 
     Ok(witness)
 }