perf(core): make TrieDb use NodeHash as key

CHANGELOG.md

@@ -2,6 +2,14 @@
 
 ## Perf
 
+### 2025-04-23
+
+- Make TrieDb trait use NodeHash as key [2517](https://github.com/lambdaclass/ethrex/pull/2517)
+
+### 2025-04-22
+
+- Transform the inlined variant of NodeHash to a constant sized array [2516](https://github.com/lambdaclass/ethrex/pull/2516)
+
 ### 2025-04-11
 
 - Removed some unnecessary clones and made some functions const: [2438](https://github.com/lambdaclass/ethrex/pull/2438)

crates/common/trie/db.rs

@@ -1,23 +1,23 @@
-use crate::error::TrieError;
+use crate::{error::TrieError, NodeHash};
 use std::{
     collections::HashMap,
     sync::{Arc, Mutex},
 };
 
 pub trait TrieDB: Send + Sync {
-    fn get(&self, key: Vec<u8>) -> Result<Option<Vec<u8>>, TrieError>;
-    fn put(&self, key: Vec<u8>, value: Vec<u8>) -> Result<(), TrieError>;
+    fn get(&self, key: NodeHash) -> Result<Option<Vec<u8>>, TrieError>;
+    fn put(&self, key: NodeHash, value: Vec<u8>) -> Result<(), TrieError>;
     // fn put_batch(&self, key: Vec<u8>, value: Vec<u8>) -> Result<(), TrieError>;
-    fn put_batch(&self, key_values: Vec<(Vec<u8>, Vec<u8>)>) -> Result<(), TrieError>;
+    fn put_batch(&self, key_values: Vec<(NodeHash, Vec<u8>)>) -> Result<(), TrieError>;
 }
 
 /// InMemory implementation for the TrieDB trait, with get and put operations.
 pub struct InMemoryTrieDB {
-    inner: Arc<Mutex<HashMap<Vec<u8>, Vec<u8>>>>,
+    inner: Arc<Mutex<HashMap<NodeHash, Vec<u8>>>>,
 }
 
 impl InMemoryTrieDB {
-    pub const fn new(map: Arc<Mutex<HashMap<Vec<u8>, Vec<u8>>>>) -> Self {
+    pub const fn new(map: Arc<Mutex<HashMap<NodeHash, Vec<u8>>>>) -> Self {
         Self { inner: map }
     }
     pub fn new_empty() -> Self {
@@ -28,7 +28,7 @@ impl InMemoryTrieDB {
 }
 
 impl TrieDB for InMemoryTrieDB {
-    fn get(&self, key: Vec<u8>) -> Result<Option<Vec<u8>>, TrieError> {
+    fn get(&self, key: NodeHash) -> Result<Option<Vec<u8>>, TrieError> {
         Ok(self
             .inner
             .lock()
@@ -37,15 +37,15 @@ impl TrieDB for InMemoryTrieDB {
             .cloned())
     }
 
-    fn put(&self, key: Vec<u8>, value: Vec<u8>) -> Result<(), TrieError> {
+    fn put(&self, key: NodeHash, value: Vec<u8>) -> Result<(), TrieError> {
         self.inner
             .lock()
             .map_err(|_| TrieError::LockError)?
             .insert(key, value);
         Ok(())
     }
 
-    fn put_batch(&self, key_values: Vec<(Vec<u8>, Vec<u8>)>) -> Result<(), TrieError> {
+    fn put_batch(&self, key_values: Vec<(NodeHash, Vec<u8>)>) -> Result<(), TrieError> {
         let mut db = self.inner.lock().map_err(|_| TrieError::LockError)?;
 
         for (key, value) in key_values {

crates/common/trie/node.rs

@@ -5,7 +5,6 @@ mod leaf;
 use std::array;
 
 pub use branch::BranchNode;
-use ethereum_types::H256;
 use ethrex_rlp::{decode::decode_bytes, error::RLPDecodeError, structs::Decoder};
 pub use extension::ExtensionNode;
 pub use leaf::LeafNode;
@@ -178,8 +177,8 @@ impl Node {
 
 fn decode_child(rlp: &[u8]) -> NodeHash {
     match decode_bytes(rlp) {
-        Ok((hash, &[])) if hash.len() == 32 => NodeHash::Hashed(H256::from_slice(hash)),
-        Ok((&[], &[])) => NodeHash::Inline(vec![]),
-        _ => NodeHash::Inline(rlp.to_vec()),
+        Ok((hash, &[])) if hash.len() == 32 => NodeHash::from_slice(hash),
+        Ok((&[], &[])) => NodeHash::default(),
+        _ => NodeHash::from_slice(rlp),
     }
 }

crates/common/trie/node/branch.rs

@@ -61,7 +61,7 @@ impl BranchNode {
             let child_hash = &self.choices[choice];
             if child_hash.is_valid() {
                 let child_node = state
-                    .get_node(child_hash.clone())?
+                    .get_node(*child_hash)?
                     .ok_or(TrieError::InconsistentTree)?;
                 child_node.get(state, path)
             } else {
@@ -93,7 +93,7 @@ impl BranchNode {
                 // Insert into existing child and then update it
                 choice_hash => {
                     let child_node = state
-                        .get_node(choice_hash.clone())?
+                        .get_node(*choice_hash)?
                         .ok_or(TrieError::InconsistentTree)?;
 
                     let child_node = child_node.insert(state, path, value)?;
@@ -138,7 +138,7 @@ impl BranchNode {
         let value = if let Some(choice_index) = path.next_choice() {
             if self.choices[choice_index].is_valid() {
                 let child_node = state
-                    .get_node(self.choices[choice_index].clone())?
+                    .get_node(self.choices[choice_index])?
                     .ok_or(TrieError::InconsistentTree)?;
                 // Remove value from child node
                 let (child_node, old_value) = child_node.remove(state, path.clone())?;
@@ -179,15 +179,14 @@ impl BranchNode {
             (1, false) => {
                 let (choice_index, child_hash) = children[0];
                 let child = state
-                    .get_node(child_hash.clone())?
+                    .get_node(*child_hash)?
                     .ok_or(TrieError::InconsistentTree)?;
                 match child {
                     // Replace self with an extension node leading to the child
-                    Node::Branch(_) => ExtensionNode::new(
-                        Nibbles::from_hex(vec![choice_index as u8]),
-                        child_hash.clone(),
-                    )
-                    .into(),
+                    Node::Branch(_) => {
+                        ExtensionNode::new(Nibbles::from_hex(vec![choice_index as u8]), *child_hash)
+                            .into()
+                    }
                     // Replace self with the child extension node, updating its path in the process
                     Node::Extension(mut extension_node) => {
                         extension_node.prefix.prepend(choice_index as u8);
@@ -207,7 +206,7 @@ impl BranchNode {
 
     /// Computes the node's hash
     pub fn compute_hash(&self) -> NodeHash {
-        NodeHash::from_encoded_raw(self.encode_raw())
+        NodeHash::from_encoded_raw(&self.encode_raw())
     }
 
     /// Encodes the node
@@ -217,7 +216,7 @@ impl BranchNode {
         for child in self.choices.iter() {
             match child {
                 NodeHash::Hashed(hash) => encoder = encoder.encode_bytes(&hash.0),
-                NodeHash::Inline(raw) if !raw.is_empty() => encoder = encoder.encode_raw(raw),
+                NodeHash::Inline(raw) if raw.1 != 0 => encoder = encoder.encode_raw(child.as_ref()),
                 _ => encoder = encoder.encode_bytes(&[]),
             }
         }
@@ -229,7 +228,7 @@ impl BranchNode {
     /// Inserts the node into the state and returns its hash
     pub fn insert_self(self, state: &mut TrieState) -> Result<NodeHash, TrieError> {
         let hash = self.compute_hash();
-        state.insert_node(self.into(), hash.clone());
+        state.insert_node(self.into(), hash);
         Ok(hash)
     }
 
@@ -253,7 +252,7 @@ impl BranchNode {
             let child_hash = &self.choices[choice];
             if child_hash.is_valid() {
                 let child_node = state
-                    .get_node(child_hash.clone())?
+                    .get_node(*child_hash)?
                     .ok_or(TrieError::InconsistentTree)?;
                 child_node.get_path(state, path, node_path)?;
             }

crates/common/trie/node/extension.rs

@@ -28,7 +28,7 @@ impl ExtensionNode {
         // Otherwise, no value is present.
         if path.skip_prefix(&self.prefix) {
             let child_node = state
-                .get_node(self.child.clone())?
+                .get_node(self.child)?
                 .ok_or(TrieError::InconsistentTree)?;
 
             child_node.get(state, path)
@@ -144,29 +144,22 @@ impl ExtensionNode {
 
     /// Computes the node's hash
     pub fn compute_hash(&self) -> NodeHash {
-        NodeHash::from_encoded_raw(self.encode_raw())
+        NodeHash::from_encoded_raw(&self.encode_raw())
     }
 
     /// Encodes the node
     pub fn encode_raw(&self) -> Vec<u8> {
         let mut buf = vec![];
         let mut encoder = Encoder::new(&mut buf).encode_bytes(&self.prefix.encode_compact());
-        match &self.child {
-            NodeHash::Inline(x) => {
-                encoder = encoder.encode_raw(x);
-            }
-            NodeHash::Hashed(x) => {
-                encoder = encoder.encode_bytes(&x.0);
-            }
-        }
+        encoder = self.child.encode(encoder);
         encoder.finish();
         buf
     }
 
     /// Inserts the node into the state and returns its hash
     pub fn insert_self(self, state: &mut TrieState) -> Result<NodeHash, TrieError> {
         let hash = self.compute_hash();
-        state.insert_node(self.into(), hash.clone());
+        state.insert_node(self.into(), hash);
         Ok(hash)
     }
 
@@ -187,7 +180,7 @@ impl ExtensionNode {
         // Continue to child
         if path.skip_prefix(&self.prefix) {
             let child_node = state
-                .get_node(self.child.clone())?
+                .get_node(self.child)?
                 .ok_or(TrieError::InconsistentTree)?;
             child_node.get_path(state, path, node_path)?;
         }

crates/common/trie/node/leaf.rs

@@ -101,7 +101,7 @@ impl LeafNode {
 
     /// Computes the node's hash
     pub fn compute_hash(&self) -> NodeHash {
-        NodeHash::from_encoded_raw(self.encode_raw())
+        NodeHash::from_encoded_raw(&self.encode_raw())
     }
 
     /// Encodes the node
@@ -118,7 +118,7 @@ impl LeafNode {
     /// Receives the offset that needs to be traversed to reach the leaf node from the canonical root, used to compute the node hash
     pub fn insert_self(self, state: &mut TrieState) -> Result<NodeHash, TrieError> {
         let hash = self.compute_hash();
-        state.insert_node(self.into(), hash.clone());
+        state.insert_node(self.into(), hash);
         Ok(hash)
     }