From b0a8fdc391baa02e3ab3b8360055f967f36c6191 Mon Sep 17 00:00:00 2001 From: clabby Date: Sat, 1 Jun 2024 17:55:46 -0400 Subject: [PATCH] feat(mpt): Simplify `TrieDB` Simplifies `TrieDB` by assuming it will be wrapped with [`revm::State`](https://github.com/bluealloy/revm/blob/fadf6fbe43a3022c1406c0d8ee062a83e951d4d3/crates/revm/src/db/states/state.rs#L29). The `State` struct in `revm` is better for our usecase, since it allows for merging a `BundleState` after a full block's execution. The previous implementation had the fatal flaw having to recompute the state root every time the `DatabaseCommit` trait implementation was invoked, which was not suitable. Examples of usage may be found within the rustdoc for `TrieDB`. --- crates/mpt/src/db/mod.rs | 400 ++++++++++++++++----------------------- crates/mpt/src/lib.rs | 2 +- crates/mpt/src/node.rs | 42 ++-- 3 files changed, 191 insertions(+), 253 deletions(-) diff --git a/crates/mpt/src/db/mod.rs b/crates/mpt/src/db/mod.rs index 238dba83a..5bd4fdd06 100644 --- a/crates/mpt/src/db/mod.rs +++ b/crates/mpt/src/db/mod.rs @@ -3,53 +3,70 @@ 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, 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, + db::BundleState, + primitives::{AccountInfo, Bytecode, HashMap}, + Database, }; -use tracing::trace; mod account; pub use account::TrieAccount; -/// A Trie DB that caches account state in-memory. When accounts that don't already exist within the -/// cache are queried, the database fetches the preimages of the trie nodes on the path to the -/// account using the `PreimageFetcher` (`PF` generic) and `CodeHashFetcher` (`CHF` generic). This -/// allows for data to be fetched in a verifiable manner given an initial trusted state root as it -/// is needed during execution. +/// A Trie DB that caches open state in-memory. When accounts that don't already exist within the +/// cached [TrieNode] are queried, the database fetches the preimages of the trie nodes on the path +/// to the account using the `PreimageFetcher` (`PF` generic) and `CodeHashFetcher` (`CHF` generic). +/// This allows for data to be fetched in a verifiable manner given an initial trusted state root +/// as it is needed during execution. +/// +/// The [TrieDB] is intended to be wrapped by a [State], which is then used by the [revm::Evm] to +/// capture state transitions during block execution. /// /// **Behavior**: -/// - When an account is queried and it does not already exist in the inner cache database, we fall -/// through to the `PreimageFetcher` to fetch the preimages of the trie nodes on the path to the -/// account. After it has been fetched, the account is inserted into the cache database and will -/// be read from there on subsequent queries. +/// - When an account is queried and the trie path has not already been opened by [Self::basic], we +/// fall through to the `PreimageFetcher` to fetch the preimages of the trie nodes on the path to +/// the account. After it has been fetched, the path will be cached until the next call to +/// [Self::state_root]. /// - When querying for the code hash of an account, the `CodeHashFetcher` is consulted to fetch the /// code hash of the account. -/// - When a changeset is committed to the database, the changes are first applied to the cache -/// database and then the trie hash is recomputed. The root hash of the trie is then persisted to -/// the struct. +/// - When a [BundleState] changeset is committed to the parent [State] database, the changes are +/// first applied to the [State]'s cache, then the trie hash is recomputed with +/// [Self::state_root]. +/// +/// *Example Construction**: +/// ```rust +/// use alloy_primitives::{Bytes, B256}; +/// use anyhow::Result; +/// use kona_mpt::TrieDB; +/// use revm::{db::states::bundle_state::BundleRetention, EvmBuilder, StateBuilder}; +/// +/// let mock_fetcher = |hash: B256| -> Result { Ok(Default::default()) }; +/// let mock_starting_root = B256::default(); +/// +/// let trie_db = TrieDB::new(mock_starting_root, mock_fetcher, mock_fetcher); +/// let mut state = StateBuilder::new_with_database(trie_db).with_bundle_update().build(); +/// let evm = EvmBuilder::default().with_db(&mut state).build(); +/// +/// // Execute your block's transactions... +/// +/// // Drop the EVM prior to merging the state transitions. +/// drop(evm); +/// +/// state.merge_transitions(BundleRetention::PlainState); +/// let bundle = state.take_bundle(); +/// let state_root = state.database.state_root(&bundle).expect("Failed to compute state root"); +/// ``` /// -/// Note: This Database implementation intentionally wraps the [InMemoryDB], rather than serving as -/// a backing database for [revm::db::CacheDB]. This is because the [revm::db::CacheDB] is designed -/// to be a cache layer on top of a [revm::DatabaseRef] implementation, and the [TrieCacheDB] has a -/// requirement to also cache the opened state trie and account storage trie nodes, which requires -/// mutable access. +/// [State]: revm::State #[derive(Debug, Clone)] -pub struct TrieCacheDB +pub struct TrieDB where PF: Fn(B256) -> Result + Copy, CHF: Fn(B256) -> Result + Copy, { - /// The underlying DB that stores the account state in-memory. - db: InMemoryDB, - /// The root hash of the trie. - root: B256, /// The [TrieNode] representation of the root node. root_node: TrieNode, /// Storage roots of accounts within the trie. @@ -60,50 +77,21 @@ where code_by_hash_fetcher: CHF, } -impl TrieCacheDB +impl TrieDB where PF: Fn(B256) -> Result + Copy, CHF: Fn(B256) -> Result + Copy, { - /// Creates a new [TrieCacheDB] with the given root node. + /// Creates a new [TrieDB] with the given root node. pub fn new(root: B256, preimage_fetcher: PF, code_by_hash_fetcher: CHF) -> Self { Self { - db: InMemoryDB::default(), - root, - root_node: TrieNode::Blinded { commitment: root }, + root_node: TrieNode::new_blinded(root), preimage_fetcher, code_by_hash_fetcher, storage_roots: Default::default(), } } - /// Returns a reference to the underlying in-memory DB. - pub fn inner_db_ref(&self) -> &InMemoryDB { - &self.db - } - - /// Returns a mutable reference to the underlying in-memory DB. - pub fn inner_db_mut(&mut self) -> &mut InMemoryDB { - &mut self.db - } - - /// 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 { - trace!("Start state root update"); - self.root_node.blind(); - trace!("State root node updated successfully"); - - let commitment = if let TrieNode::Blinded { commitment } = self.root_node { - commitment - } else { - anyhow::bail!("Root node is not a blinded node") - }; - - self.root = commitment; - Ok(commitment) - } - /// Consumes `Self` and takes the the current state root of the trie DB. pub fn take_root_node(self) -> TrieNode { self.root_node @@ -141,71 +129,85 @@ where &mut self.storage_roots } - /// Loads an account from the trie by consulting the `PreimageFetcher` to fetch the preimages of - /// 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 + /// Applies a [BundleState] changeset to the [TrieNode] and recomputes the state root hash. /// /// # Takes - /// - `address`: The address of the account to load. + /// - `bundle`: The [BundleState] changeset to apply to the trie DB. /// /// # 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, self.preimage_fetcher)?; - let trie_account = TrieAccount::decode(&mut trie_account_rlp.as_ref()) - .map_err(|e| anyhow!("Error decoding trie account: {e}"))?; + /// - `Ok(B256)`: The new state root hash of the trie DB. + /// - `Err(_)`: If the state root hash could not be computed. + pub fn state_root(&mut self, bundle: &BundleState) -> Result { + // Update the accounts in the trie with the changeset. + self.update_accounts(bundle)?; - // Insert the account's storage root into the cache. - self.storage_roots - .insert(address, TrieNode::Blinded { commitment: trie_account.storage_root }); - - // If the account's code hash is not empty, fetch the bytecode and insert it into the cache. - let code = (trie_account.code_hash != KECCAK_EMPTY) - .then(|| { - let code = Bytecode::new_raw((self.code_by_hash_fetcher)(trie_account.code_hash)?); - Ok::<_, anyhow::Error>(code) - }) - .transpose()?; - - // Return a partial DB account. The storage and code are not loaded out-right, and are - // loaded optimistically in the `Database` + `DatabaseRef` trait implementations. - let mut info = AccountInfo { - balance: trie_account.balance, - nonce: trie_account.nonce, - code_hash: trie_account.code_hash, - code, - }; - self.insert_contract(&mut info); + // Recompute the root hash of the trie. + self.root_node.blind(); - Ok(DbAccount { info, ..Default::default() }) + // Extract the new state root from the root node. + self.root_node.blinded_commitment().ok_or(anyhow!("State root node is not a blinded node")) } - /// Inserts the account's code into the cache. - /// - /// 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. + /// Modifies the accounts in the storage trie with the given [BundleState] changeset. /// /// # 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 { - account.code_hash = code.hash_slow(); - } - self.db.contracts.entry(account.code_hash).or_insert_with(|| code.clone()); + /// - `bundle`: The [BundleState] changeset to apply to the trie DB. + /// + /// # Returns + /// - `Ok(())` if the accounts were successfully updated. + /// - `Err(_)` if the accounts could not be updated. + fn update_accounts(&mut self, bundle: &BundleState) -> Result<()> { + for (address, bundle_account) in bundle.state() { + let account_info = + bundle_account.account_info().ok_or(anyhow!("Account info not found"))?; + 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 + let acc_storage_root = self + .storage_roots + .get_mut(address) + .ok_or(anyhow!("Storage root not found for account"))?; + bundle_account.storage.iter().try_for_each(|(index, value)| { + Self::change_storage( + acc_storage_root, + *index, + value.present_value, + self.preimage_fetcher, + ) + })?; + + // Recompute the account storage root. + acc_storage_root.blind(); + + let commitment = acc_storage_root + .blinded_commitment() + .ok_or(anyhow!("Storage root node is not a blinded node"))?; + trie_account.storage_root = commitment; + + // RLP encode the trie account for insertion. + let mut account_buf = Vec::with_capacity(trie_account.length()); + trie_account.encode(&mut account_buf); + + // Insert or update the account in the trie. + let account_path = Nibbles::unpack(keccak256(address.as_slice())); + if let Ok(account_rlp_ref) = self.root_node.open(&account_path, self.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(), self.preimage_fetcher)?; } } - if account.code_hash == B256::ZERO { - account.code_hash = KECCAK_EMPTY; - } + + Ok(()) } - /// Modifies a storage slot of an account in the trie DB. + /// Modifies a storage slot of an account in the Merkle Patricia Trie. /// /// # Takes /// - `address`: The address of the account. @@ -215,23 +217,20 @@ where /// # 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, + fn change_storage( + storage_root: &mut TrieNode, index: U256, value: U256, + preimage_fetcher: PF, ) -> 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()); - + // RLP encode the storage slot value. let mut rlp_buf = Vec::with_capacity(value.length()); value.encode(&mut rlp_buf); + // Insert or update the storage slot in the trie. + let hashed_slot_key = keccak256(index.to_be_bytes::<32>().as_slice()); if let Ok(storage_slot_rlp) = - storage_root.open(&Nibbles::unpack(hashed_slot_key), self.preimage_fetcher) + storage_root.open(&Nibbles::unpack(hashed_slot_key), preimage_fetcher) { // If the storage slot already exists, update it. *storage_slot_rlp = rlp_buf.into(); @@ -240,7 +239,7 @@ where storage_root.insert( &Nibbles::unpack(hashed_slot_key), rlp_buf.into(), - self.preimage_fetcher, + preimage_fetcher, )?; } @@ -248,57 +247,7 @@ where } } -impl DatabaseCommit for TrieCacheDB -where - PF: Fn(B256) -> Result + Copy, - CHF: Fn(B256) -> Result + Copy, -{ - 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, 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(), 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"); - } -} - -impl Database for TrieCacheDB +impl Database for TrieDB where PF: Fn(B256) -> Result + Copy, CHF: Fn(B256) -> Result + Copy, @@ -306,82 +255,51 @@ where type Error = anyhow::Error; fn basic(&mut self, address: Address) -> Result, Self::Error> { - let basic = match self.db.accounts.entry(address) { - Entry::Occupied(entry) => entry.into_mut(), - Entry::Vacant(_) => { - let account = self.load_account_from_trie(address)?; - if let Some(ref code) = account.info.code { - self.db.contracts.insert(account.info.code_hash, code.clone()); - } - self.db.accounts.insert(address, account); - self.db - .accounts - .get_mut(&address) - .ok_or(anyhow!("Account not found in cache: {address}"))? - } + // Fetch the account from the trie. + let hashed_address_nibbles = Nibbles::unpack(keccak256(address.as_slice())); + let Ok(trie_account_rlp) = + self.root_node.open(&hashed_address_nibbles, self.preimage_fetcher) + else { + // If the account does not exist in the trie, return `Ok(None)`. + return Ok(None); }; - Ok(basic.info()) + + // Decode the trie account from the RLP bytes. + let trie_account = TrieAccount::decode(&mut trie_account_rlp.as_ref()) + .map_err(|e| anyhow!("Error decoding trie account: {e}"))?; + + // Insert the account's storage root into the cache. + self.storage_roots.insert(address, TrieNode::new_blinded(trie_account.storage_root)); + + // Return a partial DB account. The storage and code are not loaded out-right, and are + // loaded optimistically in the `Database` + `DatabaseRef` trait implementations. + Ok(Some(AccountInfo { + balance: trie_account.balance, + nonce: trie_account.nonce, + code_hash: trie_account.code_hash, + code: None, + })) } 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(_) => anyhow::bail!("Code hash not found in cache: {code_hash}"), - } + (self.code_by_hash_fetcher)(code_hash) + .map(Bytecode::new_raw) + .map_err(|e| anyhow!("Failed to fetch code by hash: {e}")) } fn storage(&mut self, address: Address, index: U256) -> Result { - match self.db.accounts.entry(address) { - Entry::Occupied(mut acc_entry) => { - let acc_entry = acc_entry.get_mut(); - match acc_entry.storage.entry(index) { - Entry::Occupied(entry) => Ok(*entry.get()), - Entry::Vacant(_) => { - if matches!( - acc_entry.account_state, - AccountState::StorageCleared | AccountState::NotExisting - ) { - Ok(U256::ZERO) - } else { - let fetcher = self.preimage_fetcher; - 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 slot_value = - storage_root.open(&Nibbles::unpack(hashed_slot_key), fetcher)?; - - let int_slot = U256::decode(&mut slot_value.as_ref()) - .map_err(|e| anyhow!("Failed to decode storage slot value: {e}"))?; - - self.db - .accounts - .get_mut(&address) - .ok_or(anyhow!("Account not found in cache: {address}"))? - .storage - .insert(index, int_slot); - Ok(int_slot) - } - } - } - } - Entry::Vacant(_) => { - // acc needs to be loaded for us to access slots. - let info = self.basic(address)?; - let (account, value) = if info.is_some() { - let value = self.storage(address, index)?; - let mut account: DbAccount = info.into(); - account.storage.insert(index, value); - (account, value) - } else { - (info.into(), U256::ZERO) - }; - self.db.accounts.insert(address, account); - Ok(value) - } - } + 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 slot_value = + storage_root.open(&Nibbles::unpack(hashed_slot_key), self.preimage_fetcher)?; + + let int_slot = U256::decode(&mut slot_value.as_ref()) + .map_err(|e| anyhow!("Failed to decode storage slot value: {e}"))?; + Ok(int_slot) } fn block_hash(&mut self, _: U256) -> Result { diff --git a/crates/mpt/src/lib.rs b/crates/mpt/src/lib.rs index 628f5a1bc..b942c5942 100644 --- a/crates/mpt/src/lib.rs +++ b/crates/mpt/src/lib.rs @@ -7,7 +7,7 @@ extern crate alloc; mod db; -pub use db::{TrieAccount, TrieCacheDB}; +pub use db::{TrieAccount, TrieDB}; mod node; pub use node::TrieNode; diff --git a/crates/mpt/src/node.rs b/crates/mpt/src/node.rs index 0d13d99db..fcb6974f3 100644 --- a/crates/mpt/src/node.rs +++ b/crates/mpt/src/node.rs @@ -82,6 +82,30 @@ pub enum TrieNode { } impl TrieNode { + /// Creates a new [TrieNode::Blinded] node. + /// + /// ## Takes + /// - `commitment` - The commitment that blinds the node + /// + /// ## Returns + /// - `Self` - The new blinded [TrieNode]. + pub fn new_blinded(commitment: B256) -> Self { + TrieNode::Blinded { commitment } + } + + /// Returns the commitment of a [TrieNode::Blinded] node, if `self` is of the + /// [TrieNode::Blinded] variant. + /// + /// ## Returns + /// - `Some(B256)` - The commitment of the blinded node + /// - `None` - `self` is not a [TrieNode::Blinded] node + pub fn blinded_commitment(&self) -> Option { + match self { + TrieNode::Blinded { commitment } => Some(*commitment), + _ => None, + } + } + /// Blinds the [TrieNode] if it is longer than an encoded [B256] string in length, and returns /// the mutated node. pub fn blind(&mut self) { @@ -469,7 +493,7 @@ impl Decodable for TrieNode { } let commitment = B256::decode(buf)?; - Ok(Self::Blinded { commitment }) + Ok(Self::new_blinded(commitment)) } } } @@ -502,7 +526,7 @@ fn encode_blinded(value: T, out: &mut dyn BufMut) { if value.length() > B256::ZERO.length() { let mut rlp_buf = Vec::with_capacity(value.length()); value.encode(&mut rlp_buf); - TrieNode::Blinded { commitment: keccak256(rlp_buf) }.encode(out); + TrieNode::new_blinded(keccak256(rlp_buf)).encode(out); } else { value.encode(out); } @@ -562,11 +586,9 @@ mod test { const BRANCH_RLP: [u8; 64] = hex!("f83ea0eb08a66a94882454bec899d3e82952dcc918ba4b35a09a84acd98019aef4345080808080808080cd308b8a746573742074687265658080808080808080"); let expected = TrieNode::Branch { stack: vec![ - TrieNode::Blinded { - commitment: b256!( - "eb08a66a94882454bec899d3e82952dcc918ba4b35a09a84acd98019aef43450" - ), - }, + TrieNode::new_blinded(b256!( + "eb08a66a94882454bec899d3e82952dcc918ba4b35a09a84acd98019aef43450" + )), TrieNode::Empty, TrieNode::Empty, TrieNode::Empty, @@ -624,7 +646,7 @@ mod test { let opened = TrieNode::Leaf { prefix: Nibbles::from_nibbles([0x00]), value: [0xFF; 64].into() }; opened.encode(&mut rlp_buf); - let blinded = TrieNode::Blinded { commitment: keccak256(&rlp_buf) }; + let blinded = TrieNode::new_blinded(keccak256(&rlp_buf)); rlp_buf.clear(); let opened_extension = @@ -674,9 +696,7 @@ mod test { } root_node.blind(); - let TrieNode::Blinded { commitment } = root_node else { - panic!("Expected blinded root node"); - }; + let commitment = root_node.blinded_commitment().unwrap(); assert_eq!(commitment, root); }