chore(l2): refactor get_embedded_node and remove hash cloning

crates/common/trie/node.rs

@@ -2,9 +2,13 @@ mod branch;
 mod extension;
 mod leaf;
 
-use std::sync::{Arc, OnceLock};
+use std::{
+    collections::BTreeMap,
+    sync::{Arc, OnceLock},
+};
 
 pub use branch::BranchNode;
+use ethereum_types::H256;
 use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
 pub use extension::ExtensionNode;
 pub use leaf::LeafNode;
@@ -83,6 +87,55 @@ impl NodeRef {
             NodeRef::Hash(hash) => *hash,
         }
     }
+
+    pub fn compute_hash_finalized(&mut self) -> &H256 {
+        let node_hash = match self {
+            NodeRef::Node(node, hash) => {
+                // uses get_or_init() instead of set() because the latter doesn't have lazy eval.
+                hash.get_or_init(|| node.compute_hash());
+                hash.get_mut().unwrap()
+            }
+            NodeRef::Hash(hash) => hash,
+        };
+        node_hash.finalize_mut()
+    }
+
+    pub fn resolve_subtrie(
+        &mut self,
+        all_nodes: &BTreeMap<H256, Vec<u8>>,
+    ) -> Result<(), TrieError> {
+        let finalized_hash = self.compute_hash_finalized();
+
+        let mut decoded_node = match all_nodes.get(finalized_hash) {
+            Some(rlp) => Node::decode(rlp)?,
+            None => return Ok(()),
+        };
+
+        match &mut decoded_node {
+            Node::Branch(node) => {
+                for choice in &mut node.choices {
+                    let NodeRef::Hash(hash) = choice else {
+                        unreachable!()
+                    };
+
+                    if hash.is_valid() {
+                        choice.resolve_subtrie(all_nodes)?;
+                    }
+                }
+            }
+            Node::Extension(node) => {
+                let NodeRef::Hash(_) = node.child else {
+                    unreachable!()
+                };
+
+                node.child.resolve_subtrie(all_nodes)?;
+            }
+            Node::Leaf(_) => {}
+        };
+
+        *self = decoded_node.into();
+        Ok(())
+    }
 }
 
 impl Default for NodeRef {

crates/common/trie/node_hash.rs

@@ -56,6 +56,24 @@ impl NodeHash {
         }
     }
 
+    /// Returns a reference to the finalized hash, hashing `self` if it's `NodeHash::Inline`, and storing the result by
+    /// changing it to the `NodeHash::Hashed` variant.
+    pub fn finalize_mut(&mut self) -> &H256 {
+        if let NodeHash::Inline(_) = self {
+            let hash = H256(
+                Keccak256::new()
+                    .chain_update(self.as_ref())
+                    .finalize()
+                    .into(),
+            );
+            *self = NodeHash::Hashed(hash);
+        }
+        let NodeHash::Hashed(hash_ref) = self else {
+            unreachable!();
+        };
+        hash_ref
+    }
+
     /// Returns true if the hash is valid
     /// The hash will only be considered invalid if it is empty
     /// Aka if it has a default value instead of being a product of hash computation

crates/common/trie/trie.rs

@@ -290,52 +290,9 @@ impl Trie {
         if root_hash == *EMPTY_TRIE_HASH {
             return Ok(NodeRef::default());
         }
-
-        let root_rlp = all_nodes.get(&root_hash).ok_or_else(|| {
-            TrieError::InconsistentTree(Box::new(InconsistentTreeError::RootNotFound(root_hash)))
-        })?;
-
-        fn get_embedded_node(
-            all_nodes: &BTreeMap<H256, Vec<u8>>,
-            cur_node_rlp: &[u8],
-        ) -> Result<Node, TrieError> {
-            let cur_node = Node::decode(cur_node_rlp)?;
-
-            Ok(match cur_node {
-                Node::Branch(mut node) => {
-                    for choice in &mut node.choices {
-                        let NodeRef::Hash(hash) = *choice else {
-                            unreachable!()
-                        };
-
-                        if hash.is_valid() {
-                            *choice = match all_nodes.get(&hash.finalize()) {
-                                Some(rlp) => get_embedded_node(all_nodes, rlp)?.into(),
-                                None => hash.into(),
-                            };
-                        }
-                    }
-
-                    (*node).into()
-                }
-                Node::Extension(mut node) => {
-                    let NodeRef::Hash(hash) = node.child else {
-                        unreachable!()
-                    };
-
-                    node.child = match all_nodes.get(&hash.finalize()) {
-                        Some(rlp) => get_embedded_node(all_nodes, rlp)?.into(),
-                        None => hash.into(),
-                    };
-
-                    node.into()
-                }
-                Node::Leaf(node) => node.into(),
-            })
-        }
-
-        let root = get_embedded_node(all_nodes, root_rlp)?;
-        Ok(root.into())
+        let mut root = NodeRef::Hash(root_hash.into());
+        root.resolve_subtrie(all_nodes)?;
+        Ok(root)
     }
 
     /// Builds a trie from a set of nodes with an empty InMemoryTrieDB as a backend because the nodes are embedded in the root.