diff --git a/Cargo.lock b/Cargo.lock index f4fa9b8096..c3a0198c95 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4319,6 +4319,7 @@ dependencies = [ "hex", "hex-literal", "libmdbx", + "rayon", "redb", "serde", "serde_json", @@ -4369,6 +4370,7 @@ dependencies = [ "libmdbx", "proptest", "rand 0.8.5", + "rayon", "serde", "serde_json", "sha3", diff --git a/Cargo.toml b/Cargo.toml index 73e5567257..a7d6865301 100644 --- a/Cargo.toml +++ b/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] diff --git a/cmd/ethrex/l2/command.rs b/cmd/ethrex/l2/command.rs index efd0731008..8ceb93e488 100644 --- a/cmd/ethrex/l2/command.rs +++ b/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::>()) .await .expect("Error applying account updates"); diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index 4b912e331a..ba508ba37a 100644 --- a/crates/common/Cargo.toml +++ b/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" diff --git a/crates/common/trie/Cargo.toml b/crates/common/trie/Cargo.toml index aa0b387b36..4af608f5ea 100644 --- a/crates/common/trie/Cargo.toml +++ b/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 - diff --git a/crates/common/trie/node/branch.rs b/crates/common/trie/node/branch.rs index 685930b817..2ccc9afb17 100644 --- a/crates/common/trie/node/branch.rs +++ b/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 diff --git a/crates/common/trie/node/extension.rs b/crates/common/trie/node/extension.rs index 9878a6f5b4..27fc3c7afe 100644 --- a/crates/common/trie/node/extension.rs +++ b/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() { diff --git a/crates/common/trie/node/leaf.rs b/crates/common/trie/node/leaf.rs index dce9a34416..31fea1bb50 100644 --- a/crates/common/trie/node/leaf.rs +++ b/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)] diff --git a/crates/common/trie/node/mod.rs b/crates/common/trie/node/mod.rs new file mode 100644 index 0000000000..6c32fe3652 --- /dev/null +++ b/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; diff --git a/crates/common/trie/node.rs b/crates/common/trie/node/node.rs similarity index 61% rename from crates/common/trie/node.rs rename to crates/common/trie/node/node.rs index 8818c54f8a..8dcaf552ab 100644 --- a/crates/common/trie/node.rs +++ b/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, OnceLock), - /// The node is in the database, referenced by its hash. - Hash(NodeHash), -} - -impl NodeRef { - pub fn get_node(&self, db: &dyn TrieDB) -> Result, 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)>) -> 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 for NodeRef { - fn from(value: Node) -> Self { - Self::Node(Arc::new(value), OnceLock::new()) - } -} - -impl From 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), - } -} diff --git a/crates/common/trie/node_hash.rs b/crates/common/trie/node/node_hash.rs similarity index 92% rename from crates/common/trie/node_hash.rs rename to crates/common/trie/node/node_hash.rs index 26e12f41c1..5f85ad2456 100644 --- a/crates/common/trie/node_hash.rs +++ b/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 for NodeHash { diff --git a/crates/common/trie/node/node_ref.rs b/crates/common/trie/node/node_ref.rs new file mode 100644 index 0000000000..2d077a71e7 --- /dev/null +++ b/crates/common/trie/node/node_ref.rs @@ -0,0 +1,101 @@ +use crate::Node; +use std::sync::{Arc, OnceLock}; + +use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode}; +use rayon::prelude::*; + +use crate::{TrieDB, error::TrieError}; + +use super::node_hash::NodeHash; + +/// A reference to a node. +#[derive(Clone, Debug)] +pub enum NodeRef { + /// The node is embedded within the reference. + Node(Arc, OnceLock), + /// The node is in the database, referenced by its hash. + Hash(NodeHash), +} + +impl NodeRef { + pub fn get_node(&self, db: &dyn TrieDB) -> Result, 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(), + } + } + + /// Returns the hash of the node, computing it if necessary. + pub fn commit(&mut self, acc: &mut Vec<(NodeHash, Vec)>) -> NodeHash { + match *self { + NodeRef::Node(ref mut node, ref mut hash) => { + match Arc::make_mut(node) { + Node::Branch(node) => { + acc.par_extend(node.choices[..].par_iter_mut().flat_map(|child| { + let mut acc = Vec::new(); + child.commit(&mut acc); + 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 for NodeRef { + fn from(value: Node) -> Self { + Self::Node(Arc::new(value), OnceLock::new()) + } +} + +impl From 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() + } +} diff --git a/crates/common/trie/rlp.rs b/crates/common/trie/node/rlp.rs similarity index 98% rename from crates/common/trie/rlp.rs rename to crates/common/trie/node/rlp.rs index 40e6694137..b5eb4892aa 100644 --- a/crates/common/trie/rlp.rs +++ b/crates/common/trie/node/rlp.rs @@ -8,7 +8,8 @@ use ethrex_rlp::{ }; use super::node::{BranchNode, ExtensionNode, LeafNode, Node}; -use crate::{NodeHash, node::NodeRef}; +use super::node_ref::NodeRef; +use crate::NodeHash; enum NodeType { Branch = 0, diff --git a/crates/common/trie/test_utils.rs b/crates/common/trie/test_utils.rs index 2094a13cf7..41ca89733c 100644 --- a/crates/common/trie/test_utils.rs +++ b/crates/common/trie/test_utils.rs @@ -7,10 +7,10 @@ macro_rules! pmt_node { branch { $( $choice:expr => $child_type:ident { $( $child_tokens:tt )* } ),+ $(,)? } $( offset $offset:expr )? ) => { - $crate::node::BranchNode::new({ + $crate::node::node::BranchNode::new({ #[allow(unused_variables)] let offset = true $( ^ $offset )?; - let mut choices = $crate::node::BranchNode::EMPTY_CHOICES; + let mut choices = $crate::node::node::BranchNode::EMPTY_CHOICES; $( let child_node: Node = pmt_node! { @($trie) $child_type { $( $child_tokens )* } @@ -27,12 +27,12 @@ macro_rules! pmt_node { with_leaf { $path:expr => $value:expr } $( offset $offset:expr )? ) => {{ - $crate::node::BranchNode::new_with_value({ + $crate::node::node::BranchNode::new_with_value({ #[allow(unused_variables)] let offset = true $( ^ $offset )?; - let mut choices = $crate::node::BranchNode::EMPTY_CHOICES; + let mut choices = $crate::node::node::BranchNode::EMPTY_CHOICES; $( - choices[$choice as usize] = $crate::node::Node::from( + choices[$choice as usize] = $crate::node::node::Node::from( pmt_node! { @($trie) $child_type { $( $child_tokens )* } offset offset @@ -50,10 +50,10 @@ macro_rules! pmt_node { #[allow(unused_variables)] let prefix = $crate::nibbles::Nibbles::from_hex($prefix.to_vec()); - $crate::node::ExtensionNode::new( + $crate::node::node::ExtensionNode::new( prefix.clone(), { - let child_node = $crate::node::Node::from(pmt_node! { @($trie) + let child_node = $crate::node::node::Node::from(pmt_node! { @($trie) $child_type { $( $child_tokens )* } }); child_node.into() @@ -67,7 +67,7 @@ macro_rules! pmt_node { $( offset $offset:expr )? ) => { { - $crate::node::LeafNode::new(Nibbles::from_hex($path), $value) + $crate::node::node::LeafNode::new(Nibbles::from_hex($path), $value) } }; } diff --git a/crates/common/trie/trie.rs b/crates/common/trie/trie.rs index 7011a3d144..cfe6006d89 100644 --- a/crates/common/trie/trie.rs +++ b/crates/common/trie/trie.rs @@ -1,10 +1,9 @@ pub mod db; pub mod error; pub mod logger; -mod nibbles; -mod node; -mod node_hash; -mod rlp; +pub mod nibbles; +pub mod node; + #[cfg(test)] mod test_utils; mod trie_iter; @@ -19,13 +18,10 @@ pub use self::db::{InMemoryTrieDB, TrieDB}; pub use self::logger::{TrieLogger, TrieWitness}; pub use self::nibbles::Nibbles; pub use self::verify_range::verify_range; -pub use self::{ - node::{Node, NodeRef}, - node_hash::NodeHash, -}; +pub use node::{node::LeafNode, node::Node, node_hash::NodeHash, node_ref::NodeRef}; pub use self::error::TrieError; -use self::{node::LeafNode, trie_iter::TrieIterator}; +use self::trie_iter::TrieIterator; use ethrex_rlp::decode::RLPDecode; use lazy_static::lazy_static; diff --git a/crates/common/trie/trie_iter.rs b/crates/common/trie/trie_iter.rs index 40b87806ad..defc9f7941 100644 --- a/crates/common/trie/trie_iter.rs +++ b/crates/common/trie/trie_iter.rs @@ -1,7 +1,5 @@ use crate::{ - PathRLP, Trie, TrieDB, ValueRLP, - nibbles::Nibbles, - node::{Node, NodeRef}, + PathRLP, Trie, TrieDB, ValueRLP, nibbles::Nibbles, node::node::Node, node::node_ref::NodeRef, }; pub struct TrieIterator { diff --git a/crates/common/trie/verify_range.rs b/crates/common/trie/verify_range.rs index 3b44edd32f..ae97087fd5 100644 --- a/crates/common/trie/verify_range.rs +++ b/crates/common/trie/verify_range.rs @@ -7,10 +7,8 @@ use ethereum_types::H256; use sha3::{Digest, Keccak256}; use crate::{ - ProofTrie, Trie, TrieError, ValueRLP, - nibbles::Nibbles, - node::{Node, NodeRef}, - node_hash::NodeHash, + ProofTrie, Trie, TrieError, ValueRLP, nibbles::Nibbles, node::node::Node, + node::node_hash::NodeHash, node::node_ref::NodeRef, }; /// Verifies that the key value range belongs to the trie with the given root given the edge proofs for the range diff --git a/crates/storage/Cargo.toml b/crates/storage/Cargo.toml index ae1b2e92ff..4e1708f6d4 100644 --- a/crates/storage/Cargo.toml +++ b/crates/storage/Cargo.toml @@ -18,13 +18,16 @@ tracing.workspace = true thiserror.workspace = true sha3.workspace = true hex.workspace = true +rayon.workspace = true serde = { version = "1.0.203", features = ["derive"] } serde_json = "1.0.117" libmdbx = { workspace = true, optional = true } redb = { workspace = true, optional = true } # NOTE: intentionally avoiding the workspace dep as it brings "full" features, breaking the provers # We only need the runtime for the blocking databases to spawn blocking tasks -tokio = { version = "1.41.1", optional = true, default-features = false, features = ["rt"] } +tokio = { version = "1.41.1", optional = true, default-features = false, features = [ + "rt", +] } bincode = "1.3.3" [features] diff --git a/crates/storage/store.rs b/crates/storage/store.rs index f0e83e8c33..d587842e26 100644 --- a/crates/storage/store.rs +++ b/crates/storage/store.rs @@ -19,6 +19,7 @@ use ethrex_common::{ use ethrex_rlp::decode::RLPDecode; use ethrex_rlp::encode::RLPEncode; use ethrex_trie::{Nibbles, NodeHash, Trie, TrieLogger, TrieNode, TrieWitness}; +use rayon::prelude::*; use sha3::{Digest as _, Keccak256}; use std::collections::{BTreeMap, HashMap}; use std::fmt::Debug; @@ -366,40 +367,54 @@ impl Store { pub async fn apply_account_updates_from_trie_batch( &self, mut state_trie: Trie, - account_updates: impl IntoIterator, + account_updates: &[AccountUpdate], ) -> Result { - let mut ret_storage_updates = Vec::new(); - let mut code_updates = Vec::new(); - for update in account_updates { - let hashed_address = hash_address(&update.address); - if update.removed { - // Remove account from trie - state_trie.remove(hashed_address)?; - continue; - } - // Add or update AccountState in the trie - // Fetch current state or create a new state to be inserted - let mut account_state = match state_trie.get(&hashed_address)? { - Some(encoded_state) => AccountState::decode(&encoded_state)?, - None => AccountState::default(), - }; - if let Some(info) = &update.info { - account_state.nonce = info.nonce; - account_state.balance = info.balance; - account_state.code_hash = info.code_hash; - // Store updated code in DB - if let Some(code) = &update.code { - code_updates.push((info.code_hash, code.clone())); + let mut hashed_addresses = Vec::with_capacity(account_updates.len()); + account_updates + .par_iter() + .map(|update| hash_address_fixed(&update.address)) + .collect_into_vec(&mut hashed_addresses); + + // contracts bytecode to update + let code_updates: Vec<_> = account_updates + .iter() + .filter_map(|update| { + update + .info + .clone() + .map(|info| info.code_hash) + .zip(update.code.clone()) + }) + .collect(); + + let mut account_states: Vec = hashed_addresses + .par_iter() + .map(|address| -> Result<_, StoreError> { + Ok(state_trie + .get(&address.0.to_vec())? + .map(|encoded_state| AccountState::decode(&encoded_state)) + .transpose()? + .unwrap_or_default()) + }) + .collect::>()?; + + let storage_updates: Vec<_> = hashed_addresses + .par_iter() + .zip_eq(account_states.par_iter_mut()) + .zip_eq(account_updates.par_iter()) + .map(|((address, state), update)| -> Result<_, StoreError> { + if update.added_storage.is_empty() { + return Ok((*address, Vec::new())); } - } - // Store the added storage in the account's storage trie and compute its new root - if !update.added_storage.is_empty() { - let mut storage_trie = self.engine.open_storage_trie( - H256::from_slice(&hashed_address), - account_state.storage_root, - )?; - for (storage_key, storage_value) in &update.added_storage { - let hashed_key = hash_key(storage_key); + let mut storage_trie = self + .engine + .open_storage_trie(*address, state.storage_root)?; + let storage_updates: Vec<_> = update + .added_storage + .par_iter() + .map(|(storage_key, storage_value)| (hash_key(storage_key), storage_value)) + .collect(); + for (hashed_key, storage_value) in storage_updates { if storage_value.is_zero() { storage_trie.remove(hashed_key)?; } else { @@ -408,17 +423,34 @@ impl Store { } let (storage_hash, storage_updates) = storage_trie.collect_changes_since_last_hash(); - account_state.storage_root = storage_hash; - ret_storage_updates.push((H256::from_slice(&hashed_address), storage_updates)); + state.storage_root = storage_hash; + Ok((*address, storage_updates)) + }) + .collect::>()?; + + for ((address, state), update) in hashed_addresses + .iter() + .zip(account_states.iter_mut()) + .zip(account_updates.iter()) + { + if let Some(info) = &update.info { + state.nonce = info.nonce; + state.code_hash = info.code_hash; + state.balance = info.balance; + } + if update.removed { + state_trie.remove(address.0.to_vec())?; + } else { + state_trie.insert(address.0.to_vec(), state.encode_to_vec())?; } - state_trie.insert(hashed_address, account_state.encode_to_vec())?; } + let (state_trie_hash, state_updates) = state_trie.collect_changes_since_last_hash(); Ok(AccountUpdatesList { state_trie_hash, state_updates, - storage_updates: ret_storage_updates, + storage_updates, code_updates, }) }