[WIP] Merkle root paralelization

Cargo.toml

@@ -1,33 +1,33 @@
 [workspace]
 members = [
-  "cmd/ef_tests/blockchain",
-  "cmd/ef_tests/state",
-  "cmd/ethrex",
-  "cmd/ethrex_replay",
-  "crates/blockchain",
-  "crates/blockchain/dev",
-  "crates/common",
-  "crates/common/rlp",
-  "crates/common/trie",
-  "crates/l2/",
-  "crates/l2/contracts",
-  "crates/l2/common",
-  "crates/l2/prover",
-  "crates/l2/prover/zkvm/interface",
-  "crates/l2/sdk",
-  "crates/l2/storage",
-  "crates/l2/networking/rpc",
-  "crates/networking/p2p",
-  "crates/networking/rpc",
-  "crates/storage",
-  "crates/vm",
-  "crates/vm/levm",
-  "crates/vm/levm/bench/revm_comparison",
-  "tooling/genesis",
-  "tooling/hive_report",
-  "tooling/load_test",
-  "tooling/loc",
-  "tooling/archive_sync"
+    "cmd/ef_tests/blockchain",
+    "cmd/ef_tests/state",
+    "cmd/ethrex",
+    "cmd/ethrex_replay",
+    "crates/blockchain",
+    "crates/blockchain/dev",
+    "crates/common",
+    "crates/common/rlp",
+    "crates/common/trie",
+    "crates/l2/",
+    "crates/l2/contracts",
+    "crates/l2/common",
+    "crates/l2/prover",
+    "crates/l2/prover/zkvm/interface",
+    "crates/l2/sdk",
+    "crates/l2/storage",
+    "crates/l2/networking/rpc",
+    "crates/networking/p2p",
+    "crates/networking/rpc",
+    "crates/storage",
+    "crates/vm",
+    "crates/vm/levm",
+    "crates/vm/levm/bench/revm_comparison",
+    "tooling/genesis",
+    "tooling/hive_report",
+    "tooling/load_test",
+    "tooling/loc",
+    "tooling/archive_sync",
 ]
 resolver = "2"
 
@@ -86,9 +86,9 @@ redb = "=2.4.0"
 snap = "1.1.1"
 k256 = { version = "0.13.3", features = ["ecdh"] }
 secp256k1 = { version = "0.29.1", default-features = false, features = [
-  "global-context",
-  "recovery",
-  "rand",
+    "global-context",
+    "recovery",
+    "rand",
 ] }
 keccak-hash = "0.11.0"
 axum = "0.8.1"
@@ -100,9 +100,10 @@ libsql = "0.9.10"
 futures = "0.3.31"
 # Changing the tag for spawned will break the TDX image build
 # When updating it try to build the TDX image and update service.nix with the new hash
-spawned-concurrency = {git = "https://github.com/lambdaclass/spawned.git", tag = "v0.1.2-alpha"}
-spawned-rt = {git = "https://github.com/lambdaclass/spawned.git", tag = "v0.1.2-alpha"}
+spawned-concurrency = { git = "https://github.com/lambdaclass/spawned.git", tag = "v0.1.2-alpha" }
+spawned-rt = { git = "https://github.com/lambdaclass/spawned.git", tag = "v0.1.2-alpha" }
 lambdaworks-crypto = "0.11.0"
+rayon = "1.5"
 tui-logger = { version = "0.17.3", features = ["tracing-support"] }
 
 [patch.crates-io]

cmd/ethrex/l2/command.rs

@@ -382,7 +382,7 @@ impl Command {
                             // Apply all account updates to trie
                             let account_updates = state_diff.to_account_updates(&new_trie)?;
                             let account_updates_list = store
-                                .apply_account_updates_from_trie_batch(new_trie, account_updates.values())
+                                .apply_account_updates_from_trie_batch(new_trie, &account_updates.values().cloned().collect::<Vec<_>>())
                                 .await
                                 .expect("Error applying account updates");
 

crates/common/Cargo.toml

@@ -27,7 +27,7 @@ crc32fast.workspace = true
 bytes.workspace = true
 hex.workspace = true
 lazy_static.workspace = true
-rayon = "1.5"
+rayon.workspace = true
 
 [dev-dependencies]
 hex-literal.workspace = true
@@ -39,5 +39,5 @@ c-kzg = ["dep:c-kzg"]
 [lib]
 path = "./common.rs"
 
-[lints.clippy] 
+[lints.clippy]
 unwrap_used = "deny"

crates/common/trie/Cargo.toml

@@ -19,6 +19,7 @@ libmdbx = { workspace = true, optional = true }
 smallvec = { version = "1.10.0", features = ["const_generics", "union"] }
 digest = "0.10.6"
 lazy_static.workspace = true
+rayon.workspace = true
 
 [features]
 default = []
@@ -40,4 +41,3 @@ path = "./trie.rs"
 [[bench]]
 name = "trie_bench"
 harness = false
-

crates/common/trie/node/branch.rs

@@ -1,8 +1,11 @@
 use ethrex_rlp::structs::Encoder;
 
-use crate::{TrieDB, ValueRLP, error::TrieError, nibbles::Nibbles, node_hash::NodeHash};
+use crate::NodeRef;
+use crate::{TrieDB, ValueRLP, error::TrieError, nibbles::Nibbles};
 
-use super::{ExtensionNode, LeafNode, Node, NodeRef, ValueOrHash};
+use super::{
+    extension::ExtensionNode, leaf::LeafNode, node::Node, node::ValueOrHash, node_hash::NodeHash,
+};
 
 /// Branch Node of an an Ethereum Compatible Patricia Merkle Trie
 /// Contains the node's value and the hash of its children nodes

crates/common/trie/node/extension.rs

@@ -1,11 +1,11 @@
 use ethrex_rlp::structs::Encoder;
 
+use crate::NodeRef;
 use crate::ValueRLP;
 use crate::nibbles::Nibbles;
-use crate::node_hash::NodeHash;
 use crate::{TrieDB, error::TrieError};
 
-use super::{BranchNode, Node, NodeRef, ValueOrHash};
+use super::{branch::BranchNode, node::Node, node::ValueOrHash, node_hash::NodeHash};
 
 /// Extension Node of an an Ethereum Compatible Patricia Merkle Trie
 /// Contains the node's prefix and a its child node hash, doesn't store any value
@@ -184,8 +184,9 @@ impl ExtensionNode {
 
 #[cfg(test)]
 mod test {
+    use super::super::node::LeafNode;
     use super::*;
-    use crate::{Trie, node::LeafNode, pmt_node};
+    use crate::{Trie, pmt_node};
 
     #[test]
     fn new() {

crates/common/trie/node/leaf.rs

@@ -1,8 +1,11 @@
 use ethrex_rlp::structs::Encoder;
 
-use crate::{ValueRLP, error::TrieError, nibbles::Nibbles, node::BranchNode, node_hash::NodeHash};
+use super::{
+    branch::BranchNode, extension::ExtensionNode, node::Node, node::ValueOrHash,
+    node_hash::NodeHash,
+};
+use crate::{ValueRLP, error::TrieError, nibbles::Nibbles};
 
-use super::{ExtensionNode, Node, ValueOrHash};
 /// Leaf Node of an an Ethereum Compatible Patricia Merkle Trie
 /// Contains the node's hash, value & path
 #[derive(Debug, Clone, Default, PartialEq)]

crates/common/trie/node/mod.rs

@@ -0,0 +1,8 @@
+pub mod branch;
+pub mod extension;
+pub mod leaf;
+#[allow(clippy::module_inception)]
+pub mod node;
+pub mod node_hash;
+pub mod node_ref;
+pub mod rlp;

crates/common/trie/node.rs → crates/common/trie/node/node.rs

@@ -1,114 +1,13 @@
-mod branch;
-mod extension;
-mod leaf;
+use std::array;
 
-use std::{
-    array,
-    sync::{Arc, OnceLock},
-};
+pub use super::branch::BranchNode;
+pub use super::extension::ExtensionNode;
+pub use super::leaf::LeafNode;
+use ethrex_rlp::{decode::decode_bytes, error::RLPDecodeError, structs::Decoder};
 
-pub use branch::BranchNode;
-use ethrex_rlp::{
-    decode::{RLPDecode, decode_bytes},
-    encode::RLPEncode,
-    error::RLPDecodeError,
-    structs::Decoder,
-};
-pub use extension::ExtensionNode;
-pub use leaf::LeafNode;
+use crate::{TrieDB, ValueRLP, error::TrieError, nibbles::Nibbles};
 
-use crate::{TrieDB, error::TrieError, nibbles::Nibbles};
-
-use super::{ValueRLP, node_hash::NodeHash};
-
-/// A reference to a node.
-#[derive(Clone, Debug)]
-pub enum NodeRef {
-    /// The node is embedded within the reference.
-    Node(Arc<Node>, OnceLock<NodeHash>),
-    /// The node is in the database, referenced by its hash.
-    Hash(NodeHash),
-}
-
-impl NodeRef {
-    pub fn get_node(&self, db: &dyn TrieDB) -> Result<Option<Node>, TrieError> {
-        match *self {
-            NodeRef::Node(ref node, _) => Ok(Some(node.as_ref().clone())),
-            NodeRef::Hash(NodeHash::Inline((data, len))) => {
-                Ok(Some(Node::decode_raw(&data[..len as usize])?))
-            }
-            NodeRef::Hash(hash @ NodeHash::Hashed(_)) => db
-                .get(hash)?
-                .map(|rlp| Node::decode(&rlp).map_err(TrieError::RLPDecode))
-                .transpose(),
-        }
-    }
-
-    pub fn is_valid(&self) -> bool {
-        match self {
-            NodeRef::Node(_, _) => true,
-            NodeRef::Hash(hash) => hash.is_valid(),
-        }
-    }
-
-    pub fn commit(&mut self, acc: &mut Vec<(NodeHash, Vec<u8>)>) -> NodeHash {
-        match *self {
-            NodeRef::Node(ref mut node, ref mut hash) => {
-                match Arc::make_mut(node) {
-                    Node::Branch(node) => {
-                        for node in &mut node.choices {
-                            node.commit(acc);
-                        }
-                    }
-                    Node::Extension(node) => {
-                        node.child.commit(acc);
-                    }
-                    Node::Leaf(_) => {}
-                }
-
-                let hash = hash.get_or_init(|| node.compute_hash());
-                acc.push((*hash, node.encode_to_vec()));
-
-                let hash = *hash;
-                *self = hash.into();
-
-                hash
-            }
-            NodeRef::Hash(hash) => hash,
-        }
-    }
-
-    pub fn compute_hash(&self) -> NodeHash {
-        match self {
-            NodeRef::Node(node, hash) => *hash.get_or_init(|| node.compute_hash()),
-            NodeRef::Hash(hash) => *hash,
-        }
-    }
-}
-
-impl Default for NodeRef {
-    fn default() -> Self {
-        Self::Hash(NodeHash::default())
-    }
-}
-
-impl From<Node> for NodeRef {
-    fn from(value: Node) -> Self {
-        Self::Node(Arc::new(value), OnceLock::new())
-    }
-}
-
-impl From<NodeHash> for NodeRef {
-    fn from(value: NodeHash) -> Self {
-        Self::Hash(value)
-    }
-}
-
-impl PartialEq for NodeRef {
-    fn eq(&self, other: &Self) -> bool {
-        self.compute_hash() == other.compute_hash()
-    }
-}
+use super::node_hash::NodeHash;
 
 pub enum ValueOrHash {
     Value(ValueRLP),
@@ -254,14 +153,14 @@ impl Node {
                     // Decode as Extension
                     ExtensionNode {
                         prefix: path,
-                        child: decode_child(&rlp_items[1]).into(),
+                        child: NodeHash::decode_child(&rlp_items[1]).into(),
                     }
                     .into()
                 }
             }
             // Branch Node
             17 => {
-                let choices = array::from_fn(|i| decode_child(&rlp_items[i]).into());
+                let choices = array::from_fn(|i| NodeHash::decode_child(&rlp_items[i]).into());
                 let (value, _) = decode_bytes(&rlp_items[16])?;
                 BranchNode {
                     choices,
@@ -286,11 +185,3 @@ impl Node {
         }
     }
 }
-
-fn decode_child(rlp: &[u8]) -> NodeHash {
-    match decode_bytes(rlp) {
-        Ok((hash, &[])) if hash.len() == 32 => NodeHash::from_slice(hash),
-        Ok((&[], &[])) => NodeHash::default(),
-        _ => NodeHash::from_slice(rlp),
-    }
-}

crates/common/trie/node_hash.rs → crates/common/trie/node/node_hash.rs

@@ -1,5 +1,10 @@
 use ethereum_types::H256;
-use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode, error::RLPDecodeError, structs::Encoder};
+use ethrex_rlp::{
+    decode::{RLPDecode, decode_bytes},
+    encode::RLPEncode,
+    error::RLPDecodeError,
+    structs::Encoder,
+};
 #[cfg(feature = "libmdbx")]
 use libmdbx::orm::{Decodable, Encodable};
 use sha3::{Digest, Keccak256};
@@ -96,6 +101,14 @@ impl NodeHash {
             NodeHash::Inline(value) => value.1 == 0,
         }
     }
+
+    pub fn decode_child(rlp: &[u8]) -> Self {
+        match decode_bytes(rlp) {
+            Ok((hash, &[])) if hash.len() == 32 => Self::from_slice(hash),
+            Ok((&[], &[])) => Self::default(),
+            _ => Self::from_slice(rlp),
+        }
+    }
 }
 
 impl From<H256> for NodeHash {