diff --git a/Cargo.lock b/Cargo.lock index 427a4bfce..084a6ee86 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -383,9 +383,9 @@ dependencies = [ [[package]] name = "alloy-trie" -version = "0.3.1" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "beb28aa4ecd32fdfa1b1bdd111ff7357dd562c6b2372694cf9e613434fcba659" +checksum = "03704f265cbbb943b117ecb5055fd46e8f41e7dc8a58b1aed20bcd40ace38c15" dependencies = [ "alloy-primitives", "alloy-rlp", diff --git a/crates/mpt/Cargo.toml b/crates/mpt/Cargo.toml index 8d6b6c365..f5bf74de3 100644 --- a/crates/mpt/Cargo.toml +++ b/crates/mpt/Cargo.toml @@ -18,7 +18,7 @@ alloy-consensus.workspace = true revm.workspace = true # External -alloy-trie = { version = "0.3.1", default-features = false } +alloy-trie = { version = "0.4.1", default-features = false } smallvec = "1.13" [dev-dependencies] diff --git a/crates/mpt/src/db/mod.rs b/crates/mpt/src/db/mod.rs index ab32002fb..58dfc5019 100644 --- a/crates/mpt/src/db/mod.rs +++ b/crates/mpt/src/db/mod.rs @@ -2,9 +2,10 @@ //! incremental updates through fetching node preimages on the fly during execution. use crate::TrieNode; +use alloc::vec::Vec; use alloy_consensus::constants::KECCAK_EMPTY; use alloy_primitives::{keccak256, Address, Bytes, B256, U256}; -use alloy_rlp::Decodable; +use alloy_rlp::{Decodable, Encodable}; use alloy_trie::Nibbles; use anyhow::{anyhow, Result}; use revm::{ @@ -12,6 +13,7 @@ use revm::{ primitives::{hash_map::Entry, Account, AccountInfo, Bytecode, HashMap}, Database, DatabaseCommit, InMemoryDB, }; +use tracing::trace; mod account; pub use account::TrieAccount; @@ -88,15 +90,16 @@ where /// Returns the current state root of the trie DB, and replaces the root node with the new /// blinded form. This action drops all the cached account state. pub fn state_root(&mut self) -> Result { - let blinded = self.root_node.clone().blind(); + trace!("Start state root update"); + self.root_node.blind(); + trace!("State root node updated successfully"); - let commitment = if let TrieNode::Blinded { commitment } = blinded { + let commitment = if let TrieNode::Blinded { commitment } = self.root_node { commitment } else { anyhow::bail!("Root node is not a blinded node") }; - self.root_node = blinded; self.root = commitment; Ok(commitment) } @@ -142,7 +145,14 @@ where /// the trie nodes on the path to the account. If the account has a non-empty storage trie /// root hash, the account's storage trie will be traversed to recover the account's storage /// slots. If the account has a non-empty - pub fn load_account_from_trie(&mut self, address: Address) -> Result { + /// + /// # Takes + /// - `address`: The address of the account to load. + /// + /// # Returns + /// - `Ok(DbAccount)`: The account loaded from the trie. + /// - `Err(_)`: If the account could not be loaded from the trie. + pub(crate) fn load_account_from_trie(&mut self, address: Address) -> Result { let hashed_address_nibbles = Nibbles::unpack(keccak256(address.as_slice())); let trie_account_rlp = self.root_node.open(&hashed_address_nibbles, 0, self.preimage_fetcher)?; @@ -178,7 +188,10 @@ where /// /// Accounts objects and code are stored separately in the cache, this will take the code from /// the account and instead map it to the code hash. - pub fn insert_contract(&mut self, account: &mut AccountInfo) { + /// + /// # Takes + /// - `account`: The account to insert the code for. + pub(crate) fn insert_contract(&mut self, account: &mut AccountInfo) { if let Some(code) = &account.code { if !code.is_empty() { if account.code_hash == KECCAK_EMPTY { @@ -192,9 +205,47 @@ where } } - /// Inserts a block hash into the cache. - pub fn insert_block_hash(&mut self, number: U256, hash: B256) { - self.db.block_hashes.insert(number, hash); + /// Modifies a storage slot of an account in the trie DB. + /// + /// # Takes + /// - `address`: The address of the account. + /// - `index`: The index of the storage slot. + /// - `value`: The new value of the storage slot. + /// + /// # Returns + /// - `Ok(())` if the storage slot was successfully modified. + /// - `Err(_)` if the storage slot could not be modified. + pub(crate) fn change_storage( + &mut self, + address: Address, + index: U256, + value: U256, + ) -> Result<()> { + let storage_root = self + .storage_roots + .get_mut(&address) + .ok_or(anyhow!("Storage root not found for account: {address}"))?; + let hashed_slot_key = keccak256(index.to_be_bytes::<32>().as_slice()); + + let mut rlp_buf = Vec::with_capacity(value.length()); + value.encode(&mut rlp_buf); + + if let Ok(storage_slot_rlp) = + storage_root.open(&Nibbles::unpack(hashed_slot_key), 0, self.preimage_fetcher) + { + // If the storage slot already exists, update it. + *storage_slot_rlp = rlp_buf.into(); + } else { + // If the storage slot does not exist, insert it. + storage_root.insert( + &Nibbles::unpack(hashed_slot_key), + rlp_buf.into(), + 0, + self.preimage_fetcher, + )?; + } + + Ok(()) } } @@ -203,8 +254,48 @@ where PF: Fn(B256) -> Result + Copy, CHF: Fn(B256) -> Result + Copy, { - fn commit(&mut self, _: HashMap) { - unimplemented!("TrieCacheDB::commit") + fn commit(&mut self, updated_accounts: HashMap) { + let preimage_fetcher = self.preimage_fetcher; + for (address, account) in updated_accounts { + let account_path = Nibbles::unpack(keccak256(address.as_slice())); + let mut trie_account = TrieAccount { + balance: account.info.balance, + nonce: account.info.nonce, + code_hash: account.info.code_hash, + ..Default::default() + }; + + // Update the account's storage root + for (index, value) in account.storage { + self.change_storage(address, index, value.present_value) + .expect("Failed to update account storage"); + } + let acc_storage_root = + self.storage_roots.get_mut(&address).expect("Storage root not found for account"); + acc_storage_root.blind(); + if let TrieNode::Blinded { commitment } = acc_storage_root { + trie_account.storage_root = *commitment; + } else { + panic!("Storage root was not blinded successfully"); + } + + // RLP encode the account. + let mut account_buf = Vec::with_capacity(trie_account.length()); + trie_account.encode(&mut account_buf); + + if let Ok(account_rlp_ref) = self.root_node.open(&account_path, 0, preimage_fetcher) { + // Update the existing account in the trie. + *account_rlp_ref = account_buf.into(); + } else { + // Insert the new account into the trie. + self.root_node + .insert(&account_path, account_buf.into(), 0, preimage_fetcher) + .expect("Failed to insert account into trie"); + } + } + + // Update the root hash of the trie. + self.state_root().expect("Failed to update state root"); } } @@ -224,7 +315,10 @@ where self.db.contracts.insert(account.info.code_hash, code.clone()); } self.db.accounts.insert(address, account); - self.db.accounts.get_mut(&address).unwrap() + self.db + .accounts + .get_mut(&address) + .ok_or(anyhow!("Account not found in cache: {address}"))? } }; Ok(basic.info()) @@ -233,7 +327,7 @@ where fn code_by_hash(&mut self, code_hash: B256) -> Result { match self.db.contracts.entry(code_hash) { Entry::Occupied(entry) => Ok(entry.get().clone()), - Entry::Vacant(_) => unreachable!("Code hash not found in cache: {code_hash}"), + Entry::Vacant(_) => anyhow::bail!("Code hash not found in cache: {code_hash}"), } } @@ -266,7 +360,7 @@ where self.db .accounts .get_mut(&address) - .expect("Must exist") + .ok_or(anyhow!("Account not found in cache: {address}"))? .storage .insert(index, int_slot); Ok(int_slot) @@ -291,10 +385,11 @@ where } } - fn block_hash(&mut self, number: U256) -> Result { - match self.db.block_hashes.entry(number) { - Entry::Occupied(entry) => Ok(*entry.get()), - Entry::Vacant(_) => anyhow::bail!("Block hash for number not found"), - } + fn block_hash(&mut self, _: U256) -> Result { + // match self.db.block_hashes.entry(number) { + // Entry::Occupied(entry) => Ok(*entry.get()), + // Entry::Vacant(_) => anyhow::bail!("Block hash for number not found"), + // } + unimplemented!("Block hash not implemented; Need to unroll the starting block hash for this operation.") } } diff --git a/crates/mpt/src/node.rs b/crates/mpt/src/node.rs index ece9a4699..9f84462e9 100644 --- a/crates/mpt/src/node.rs +++ b/crates/mpt/src/node.rs @@ -84,13 +84,11 @@ pub enum TrieNode { impl TrieNode { /// Blinds the [TrieNode] if it is longer than an encoded [B256] string in length, and returns /// the mutated node. - pub fn blind(self) -> Self { + pub fn blind(&mut self) { if self.length() > B256::ZERO.length() { let mut rlp_buf = Vec::with_capacity(self.length()); self.encode(&mut rlp_buf); - TrieNode::Blinded { commitment: keccak256(rlp_buf) } - } else { - self + *self = TrieNode::Blinded { commitment: keccak256(rlp_buf) } } } @@ -142,9 +140,9 @@ impl TrieNode { } } TrieNode::Leaf { prefix, value } => { - // If the key length is one, it only contains the prefix and no shared nibbles. - // Return the key and value. - if prefix.len() == 1 || nibble_offset + prefix.len() >= path.len() { + // If the key length is 0 or the shared nibbles overflow the remaining path, return + // the key and value. + if prefix.len() == 0 || nibble_offset + prefix.len() >= path.len() { return Ok(value); } @@ -300,7 +298,7 @@ impl TrieNode { /// Returns the RLP payload length of the [TrieNode]. pub(crate) fn payload_length(&self) -> usize { match self { - TrieNode::Empty => 1, + TrieNode::Empty => 0, TrieNode::Blinded { commitment } => commitment.len(), TrieNode::Leaf { prefix, value } => { let encoded_key_len = prefix.length() / 2 + 1; @@ -308,7 +306,7 @@ impl TrieNode { } TrieNode::Extension { prefix, node } => { let encoded_key_len = prefix.length() / 2 + 1; - encoded_key_len + node.length() + encoded_key_len + blinded_length(node) } TrieNode::Branch { stack } => { // In branch nodes, if an element is longer than an encoded 32 byte string, it is @@ -333,7 +331,7 @@ impl TrieNode { let path = Bytes::decode(buf).map_err(|e| anyhow!("Failed to decode: {e}"))?; let first_nibble = path[0] >> NIBBLE_WIDTH; let first = match first_nibble { - PREFIX_EXTENSION_ODD | PREFIX_LEAF_ODD => Some(path[0] & 0x0f), + PREFIX_EXTENSION_ODD | PREFIX_LEAF_ODD => Some(path[0] & 0x0F), PREFIX_EXTENSION_EVEN | PREFIX_LEAF_EVEN => None, _ => anyhow::bail!("Unexpected path identifier in high-order nibble"), }; @@ -350,7 +348,7 @@ impl TrieNode { }) } PREFIX_LEAF_EVEN | PREFIX_LEAF_ODD => { - // leaf node + // Leaf node let value = Bytes::decode(buf).map_err(|e| anyhow!("Failed to decode: {e}"))?; Ok(TrieNode::Leaf { prefix: unpack_path_to_nibbles(first, path[1..].as_ref()), @@ -385,8 +383,14 @@ impl Encodable for TrieNode { // In branch nodes, if an element is longer than 32 bytes in length, it is blinded. // Assuming we have an open trie node, we must re-hash the elements // that are longer than 32 bytes in length. - let blinded_nodes = - stack.iter().cloned().map(|node| node.blind()).collect::>(); + let blinded_nodes = stack + .iter() + .cloned() + .map(|mut node| { + node.blind(); + node + }) + .collect::>(); blinded_nodes.encode(out); } } @@ -457,6 +461,12 @@ impl Decodable for TrieNode { /// Returns the encoded length of an [Encodable] value, blinding it if it is longer than an encoded /// [B256] string in length. +/// +/// ## Takes +/// - `value` - The value to encode +/// +/// ## Returns +/// - `usize` - The encoded length of the value fn blinded_length(value: T) -> usize { if value.length() > B256::ZERO.length() { B256::ZERO.length() @@ -467,6 +477,10 @@ fn blinded_length(value: T) -> usize { /// Encodes a value into an RLP stream, blidning it with a [keccak256] commitment if it is longer /// than an encoded [B256] string in length. +/// +/// ## Takes +/// - `value` - The value to encode +/// - `out` - The RLP stream to write the encoded value to fn encode_blinded(value: T, out: &mut dyn BufMut) { if value.length() > B256::ZERO.length() { let mut rlp_buf = Vec::with_capacity(value.length()); @@ -479,6 +493,13 @@ fn encode_blinded(value: T, out: &mut dyn BufMut) { /// Walks through a RLP list's elements and returns the total number of elements in the list. /// Returns [alloy_rlp::Error::UnexpectedString] if the RLP stream is not a list. +/// +/// ## Takes +/// - `buf` - The RLP stream to walk through +/// +/// ## Returns +/// - `Ok(usize)` - The total number of elements in the list +/// - `Err(_)` - The RLP stream is not a list fn rlp_list_element_length(buf: &mut &[u8]) -> alloy_rlp::Result { let header = Header::decode(buf)?; if !header.list { @@ -635,7 +656,8 @@ mod test { assert_eq!(v, encoded_value.as_slice()); } - let TrieNode::Blinded { commitment } = root_node.blind() else { + root_node.blind(); + let TrieNode::Blinded { commitment } = root_node else { panic!("Expected blinded root node"); }; assert_eq!(commitment, root); diff --git a/crates/mpt/src/util.rs b/crates/mpt/src/util.rs index e7c15201c..7d691af6b 100644 --- a/crates/mpt/src/util.rs +++ b/crates/mpt/src/util.rs @@ -2,7 +2,7 @@ use alloc::vec::Vec; use alloy_rlp::{BufMut, Encodable}; -use alloy_trie::{HashBuilder, Nibbles}; +use alloy_trie::{proof::ProofRetainer, HashBuilder, Nibbles}; /// Compute a trie root of the collection of items with a custom encoder. pub fn ordered_trie_with_encoder(items: &[T], mut encode: F) -> HashBuilder @@ -23,7 +23,7 @@ where }) .collect::>(); - let mut hb = HashBuilder::default().with_proof_retainer(path_nibbles); + let mut hb = HashBuilder::default().with_proof_retainer(ProofRetainer::new(path_nibbles)); for i in 0..items_len { let index = adjust_index_for_rlp(i, items_len);