perf(l1): decoded TrieLayerCache: skip Node::decode() on cache hits

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Perf
 
+### 2026-03-12
+
+- Store decoded `Arc<Node>` in TrieLayerCache to eliminate RLP decoding on cache hits, and split cache into separate account/storage/leaf maps
+
 ### 2026-03-05
 
 - Switch hot EVM and mempool HashMaps to FxHashMap for faster hashing [#6303](https://github.com/lambdaclass/ethrex/pull/6303)

crates/blockchain/blockchain.rs

@@ -83,7 +83,7 @@ use ethrex_storage::{
     AccountUpdatesList, Store, UpdateBatch, error::StoreError, hash_address, hash_key,
 };
 use ethrex_trie::node::{BranchNode, ExtensionNode, LeafNode};
-use ethrex_trie::{Nibbles, Node, NodeRef, Trie, TrieError, TrieNode};
+use ethrex_trie::{Nibbles, Node, NodeRef, Trie, TrieCommitEntry, TrieError};
 use ethrex_vm::backends::CachingDatabase;
 use ethrex_vm::backends::levm::LEVM;
 use ethrex_vm::backends::levm::db::DatabaseLogger;
@@ -262,22 +262,22 @@ enum MerklizationRequest {
 struct CollectedStateMsg {
     index: u8,
     subroot: Box<BranchNode>,
-    state_nodes: Vec<TrieNode>,
-    storage_nodes: Vec<(H256, Vec<TrieNode>)>,
+    state_nodes: Vec<TrieCommitEntry>,
+    storage_nodes: Vec<(H256, Vec<TrieCommitEntry>)>,
 }
 
 struct CollectedStorageMsg {
     index: u8,
     prefix: H256,
     subroot: Box<BranchNode>,
-    nodes: Vec<TrieNode>,
+    nodes: Vec<TrieCommitEntry>,
 }
 
 #[derive(Default)]
 struct PreMerkelizedAccountState {
     info: Option<AccountInfo>,
     storage_root: Option<Box<BranchNode>>,
-    nodes: Vec<TrieNode>,
+    nodes: Vec<TrieCommitEntry>,
 }
 
 /// Work item for BAL state trie shard workers.
@@ -688,7 +688,7 @@ impl Blockchain {
             state.nodes.extend(nodes);
         }
 
-        let mut storage_updates: Vec<(H256, Vec<TrieNode>)> = Default::default();
+        let mut storage_updates: Vec<(H256, Vec<TrieCommitEntry>)> = Default::default();
 
         for (hashed_account, state) in account_state {
             let bucket = hashed_account.as_fixed_bytes()[0] >> 4;
@@ -728,7 +728,10 @@ impl Blockchain {
                 let hash = root.commit(Nibbles::default(), &mut state_updates);
                 hash.finalize()
             } else {
-                state_updates.push((Nibbles::default(), vec![RLP_NULL]));
+                state_updates.push(TrieCommitEntry::LeafValue {
+                    path: Nibbles::default(),
+                    value: vec![RLP_NULL],
+                });
                 *EMPTY_TRIE_HASH
             };
 
@@ -845,7 +848,7 @@ impl Blockchain {
 
         // Compute storage roots in parallel
         let mut storage_roots: Vec<Option<H256>> = vec![None; accounts.len()];
-        let mut storage_updates: Vec<(H256, Vec<TrieNode>)> = Vec::new();
+        let mut storage_updates: Vec<(H256, Vec<TrieCommitEntry>)> = Vec::new();
 
         std::thread::scope(|s| -> Result<(), StoreError> {
             let accounts_ref = &accounts;
@@ -861,8 +864,8 @@ impl Blockchain {
                             .name(format!("bal_storage_worker_{worker_id}"))
                             .spawn_scoped(
                                 s,
-                                move || -> Result<Vec<(usize, H256, Vec<TrieNode>)>, StoreError> {
-                                    let mut results: Vec<(usize, H256, Vec<TrieNode>)> = Vec::new();
+                                move || -> Result<Vec<(usize, H256, Vec<TrieCommitEntry>)>, StoreError> {
+                                    let mut results: Vec<(usize, H256, Vec<TrieCommitEntry>)> = Vec::new();
                                     // Open one state trie per worker for storage root lookups
                                     let state_trie =
                                         self.storage.open_state_trie(parent_state_root)?;
@@ -879,7 +882,10 @@ impl Blockchain {
                                             results.push((
                                                 idx,
                                                 *EMPTY_TRIE_HASH,
-                                                vec![(Nibbles::default(), vec![RLP_NULL])],
+                                                vec![TrieCommitEntry::LeafValue {
+                                                    path: Nibbles::default(),
+                                                    value: vec![RLP_NULL],
+                                                }],
                                             ));
                                             continue;
                                         }
@@ -967,7 +973,7 @@ impl Blockchain {
                         .name(format!("bal_state_shard_{index}"))
                         .spawn_scoped(
                             s,
-                            move || -> Result<(Box<BranchNode>, Vec<TrieNode>), StoreError> {
+                            move || -> Result<(Box<BranchNode>, Vec<TrieCommitEntry>), StoreError> {
                                 let mut state_trie =
                                     self.storage.open_state_trie(parent_state_root)?;
 
@@ -1036,7 +1042,10 @@ impl Blockchain {
                 let hash = root.commit(Nibbles::default(), &mut state_updates);
                 hash.finalize()
             } else {
-                state_updates.push((Nibbles::default(), vec![RLP_NULL]));
+                state_updates.push(TrieCommitEntry::LeafValue {
+                    path: Nibbles::default(),
+                    value: vec![RLP_NULL],
+                });
                 *EMPTY_TRIE_HASH
             };
 
@@ -1186,7 +1195,10 @@ impl Blockchain {
                             let hash = root.commit(Nibbles::default(), &mut state.nodes);
                             storage_root = Some(hash.finalize());
                         } else {
-                            state.nodes.push((Nibbles::default(), vec![RLP_NULL]));
+                            state.nodes.push(TrieCommitEntry::LeafValue {
+                                path: Nibbles::default(),
+                                value: vec![RLP_NULL],
+                            });
                             storage_root = Some(*EMPTY_TRIE_HASH);
                         }
                     }
@@ -2799,15 +2811,18 @@ fn branchify(node: Node) -> Box<BranchNode> {
     }
 }
 
-fn collect_trie(index: u8, mut trie: Trie) -> Result<(Box<BranchNode>, Vec<TrieNode>), TrieError> {
+fn collect_trie(
+    index: u8,
+    mut trie: Trie,
+) -> Result<(Box<BranchNode>, Vec<TrieCommitEntry>), TrieError> {
     let root = branchify(
         trie.root_node()?
             .map(Arc::unwrap_or_clone)
             .unwrap_or_else(|| Node::Branch(Box::default())),
     );
     trie.root = Node::Branch(root).into();
     let (_, mut nodes) = trie.collect_changes_since_last_hash();
-    nodes.retain(|(nib, _)| nib.as_ref().first() == Some(&index));
+    nodes.retain(|entry| entry.path().as_ref().first() == Some(&index));
 
     let Some(Node::Branch(root)) = trie.root_node()?.map(Arc::unwrap_or_clone) else {
         return Err(TrieError::InvalidInput);

crates/common/trie/db.rs

@@ -13,6 +13,12 @@ pub type NodeMap = Arc<Mutex<BTreeMap<Vec<u8>, Vec<u8>>>>;
 pub trait TrieDB: Send + Sync {
     fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError>;
     fn put_batch(&self, key_values: Vec<(Nibbles, Vec<u8>)>) -> Result<(), TrieError>;
+
+    /// Try to get a decoded node directly from a cache.
+    /// Default returns None — callers fall back to get() + Node::decode().
+    fn get_node(&self, _key: Nibbles) -> Result<Option<Arc<Node>>, TrieError> {
+        Ok(None)
+    }
     // TODO: replace putbatch with this function.
     fn put_batch_no_alloc(&self, key_values: &[(Nibbles, Node)]) -> Result<(), TrieError> {
         self.put_batch(
@@ -77,7 +83,10 @@ impl InMemoryTrieDB {
 
         let hashed_nodes = hashed_nodes
             .into_iter()
-            .map(|(k, v)| (k.into_vec(), v))
+            .map(|entry| {
+                let (k, v) = entry.into_rlp_pair();
+                (k.into_vec(), v)
+            })
             .collect();
 
         let in_memory_trie = Arc::new(Mutex::new(hashed_nodes));

crates/common/trie/logger.rs

@@ -33,6 +33,15 @@ impl TrieLogger {
 }
 
 impl TrieDB for TrieLogger {
+    fn get_node(&self, key: Nibbles) -> Result<Option<Arc<Node>>, TrieError> {
+        if let Some(node) = self.inner_db.get_node(key)? {
+            let mut lock = self.witness.lock().map_err(|_| TrieError::LockError)?;
+            lock.insert(node.compute_hash(), (*node).clone());
+            return Ok(Some(node));
+        }
+        Ok(None)
+    }
+
     fn get(&self, key: Nibbles) -> Result<Option<Vec<u8>>, TrieError> {
         let result = self.inner_db.get(key)?;
         if let Some(result) = result.as_ref()

crates/common/trie/node.rs

@@ -16,7 +16,7 @@ use rkyv::{
     with::Skip,
 };
 
-use crate::{NodeRLP, TrieDB, error::TrieError, nibbles::Nibbles};
+use crate::{NodeRLP, TrieCommitEntry, TrieDB, error::TrieError, nibbles::Nibbles};
 
 use super::{ValueRLP, node_hash::NodeHash};
 
@@ -58,11 +58,17 @@ impl NodeRef {
             NodeRef::Hash(hash @ NodeHash::Inline(_)) => {
                 Ok(Some(Arc::new(Node::decode(hash.as_ref())?)))
             }
-            NodeRef::Hash(_) => db
-                .get(path)?
-                .filter(|rlp| !rlp.is_empty())
-                .map(|rlp| Ok(Arc::new(Node::decode(&rlp)?)))
-                .transpose(),
+            NodeRef::Hash(_) => {
+                // Try decoded cache first
+                if let Some(node) = db.get_node(path.clone())? {
+                    return Ok(Some(node));
+                }
+                // Fall back to raw bytes + decode
+                db.get(path)?
+                    .filter(|rlp| !rlp.is_empty())
+                    .map(|rlp| Ok(Arc::new(Node::decode(&rlp)?)))
+                    .transpose()
+            }
         }
     }
 
@@ -78,14 +84,19 @@ impl NodeRef {
             NodeRef::Hash(hash @ NodeHash::Inline(_)) => {
                 Ok(Some(Arc::new(Node::decode(hash.as_ref())?)))
             }
-            NodeRef::Hash(hash @ NodeHash::Hashed(_)) => db
-                .get(path)?
-                .filter(|rlp| !rlp.is_empty())
-                .and_then(|rlp| match Node::decode(&rlp) {
-                    Ok(node) => (node.compute_hash() == *hash).then_some(Ok(Arc::new(node))),
-                    Err(err) => Some(Err(TrieError::RLPDecode(err))),
-                })
-                .transpose(),
+            NodeRef::Hash(hash @ NodeHash::Hashed(_)) => {
+                // Try decoded cache first (cache contents are trusted)
+                if let Some(node) = db.get_node(path.clone())? {
+                    return Ok(Some(node));
+                }
+                db.get(path)?
+                    .filter(|rlp| !rlp.is_empty())
+                    .and_then(|rlp| match Node::decode(&rlp) {
+                        Ok(node) => (node.compute_hash() == *hash).then_some(Ok(Arc::new(node))),
+                        Err(err) => Some(Err(TrieError::RLPDecode(err))),
+                    })
+                    .transpose()
+            }
         }
     }
 
@@ -108,6 +119,11 @@ impl NodeRef {
                 self.get_node_mut(db, path)
             }
             NodeRef::Hash(hash @ NodeHash::Hashed(_)) => {
+                // Try decoded cache first
+                if let Some(node) = db.get_node(path.clone())? {
+                    *self = NodeRef::Node(node, OnceLock::from(*hash));
+                    return self.get_node_mut(db, path);
+                }
                 let Some(node) = db
                     .get(path.clone())?
                     .filter(|rlp| !rlp.is_empty())
@@ -129,30 +145,37 @@ impl NodeRef {
         }
     }
 
-    pub fn commit(&mut self, path: Nibbles, acc: &mut Vec<(Nibbles, Vec<u8>)>) -> NodeHash {
+    pub fn commit(&mut self, path: Nibbles, acc: &mut Vec<TrieCommitEntry>) -> NodeHash {
         match *self {
             NodeRef::Node(ref mut node, ref mut hash) => {
                 if let Some(hash) = hash.get() {
                     return *hash;
                 }
                 match Arc::make_mut(node) {
-                    Node::Branch(node) => {
-                        for (choice, node) in &mut node.choices.iter_mut().enumerate() {
-                            node.commit(path.append_new(choice as u8), acc);
+                    Node::Branch(n) => {
+                        for (choice, child) in &mut n.choices.iter_mut().enumerate() {
+                            child.commit(path.append_new(choice as u8), acc);
                         }
                     }
-                    Node::Extension(node) => {
-                        node.child.commit(path.concat(&node.prefix), acc);
+                    Node::Extension(n) => {
+                        n.child.commit(path.concat(&n.prefix), acc);
                     }
                     Node::Leaf(_) => {}
                 }
                 let mut buf = Vec::new();
                 node.encode(&mut buf);
                 let hash = *hash.get_or_init(|| NodeHash::from_encoded(&buf));
                 if let Node::Leaf(leaf) = node.as_ref() {
-                    acc.push((path.concat(&leaf.partial), leaf.value.clone()));
+                    acc.push(TrieCommitEntry::LeafValue {
+                        path: path.concat(&leaf.partial),
+                        value: leaf.value.clone(),
+                    });
                 }
-                acc.push((path, buf));
+                acc.push(TrieCommitEntry::Node {
+                    path,
+                    node: node.clone(),
+                    encoded: buf,
+                });
 
                 hash
             }