perf(core): make TrieDb use NodeHash as key

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Perf
 
+### 2025-04-28
+
+- Make TrieDb trait use NodeHash as key [2517](https://github.com/lambdaclass/ethrex/pull/2517)
+
 ### 2025-04-22
 
 - Avoid calculating state transitions after every block in bulk mode [2519](https://github.com/lambdaclass/ethrex/pull/2519)

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_hash.rs

@@ -87,6 +87,20 @@ impl NodeHash {
         }
         encoder
     }
+
+    pub fn len(&self) -> usize {
+        match self {
+            NodeHash::Hashed(h256) => h256.as_bytes().len(),
+            NodeHash::Inline(value) => value.1 as usize,
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        match self {
+            NodeHash::Hashed(h256) => h256.as_bytes().is_empty(),
+            NodeHash::Inline(value) => value.1 == 0,
+        }
+    }
 }
 
 impl From<Vec<u8>> for NodeHash {