feat(mpt): Trie DB commit (#196)

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]

crates/mpt/src/db/mod.rs

@@ -2,16 +2,18 @@
 //! 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::{
     db::{AccountState, DbAccount},
     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<B256> {
-        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<DbAccount> {
+    ///
+    /// # 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<DbAccount> {
         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<Bytes> + Copy,
     CHF: Fn(B256) -> Result<Bytes> + Copy,
 {
-    fn commit(&mut self, _: HashMap<Address, Account>) {
-        unimplemented!("TrieCacheDB::commit")
+    fn commit(&mut self, updated_accounts: HashMap<Address, Account>) {
+        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<Bytecode, Self::Error> {
         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<B256, Self::Error> {
-        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<B256, Self::Error> {
+        // 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.")
     }
 }

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,15 +298,15 @@ 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;
                 encoded_key_len + value.length()
             }
             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::<Vec<TrieNode>>();
+                let blinded_nodes = stack
+                    .iter()
+                    .cloned()
+                    .map(|mut node| {
+                        node.blind();
+                        node
+                    })
+                    .collect::<Vec<TrieNode>>();
                 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<T: Encodable>(value: T) -> usize {
     if value.length() > B256::ZERO.length() {
         B256::ZERO.length()
@@ -467,6 +477,10 @@ fn blinded_length<T: Encodable>(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<T: Encodable>(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<T: Encodable>(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<usize> {
     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);

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<T, F>(items: &[T], mut encode: F) -> HashBuilder
@@ -23,7 +23,7 @@ where
         })
         .collect::<Vec<_>>();
 
-    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);