diff --git a/src/storage.rs b/src/storage.rs index 2638784e..294a2669 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -1,6 +1,7 @@ pub mod debug; pub mod engine; -pub mod overlay_root; +pub mod overlay; +pub mod overlay_traversal; pub mod proofs; mod test_utils; pub mod value; diff --git a/src/storage/overlay.rs b/src/storage/overlay.rs new file mode 100644 index 00000000..36cfd6a7 --- /dev/null +++ b/src/storage/overlay.rs @@ -0,0 +1,939 @@ +use crate::{ + context::TransactionContext, + overlay::OverlayState, + path::{AddressPath, RawPath, StoragePath}, + storage::{ + proofs::AccountProof, + engine::{Error, StorageEngine}, + overlay_traversal::OverlayTrie + }, +}; +use alloy_primitives::{map::{HashMap, B256Map}, B256}; +use alloy_trie::{Nibbles, BranchNodeCompact}; +use std::collections::HashSet; + +/// Represents the result of applying an overlay to the state root. +/// Contains the new root hash and any updated branch nodes. +#[derive(Debug)] +pub struct OverlayedRoot { + pub root: B256, + pub updated_branch_nodes: HashMap, + pub storage_branch_updates: B256Map>, +} + +impl OverlayedRoot { + pub fn new( + root: B256, + updated_branch_nodes: HashMap, + storage_branch_updates: B256Map>, + ) -> Self { + Self { root, updated_branch_nodes, storage_branch_updates } + } + + pub fn new_hash(root: B256) -> Self { + Self { + root, + updated_branch_nodes: HashMap::default(), + storage_branch_updates: B256Map::default(), + } + } +} + +impl StorageEngine { + /// Creates an `OverlayTrie` builder for applying overlay state to the persistent trie. + /// + /// This is a low-level builder that allows fine-grained control over the overlay + /// traversal process, including specifying proof targets. For most use cases, + /// prefer the higher-level methods like `compute_state_root_with_overlay()`, + /// `compute_account_with_proof_with_overlay()` or `compute_storage_with_proof_with_overlay()`. + /// + /// # Arguments + /// * `context` - Transaction context containing the base state + /// * `overlay` - Overlay state containing pending modifications + /// + /// # Returns + /// An `OverlayTrie` builder that can be configured and executed + /// + /// # Example + /// ```ignore + /// let trie = engine.overlay_trie(context, overlay) + /// .with_proof_targets(targets) + /// .compute_root()?; + /// ``` + pub fn overlay_trie<'a>( + &'a self, + context: &'a TransactionContext, + overlay: OverlayState, + ) -> OverlayTrie<'a> { + OverlayTrie::new(self, context, overlay) + } + + /// Computes a new state root by applying overlay state to the persistent trie. + /// + /// This method merges the uncommitted overlay state with the persistent storage + /// to compute the resulting state root hash. + /// + /// # Arguments + /// * `context` - Transaction context containing the base state root + /// * `overlay` - Overlay state containing pending account/storage modifications + /// + /// # Returns + /// * `Ok(OverlayedRoot)` - Contains the new root hash and updated branch nodes + /// * `Err(Error)` - If an error occurs during traversal or hashing + pub fn compute_state_root_with_overlay( + &self, + context: &TransactionContext, + overlay: OverlayState, + ) -> Result { + if overlay.is_empty() { + return Ok(OverlayedRoot::new_hash(context.root_node_hash)); + } + + self.overlay_trie(context, overlay).compute_root() + } + + /// Computes an account proof for the given address by applying overlay state to the persistent trie. + /// + /// This method merges the uncommitted overlay state with the persistent storage + /// to compute the account proof. + /// + /// # Arguments + /// * `context` - Transaction context containing the base state + /// * `address_path` - Path to the account address + /// * `overlay` - Overlay state containing pending modifications + /// + /// # Returns + /// * `Ok(Some(AccountProof))` - If the account exists in base state or overlay + /// * `Ok(None)` - If the account doesn't exist (was tombstoned) + /// * `Err(Error)` - If an error occurred during proof generation + pub fn compute_account_with_proof_with_overlay( + &self, + context: &TransactionContext, + address_path: AddressPath, + overlay: OverlayState, + ) -> Result, Error> { + let account_nibbles = Nibbles::from(address_path.clone()); + let mut prover = self + .overlay_trie(context, overlay) + .with_proof_targets(HashSet::from([account_nibbles])) + .compute_proof()?; + + Ok(prover.account_proof(address_path)) + } + + /// Computes both account and storage proofs for the given storage slot by applying overlay state to the persistent trie. + /// + /// This method merges the uncommitted overlay state with the persistent storage to compute the proofs. + /// + /// # Arguments + /// * `context` - Transaction context containing the base state + /// * `storage_path` - Path to the storage slot + /// * `overlay` - Overlay state containing pending modifications + /// + /// # Returns + /// * `Ok(Some(AccountProof))` - If the account exists, with storage proof included + /// * `Ok(None)` - If the account doesn't exist + /// * `Err(Error)` - If an error occurred during proof generation + pub fn compute_storage_with_proof_with_overlay( + &self, + context: &TransactionContext, + storage_path: StoragePath, + overlay: OverlayState, + ) -> Result, Error> { + let account_nibbles = Nibbles::from(storage_path.get_address().clone()); + let target_nibbles = HashSet::from([account_nibbles]); + + let slot_nibbles = *storage_path.get_slot(); + let raw_slot_path = RawPath::from(slot_nibbles); + let mut storage_targets = HashMap::default(); + storage_targets.insert(account_nibbles, HashSet::from([slot_nibbles])); + + let mut prover = self + .overlay_trie(context, overlay) + .with_proof_targets(target_nibbles) + .with_storage_proof_targets(storage_targets) + .compute_proof()?; + + let account_proof_opt = prover.account_proof(storage_path.get_address().clone()); + + if let Some(mut account_proof) = account_proof_opt { + let storage_proof_opt = prover.storage_proof(storage_path); + + if let Some(storage_proof) = storage_proof_opt { + account_proof + .storage_proofs + .insert(B256::from_slice(&raw_slot_path.pack::<32>()), storage_proof); + } + Ok(Some(account_proof)) + } else { + Ok(None) + } + } +} + +#[cfg(test)] +mod tests { + use alloy_primitives::{address, B256, U256}; + use alloy_trie::{KECCAK_EMPTY, EMPTY_ROOT_HASH, Nibbles}; + + use crate::{ + account::Account, overlay::{OverlayStateMut, OverlayValue}, + path::{AddressPath, RawPath, StoragePath}, + storage::test_utils::{create_test_account, create_test_engine, verify_account_proof}, + }; + + + #[test] + fn test_compute_account_from_overlay() { + let (storage_engine, mut context) = create_test_engine(2000); + + // 1. insert a single account to storage + let address = address!("0x0000000000000000000000000000000000000001"); + let path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + storage_engine + .set_values( + &mut context, + vec![(path.clone().into(), Some(account.clone().into()))].as_mut(), + ) + .unwrap(); + + // 2. create overlay with updated account + let mut overlay_mut = OverlayStateMut::new(); + let new_account = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + overlay_mut.insert(path.clone().into(), Some(OverlayValue::Account(new_account.clone()))); + let overlay = overlay_mut.freeze(); + + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + let proof = storage_engine.compute_account_with_proof_with_overlay(&context, path.clone(), overlay).unwrap().unwrap(); + assert_eq!(proof.account, new_account); + + // The proof logic uses Alloy Trie's HashBuilder proof retention. + // For a single leaf node, it might generate just the leaf if root is the leaf (small tree). + // Since we have 1 account in DB and we modify it in overlay, the structure is still 1 node. + + verify_account_proof(&proof, overlay_root); + } + + #[test] + fn test_compute_account_from_overlay_new_account() { + let (storage_engine, context) = create_test_engine(2000); + + // 1. Storage is empty + + // 2. Overlay adds a new account + let address = address!("0x0000000000000000000000000000000000000001"); + let path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(path.clone().into(), Some(OverlayValue::Account(account.clone()))); + let overlay = overlay_mut.freeze(); + + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + let proof = storage_engine.compute_account_with_proof_with_overlay(&context, path.clone(), overlay).unwrap().unwrap(); + assert_eq!(proof.account, account); + + verify_account_proof(&proof, overlay_root); + } + + #[test] + fn test_compute_account_from_overlay_tombstone() { + let (storage_engine, mut context) = create_test_engine(2000); + + // 1. insert account to storage + let address = address!("0x0000000000000000000000000000000000000001"); + let path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + storage_engine + .set_values( + &mut context, + vec![(path.clone().into(), Some(account.clone().into()))].as_mut(), + ) + .unwrap(); + + // 2. delete account in overlay + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(path.clone().into(), None); + let overlay = overlay_mut.freeze(); + + // 3. compute should return None + let proof = storage_engine + .compute_account_with_proof_with_overlay(&context, path.clone(), overlay) + .unwrap(); + assert!(proof.is_none()); + } + + #[test] + fn test_compute_account_from_overlay_partial_update() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address1 = address!("0x0000000000000000000000000000000000000001"); + let path1 = AddressPath::for_address(address1); + let account1 = create_test_account(1, 1); + + // Another account to force branch structure later if needed, but let's stick to 1 for simplicity first + // actually let's add a second account to DB so we have a branch + let address2 = address!("0x0000000000000000000000000000000000000002"); + let path2 = AddressPath::for_address(address2); + let account2 = create_test_account(2, 2); + + storage_engine + .set_values( + &mut context, + vec![ + (path1.clone().into(), Some(account1.clone().into())), + (path2.clone().into(), Some(account2.clone().into())) + ].as_mut(), + ) + .unwrap(); + + // Update Account 1 in overlay + let mut overlay_mut = OverlayStateMut::new(); + let new_account1 = Account::new(10, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + overlay_mut.insert(path1.clone().into(), Some(OverlayValue::Account(new_account1.clone()))); + let overlay = overlay_mut.freeze(); + + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // Get proof for Account 1 (overlayed) + let proof1 = storage_engine.compute_account_with_proof_with_overlay(&context, path1.clone(), overlay.clone()).unwrap().unwrap(); + assert_eq!(proof1.account, new_account1); + verify_account_proof(&proof1, overlay_root); + + // Get proof for Account 2 (not overlayed, but proof path affected by root change if any) + // Note: compute_account_with_proof_with_overlay should work for non-overlayed accounts too if they are in the trie + // derived from DB + overlay. + + // Wait, compute_account_with_proof_with_overlay logic builds the overlay trie. + // If account 2 is stored in DB, it will be included in the trie build process. + let proof2 = storage_engine.compute_account_with_proof_with_overlay(&context, path2.clone(), overlay).unwrap().unwrap(); + assert_eq!(proof2.account, account2); + verify_account_proof(&proof2, overlay_root); + } + + #[test] + fn test_compute_account_from_overlay_branch_creation() { + let (storage_engine, mut context) = create_test_engine(2000); + + // Construct paths explicitly to force branching + // Path A: 0x1000... (starts with nibble 1, then 0) + let mut bytes_a = [0u8; 32]; + bytes_a[0] = 0x10; + let path_a = AddressPath::new(Nibbles::unpack(B256::from(bytes_a))); + + // Path B: 0x1100... (starts with nibble 1, then 1) + let mut bytes_b = [0u8; 32]; + bytes_b[0] = 0x11; + let path_b = AddressPath::new(Nibbles::unpack(B256::from(bytes_b))); + + let account_a = create_test_account(1, 1); + let account_b = create_test_account(2, 2); + + // 1. Put Account A in DB + storage_engine.set_values( + &mut context, + vec![(path_a.clone().into(), Some(account_a.clone().into()))].as_mut(), + ).unwrap(); + + // 2. Put Account B in Overlay + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(path_b.clone().into(), Some(OverlayValue::Account(account_b.clone()))); + let overlay = overlay_mut.freeze(); + + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let root = output.root; + + // 3. Verify proofs for both + let proof_a = storage_engine.compute_account_with_proof_with_overlay(&context, path_a.clone(), overlay.clone()).unwrap().unwrap(); + verify_account_proof(&proof_a, root); + assert_eq!(proof_a.account, account_a); + + let proof_b = storage_engine.compute_account_with_proof_with_overlay(&context, path_b.clone(), overlay).unwrap().unwrap(); + verify_account_proof(&proof_b, root); + assert_eq!(proof_b.account, account_b); + } + + #[test] + fn test_compute_storage_from_overlay() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + let slot_key = B256::from(U256::from(99)); + let storage_path = StoragePath::for_address_path_and_slot(account_path.clone(), slot_key.into()); + let storage_value = U256::from(8888); + + // 1. Setup DB with account + storage_engine + .set_values( + &mut context, + vec![(account_path.clone().into(), Some(account.clone().into()))].as_mut(), + ) + .unwrap(); + + // 2. Overlay adds storage + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(storage_path.clone().into(), Some(OverlayValue::Storage(storage_value))); + + let overlay = overlay_mut.freeze(); + + // Compute root + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // Compute proof + let proof = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path.clone(), + overlay + ).unwrap().unwrap(); + + verify_account_proof(&proof, overlay_root); + + assert!(!proof.storage_proofs.is_empty()); + let sp_key = B256::from_slice(&RawPath::from(storage_path.get_slot().clone()).pack::<32>()); + let sp = proof.storage_proofs.get(&sp_key) + .expect("storage proof not found"); + assert_eq!(sp.value, storage_value); + } + + #[test] + fn test_compute_storage_db_updated_in_overlay() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + let slot_key = B256::from(U256::from(10)); + let storage_path = StoragePath::for_address_path_and_slot(account_path.clone(), slot_key.into()); + let db_value = U256::from(1111); + let overlay_value = U256::from(2222); + + // 1. Setup DB with account and storage + storage_engine + .set_values( + &mut context, + vec![ + (account_path.clone().into(), Some(account.clone().into())), + (storage_path.clone().into(), Some(db_value.into())) + ].as_mut(), + ) + .unwrap(); + + // 2. Overlay updates storage + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(storage_path.clone().into(), Some(OverlayValue::Storage(overlay_value))); + + let overlay = overlay_mut.freeze(); + + // Compute root + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // Compute proof + let proof = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path.clone(), + overlay + ).unwrap().unwrap(); + + verify_account_proof(&proof, overlay_root); + + let sp_key = B256::from_slice(&RawPath::from(storage_path.get_slot().clone()).pack::<32>()); + let sp = proof.storage_proofs.get(&sp_key).unwrap(); + assert_eq!(sp.value, overlay_value); + } + + #[test] + fn test_compute_storage_in_db_not_overlay() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + let slot_key = B256::from(U256::from(10)); + let storage_path = StoragePath::for_address_path_and_slot(account_path.clone(), slot_key.into()); + let db_value = U256::from(1111); + + // 1. Setup DB with account and storage + storage_engine + .set_values( + &mut context, + vec![ + (account_path.clone().into(), Some(account.clone().into())), + (storage_path.clone().into(), Some(db_value.into())) + ].as_mut(), + ) + .unwrap(); + + // 2. Overlay has changes to OTHER slot + let other_slot_key = B256::from(U256::from(20)); + let other_storage_path = StoragePath::for_address_path_and_slot(account_path.clone(), other_slot_key.into()); + let other_value = U256::from(9999); + + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(other_storage_path.into(), Some(OverlayValue::Storage(other_value))); + + let overlay = overlay_mut.freeze(); + + // Compute root + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // Compute proof for the DB-only slot + let proof = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path.clone(), + overlay + ).unwrap().unwrap(); + + verify_account_proof(&proof, overlay_root); + + let sp_key = B256::from_slice(&RawPath::from(storage_path.get_slot().clone()).pack::<32>()); + let sp = proof.storage_proofs.get(&sp_key).unwrap(); + assert_eq!(sp.value, db_value); + } + + #[test] + fn test_compute_storage_deleted_in_overlay() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + let slot_key = B256::from(U256::from(10)); + let storage_path = StoragePath::for_address_path_and_slot(account_path.clone(), slot_key.into()); + let db_value = U256::from(1111); + + // 1. Setup DB with account and storage + storage_engine + .set_values( + &mut context, + vec![ + (account_path.clone().into(), Some(account.clone().into())), + (storage_path.clone().into(), Some(db_value.into())) + ].as_mut(), + ) + .unwrap(); + + // 2. Overlay deletes the slot (sets to 0) + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(storage_path.clone().into(), Some(OverlayValue::Storage(U256::ZERO))); + + let overlay = overlay_mut.freeze(); + + // Compute root + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // Compute proof + let proof = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path.clone(), + overlay + ).unwrap().unwrap(); + + verify_account_proof(&proof, overlay_root); + + assert!(!proof.storage_proofs.is_empty()); + let sp_key = B256::from_slice(&RawPath::from(storage_path.get_slot().clone()).pack::<32>()); + let sp = proof.storage_proofs.get(&sp_key) + .expect("storage proof not found"); + assert_eq!(sp.value, U256::ZERO); + } + + #[test] + fn test_compute_storage_non_existent_account() { + let (storage_engine, context) = create_test_engine(2000); + let overlay_mut = OverlayStateMut::new(); + let overlay = overlay_mut.freeze(); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let slot_key = B256::from(U256::from(10)); + let storage_path = StoragePath::for_address_path_and_slot(account_path, slot_key.into()); + + // Account doesn't exist in DB or Overlay + let proof = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path, + overlay + ).unwrap(); + + assert!(proof.is_none()); + } + + #[test] + fn test_compute_storage_non_existent_slot() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + // 1. Setup DB with account ONLY (no storage) + storage_engine + .set_values( + &mut context, + vec![ + (account_path.clone().into(), Some(account.clone().into())), + ].as_mut(), + ) + .unwrap(); + + // 2. Empty Overlay + let overlay_mut = OverlayStateMut::new(); + let overlay = overlay_mut.freeze(); + + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // 3. Request proof for random slot + let slot_key = B256::from(U256::from(555)); + let storage_path = StoragePath::for_address_path_and_slot(account_path, slot_key.into()); + + let proof = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path.clone(), + overlay + ).unwrap().unwrap(); + + // Should find account + verify_account_proof(&proof, overlay_root); + + // Storage proof should exist giving value 0 (Exclusion Proof) + assert!(!proof.storage_proofs.is_empty()); + let sp_key = B256::from_slice(&RawPath::from(storage_path.get_slot().clone()).pack::<32>()); + let sp = proof.storage_proofs.get(&sp_key) + .expect("storage proof not found"); + assert_eq!(sp.value, U256::ZERO); + } + + #[test] + fn test_compute_storage_branch_creation() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + // Slot A: Starts with 0x01... + let mut slot_a_bytes = [0u8; 32]; slot_a_bytes[0] = 0x10; + let slot_a = B256::from(slot_a_bytes); + let storage_path_a = StoragePath::for_address_path_and_slot(account_path.clone(), slot_a.into()); + let val_a = U256::from(100); + + // Slot B: Starts with 0x02... (Forces branch at first nibble if they share prefix, or just specific structure) + // Let's use 0x11... to share first nibble '1' but differ at second '1'. + let mut slot_b_bytes = [0u8; 32]; slot_b_bytes[0] = 0x11; + let slot_b = B256::from(slot_b_bytes); + let storage_path_b = StoragePath::for_address_path_and_slot(account_path.clone(), slot_b.into()); + let val_b = U256::from(200); + + // 1. Setup DB with Account and Slot A + storage_engine + .set_values( + &mut context, + vec![ + (account_path.clone().into(), Some(account.clone().into())), + (storage_path_a.clone().into(), Some(val_a.into())), + ].as_mut(), + ) + .unwrap(); + + // 2. Overlay adds Slot B (Creating a branch node where previously there might have been leaf/extension) + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(storage_path_b.clone().into(), Some(OverlayValue::Storage(val_b))); + let overlay = overlay_mut.freeze(); + + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // 3. Verify Proof for Slot A (from DB, now under branch) + let proof_a = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path_a.clone(), + overlay.clone() + ).unwrap().unwrap(); + + verify_account_proof(&proof_a, overlay_root); + let sp_key_a = B256::from_slice(&RawPath::from(storage_path_a.get_slot().clone()).pack::<32>()); + assert_eq!(proof_a.storage_proofs.get(&sp_key_a).unwrap().value, val_a); + + // 4. Verify Proof for Slot B (from Overlay) + let proof_b = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path_b.clone(), + overlay + ).unwrap().unwrap(); + + verify_account_proof(&proof_b, overlay_root); + let sp_key_b = B256::from_slice(&RawPath::from(storage_path_b.get_slot().clone()).pack::<32>()); + assert_eq!(proof_b.storage_proofs.get(&sp_key_b).unwrap().value, val_b); + } + + #[test] + fn test_compute_storage_with_tombstoned_account() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + let slot_key = B256::from(U256::from(10)); + let storage_path = StoragePath::for_address_path_and_slot(account_path.clone(), slot_key.into()); + let db_value = U256::from(1111); + + // 1. Setup DB with account and storage + storage_engine + .set_values( + &mut context, + vec![ + (account_path.clone().into(), Some(account.clone().into())), + (storage_path.clone().into(), Some(db_value.into())) + ].as_mut(), + ) + .unwrap(); + + // 2. Overlay DELETEs the entire account + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(account_path.into(), None); + let overlay = overlay_mut.freeze(); + + // 3. Request proof for the storage slot + let proof = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path, + overlay + ).unwrap(); + + // Should be None because the account is gone + assert!(proof.is_none()); + } + + #[test] + fn test_compute_storage_and_account_modified() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + // Original Account: Nonce 1 + let account = create_test_account(1, 1); + + let slot_key = B256::from(U256::from(10)); + let storage_path = StoragePath::for_address_path_and_slot(account_path.clone(), slot_key.into()); + // Original Storage: 1111 + let db_value = U256::from(1111); + + storage_engine + .set_values( + &mut context, + vec![ + (account_path.clone().into(), Some(account.clone().into())), + (storage_path.clone().into(), Some(db_value.into())) + ].as_mut(), + ) + .unwrap(); + + // 2. Overlay updates BOTH Account (Nonce -> 2) and Storage (1111 -> 2222) + let mut overlay_mut = OverlayStateMut::new(); + + // Construct new account manually to ensure formatting + let new_account = Account::new(2, U256::from(1), EMPTY_ROOT_HASH, KECCAK_EMPTY); + // Warning: In a real Scenario, the storage_root in this OverlayValue::Account is usually ignored + // or re-calculated by the state root computation, but we pass it as placeholder. + overlay_mut.insert(account_path.clone().into(), Some(OverlayValue::Account(new_account.clone()))); + + let overlay_value = U256::from(2222); + overlay_mut.insert(storage_path.clone().into(), Some(OverlayValue::Storage(overlay_value))); + + let overlay = overlay_mut.freeze(); + + // Compute state root to establish the ground truth + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // 3. Request Proof + let proof = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path.clone(), + overlay + ).unwrap().unwrap(); + + // 4. Verify + verify_account_proof(&proof, overlay_root); + + // Check Account Nonce updated + assert_eq!(proof.account.nonce, 2); + + // Check Storage Value updated + let sp_key = B256::from_slice(&RawPath::from(storage_path.get_slot().clone()).pack::<32>()); + assert_eq!(proof.storage_proofs.get(&sp_key).unwrap().value, overlay_value); + } + + #[test] + fn test_compute_storage_collapsing_branch() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + // Setup two slots creating a branch in DB + // Slot A: 0x10... + let mut slot_a_bytes = [0u8; 32]; slot_a_bytes[0] = 0x10; + let slot_a = B256::from(slot_a_bytes); + let storage_path_a = StoragePath::for_address_path_and_slot(account_path.clone(), slot_a.into()); + let val_a = U256::from(100); + + // Slot B: 0x11... + let mut slot_b_bytes = [0u8; 32]; slot_b_bytes[0] = 0x11; + let slot_b = B256::from(slot_b_bytes); + let storage_path_b = StoragePath::for_address_path_and_slot(account_path.clone(), slot_b.into()); + let val_b = U256::from(200); + + storage_engine + .set_values( + &mut context, + vec![ + (account_path.clone().into(), Some(account.clone().into())), + (storage_path_a.clone().into(), Some(val_a.into())), + (storage_path_b.clone().into(), Some(val_b.into())), + ].as_mut(), + ) + .unwrap(); + + // 2. Overlay DELETES Slot B + // This should cause the trie node for Slot A to effectively merge/simplify if optimizing + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(storage_path_b.clone().into(), Some(OverlayValue::Storage(U256::ZERO))); + let overlay = overlay_mut.freeze(); + + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // 3. Request Proof for Slot A + let proof_a = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path_a.clone(), + overlay + ).unwrap().unwrap(); + + verify_account_proof(&proof_a, overlay_root); + let sp_key_a = B256::from_slice(&RawPath::from(storage_path_a.get_slot().clone()).pack::<32>()); + assert_eq!(proof_a.storage_proofs.get(&sp_key_a).unwrap().value, val_a); + } + + #[test] + fn test_compute_storage_exclusion_with_neighbor() { + let (storage_engine, mut context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + let account = create_test_account(1, 1); + + // Slot A (Exists in DB): 0x10... + let mut slot_a_bytes = [0u8; 32]; slot_a_bytes[0] = 0x10; + let slot_a = B256::from(slot_a_bytes); + let storage_path_a = StoragePath::for_address_path_and_slot(account_path.clone(), slot_a.into()); + let val_a = U256::from(100); + + // Slot B (Missing): 0x11... (Shares prefix '1' with A, differs at second nibble) + let mut slot_b_bytes = [0u8; 32]; slot_b_bytes[0] = 0x11; + let slot_b = B256::from(slot_b_bytes); + let storage_path_b = StoragePath::for_address_path_and_slot(account_path.clone(), slot_b.into()); + + // 1. Setup DB with Slot A + storage_engine + .set_values( + &mut context, + vec![ + (account_path.clone().into(), Some(account.clone().into())), + (storage_path_a.clone().into(), Some(val_a.into())), + ].as_mut(), + ) + .unwrap(); + + // 2. Empty Overlay + let overlay_mut = OverlayStateMut::new(); + let overlay = overlay_mut.freeze(); + + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // 3. Request Proof for MISSING Slot B + // This should return an inclusion proof for the branch node + exclusion for the leaf + let proof_b = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path_b.clone(), + overlay + ).unwrap().unwrap(); + + // Verify Account Proof + verify_account_proof(&proof_b, overlay_root); + + // Verify Storage Exclusion Proof + assert!(!proof_b.storage_proofs.is_empty()); + let sp_key_b = B256::from_slice(&RawPath::from(storage_path_b.get_slot().clone()).pack::<32>()); + let sp = proof_b.storage_proofs.get(&sp_key_b).unwrap(); + + assert_eq!(sp.value, U256::ZERO); + // The proof should likely contain the branch node common to A and B + assert!(sp.proof.len() >= 1); + } + + #[test] + fn test_compute_storage_from_new_account_in_overlay() { + let (storage_engine, context) = create_test_engine(2000); + + let address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(address); + // New Account + let account = create_test_account(1, 1); + + let slot_key = B256::from(U256::from(10)); + let storage_path = StoragePath::for_address_path_and_slot(account_path.clone(), slot_key.into()); + let val = U256::from(12345); + + // 1. Storage is EMPTY + + // 2. Overlay creates Account AND Storage + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(account_path.clone().into(), Some(OverlayValue::Account(account.clone()))); + overlay_mut.insert(storage_path.clone().into(), Some(OverlayValue::Storage(val))); + + let overlay = overlay_mut.freeze(); + + let output = storage_engine.compute_state_root_with_overlay(&context, overlay.clone()).unwrap(); + let overlay_root = output.root; + + // 3. Request Proof for Storage + let proof = storage_engine.compute_storage_with_proof_with_overlay( + &context, + storage_path.clone(), + overlay + ).unwrap().unwrap(); + + verify_account_proof(&proof, overlay_root); + + let sp_key = B256::from_slice(&RawPath::from(storage_path.get_slot().clone()).pack::<32>()); + assert_eq!(proof.storage_proofs.get(&sp_key).unwrap().value, val); + + // Also verify account matches (it should have the computed storage root, not empty) + assert_ne!(proof.account.storage_root, EMPTY_ROOT_HASH); + assert_eq!(proof.account.nonce, account.nonce); + assert_eq!(proof.account.balance, account.balance); + } +} \ No newline at end of file diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_traversal.rs similarity index 76% rename from src/storage/overlay_root.rs rename to src/storage/overlay_traversal.rs index 1b36c7c2..da297463 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_traversal.rs @@ -6,271 +6,289 @@ use crate::{ node::{encode_account_leaf, Node, NodeKind}, overlay::{OverlayState, OverlayValue}, page::SlottedPage, - path::RawPath, + path::{RawPath, AddressPath, StoragePath}, pointer::Pointer, - storage::engine::{Error, StorageEngine}, + storage::{ + engine::{Error, StorageEngine}, + proofs::{AccountProof, StorageProof}, + overlay::OverlayedRoot, + }, }; +use std::collections::{HashSet, BTreeMap}; use alloy_primitives::{ map::{B256Map, HashMap}, B256, U256, }; -use alloy_rlp::encode_fixed_size; -use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, EMPTY_ROOT_HASH}; +use alloy_rlp::{encode_fixed_size, Decodable}; +use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, EMPTY_ROOT_HASH, nodes::TrieNode, TrieAccount}; +use alloy_trie::proof::{ProofNodes, ProofRetainer}; use arrayvec::ArrayVec; +/// Builder for constructing a new trie by applying an overlay state on top of +/// a persistent storage engine. +/// +/// This builder traverses the existing trie and the overlay simultaneously, +/// constructing a new Merkle Trie (via `OverlayProver`) that represents the merged state. +/// It supports: +/// - Efficient updates by skipping unchanged subtrees (using structural sharing). +/// - Generating Merkle proofs for specific targets during the build process. +// # TODO: Semantic Refactoring Needed +// +// This struct is used for both state tries and storage tries, but the fields have +// different meanings depending on context: +// - **For state tries**: `proof_targets` = account paths, `storage_proof_targets` = account → storage slots +// - **For storage tries**: `proof_targets` = storage slot paths, `storage_proof_targets` = unused (empty) #[derive(Debug)] -enum TriePosition<'a> { - Node(RawPath, Rc>, Node), - Pointer(RawPath, Rc>, Pointer, bool), - None, -} - -struct TraversalStack<'a> { - stack: Vec<(TriePosition<'a>, OverlayState)>, -} - -impl<'a> TraversalStack<'a> { - fn new() -> Self { - Self { stack: vec![] } - } - - fn push_node( - &mut self, - path: RawPath, - node: Node, - page: Rc>, - overlay: OverlayState, - ) { - self.push(TriePosition::Node(path, page, node), overlay); - } - - fn push_pointer( - &mut self, - path: RawPath, - pointer: Pointer, - page: Rc>, - can_add_by_hash: bool, - overlay: OverlayState, - ) { - self.push(TriePosition::Pointer(path, page, pointer, can_add_by_hash), overlay); - } - - fn push_overlay(&mut self, overlay: OverlayState) { - self.push(TriePosition::None, overlay); - } - - fn push(&mut self, position: TriePosition<'a>, overlay: OverlayState) { - self.stack.push((position, overlay)); - } - - fn pop(&mut self) -> Option<(TriePosition<'a>, OverlayState)> { - self.stack.pop() - } +pub struct OverlayTrie<'a> { + engine: &'a StorageEngine, + context: &'a TransactionContext, + overlay: OverlayState, + /// Proof targets - account paths for state tries, storage slot paths for storage tries + proof_targets: HashSet, + /// Storage proof targets - only used for state tries (maps account paths to storage slot paths) + storage_proof_targets: HashMap>, } -#[derive(Debug)] -pub struct OverlayedRoot { - pub root: B256, - pub updated_branch_nodes: HashMap, - pub storage_branch_updates: B256Map>, -} - -impl OverlayedRoot { +impl<'a> OverlayTrie<'a> { pub fn new( - root: B256, - updated_branch_nodes: HashMap, - storage_branch_updates: B256Map>, + engine: &'a StorageEngine, + context: &'a TransactionContext, + overlay: OverlayState, ) -> Self { - Self { root, updated_branch_nodes, storage_branch_updates } - } - - pub fn new_hash(root: B256) -> Self { - Self { - root, - updated_branch_nodes: HashMap::default(), - storage_branch_updates: B256Map::default(), - } - } -} -struct RootBuilder { - hash_builder: HashBuilder, - storage_branch_updates: B256Map>, -} - -impl RootBuilder { - fn new() -> Self { Self { - hash_builder: HashBuilder::default().with_updates(true), - storage_branch_updates: B256Map::default(), + engine, + context, + overlay, + proof_targets: HashSet::new(), + storage_proof_targets: HashMap::default(), } } - fn add_leaf(&mut self, key: Nibbles, value: &[u8]) { - self.hash_builder.add_leaf(key, value); + pub fn with_proof_targets(mut self, targets: HashSet) -> Self { + self.proof_targets = targets; + self } - fn add_branch(&mut self, key: Nibbles, value: B256, stored_in_database: bool) { - self.hash_builder.add_branch(key, value, stored_in_database); + pub fn with_storage_proof_targets( + mut self, + targets: HashMap>, + ) -> Self { + self.storage_proof_targets = targets; + self } - fn add_storage_branch_updates( - &mut self, - account: B256, - updates: HashMap, - ) { - self.storage_branch_updates.insert(account, updates); + pub fn compute_root(self) -> Result { + Ok(self.build()?.root()) } - fn finalize(self) -> OverlayedRoot { - let (mut hash_builder, updated_branch_nodes) = self.hash_builder.split(); - OverlayedRoot::new(hash_builder.root(), updated_branch_nodes, self.storage_branch_updates) + pub fn compute_proof(self) -> Result { + self.build() } -} -impl StorageEngine { - pub fn compute_state_root_with_overlay( - &self, - context: &TransactionContext, - overlay: OverlayState, - ) -> Result { - if overlay.is_empty() { - return Ok(OverlayedRoot::new_hash(context.root_node_hash)); - } - - let mut root_builder = RootBuilder::new(); + fn build(self) -> Result { + let mut root_builder = if !self.proof_targets.is_empty() { + OverlayProver::default() + .with_proof_retainer(self.proof_targets.clone()) + .with_updates(true) + } else { + OverlayProver::default() + }; - let root_page = if let Some(root_page_id) = context.root_node_page_id { - let page = self.get_page(context, root_page_id)?; + let root_page = if let Some(root_page_id) = self.context.root_node_page_id { + let page = self.engine.get_page(self.context, root_page_id)?; SlottedPage::try_from(page).unwrap() } else { - self.add_overlay_to_root_builder(&mut root_builder, &overlay); - return Ok(root_builder.finalize()); + self.add_overlay_to_root_builder(&mut root_builder, &self.overlay)?; + // Finalize the hash builder state by computing the root + let _ = root_builder.hash_builder.root(); + return Ok(root_builder); }; let root_node: Node = root_page.get_value(0)?; let mut stack = TraversalStack::new(); - stack.push_node(root_node.prefix().into(), root_node, Rc::new(root_page), overlay); + stack.push_node( + root_node.prefix().into(), + root_node, + Rc::new(root_page), + self.overlay.clone(), + ); + + self.compute_traversal(&mut stack, &mut root_builder)?; - self.compute_root_with_overlay(context, &mut stack, &mut root_builder)?; + // Finalize the hash builder state by computing the root. + // The root value itself is not needed here as it will be recomputed + // when the caller accesses the OverlayProver. + let _ = root_builder.hash_builder.root(); - Ok(root_builder.finalize()) + Ok(root_builder) } - fn compute_root_with_overlay<'a>( - &'a self, - context: &TransactionContext, + /// Main traversal loop: processes positions from the stack until empty. + /// + /// This is the coordinator function that delegates to specialized handlers + /// based on the type of trie position being processed. + fn compute_traversal( + &self, stack: &mut TraversalStack<'a>, - root_builder: &mut RootBuilder, + root_builder: &mut OverlayProver, ) -> Result<(), Error> { - // Depth first traversal of the trie, starting at the root node. - // This applies any overlay state to the trie, taking precedence over the trie's own values. - // Whenever a branch or leaf is known to be the final unchanged value, we can add it to the - // hash builder. while let Some((position, overlay)) = stack.pop() { match position { TriePosition::None => { - // No trie position, process whatever is in the overlay - self.add_overlay_to_root_builder(root_builder, &overlay); + self.add_overlay_to_root_builder(root_builder, &overlay)?; } TriePosition::Pointer(path, page, pointer, can_add_by_hash) => { - if overlay.is_empty() && can_add_by_hash { - if let Some(hash) = pointer.rlp().as_hash() { - // No overlay, just add the pointer by hash - root_builder.add_branch(path.try_into().unwrap(), hash, true); - continue; - } - } - // We have an overlay, need to process the child - self.process_overlayed_child( - context, - overlay, - root_builder, - path, - &pointer, - page, - stack, - )?; + self.process_pointer(path, page, pointer, can_add_by_hash, overlay, stack, root_builder)?; } TriePosition::Node(path, page, node) => { - let (pre_overlay, matching_overlay, post_overlay) = - overlay.sub_slice_by_prefix(&path); - if pre_overlay.contains_prefix_of(&path) { - // The pre_overlay invalidates the current node, so we can simply add the - // full overlay. We need to process it all together, - // as the post_overlay may contain descendants of the pre_overlay. - self.add_overlay_to_root_builder(root_builder, &overlay); - continue; - } - - self.add_overlay_to_root_builder(root_builder, &pre_overlay); - // Defer the post_overlay to be processed after the node is traversed - stack.push_overlay(post_overlay); - - match node.into_kind() { - NodeKind::Branch { children } => { - if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = - matching_overlay.first() - { - if overlay_path == path { - // the overlay invalidates the current node, so just add this - // and skip the rest of the db traversal - self.add_overlay_to_root_builder( - root_builder, - &matching_overlay, - ); - continue; - } - } - self.process_branch_node_with_overlay( - matching_overlay, - &path, - children, - page, - stack, - )?; - } - NodeKind::AccountLeaf { - nonce_rlp, - balance_rlp, - code_hash, - storage_root, - } => { - self.process_account_leaf_with_overlay( - context, - &matching_overlay, - root_builder, - &path, - page, - nonce_rlp, - balance_rlp, - code_hash, - storage_root, - )?; - } - NodeKind::StorageLeaf { value_rlp } => { - if let Some((overlay_path, _)) = matching_overlay.first() { - if overlay_path == path { - // the overlay invalidates the current node, so just add this - // and skip the rest of the db traversal - self.add_overlay_to_root_builder( - root_builder, - &matching_overlay, - ); - continue; - } - } - // Leaf node, add it to the hash builder - root_builder.add_leaf(path.try_into().unwrap(), &value_rlp); - } - } + self.process_node(path, page, node, overlay, stack, root_builder)?; } } } Ok(()) } - fn process_branch_node_with_overlay<'a>( - &'a self, + /// Processes a pointer to a child node, potentially adding by hash if optimization applies. + fn process_pointer( + &self, + path: RawPath, + page: Rc>, + pointer: Pointer, + can_add_by_hash: bool, + overlay: OverlayState, + stack: &mut TraversalStack<'a>, + root_builder: &mut OverlayProver, + ) -> Result<(), Error> { + // if overlay is empty and we can add by hash, skip traversing this subtree + if overlay.is_empty() && can_add_by_hash { + if let Some(hash) = pointer.rlp().as_hash() { + root_builder.add_branch(path.try_into().unwrap(), hash, true); + return Ok(()); + } + } + + self.process_overlayed_child(overlay, root_builder, path, &pointer, page, stack) + } + + /// Processes a trie node with potential overlay modifications. + fn process_node( + &self, + path: RawPath, + page: Rc>, + node: Node, + overlay: OverlayState, + stack: &mut TraversalStack<'a>, + root_builder: &mut OverlayProver, + ) -> Result<(), Error> { + // Split overlay into: entries before this path, matching this path, and after + let (pre_overlay, matching_overlay, post_overlay) = overlay.sub_slice_by_prefix(&path); + + // If pre_overlay contains a prefix of this path, it means there's an overlay entry + // that shadows this entire subtree - skip the node and just add overlay + if pre_overlay.contains_prefix_of(&path) { + return self.add_overlay_to_root_builder(root_builder, &overlay); + } + + // Add entries before this node, schedule entries after for later + self.add_overlay_to_root_builder(root_builder, &pre_overlay)?; + stack.push_overlay(post_overlay); + + // Process this node based on its type + match node.into_kind() { + NodeKind::Branch { children } => { + self.handle_branch_node(path, page, children, matching_overlay, stack, root_builder)?; + } + NodeKind::AccountLeaf { + nonce_rlp, + balance_rlp, + code_hash, + storage_root, + } => { + self.handle_account_leaf( + path, + page, + nonce_rlp, + balance_rlp, + code_hash, + storage_root, + matching_overlay, + root_builder, + )?; + } + NodeKind::StorageLeaf { value_rlp } => { + self.handle_storage_leaf(path, value_rlp, matching_overlay, root_builder)?; + } + } + + Ok(()) + } + + /// Handles a branch node, processing its children with overlay modifications. + fn handle_branch_node( + &self, + path: RawPath, + page: Rc>, + children: [Option; 16], + matching_overlay: OverlayState, + stack: &mut TraversalStack<'a>, + root_builder: &mut OverlayProver, + ) -> Result<(), Error> { + // Check if overlay has a hash node that replaces this entire branch + if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = matching_overlay.first() { + if overlay_path == path { + return self.add_overlay_to_root_builder(root_builder, &matching_overlay); + } + } + + self.process_branch_node_with_overlay(matching_overlay, &path, children, page, stack) + } + + /// Handles an account leaf node with potential overlay modifications. + fn handle_account_leaf( + &self, + path: RawPath, + page: Rc>, + nonce_rlp: ArrayVec, + balance_rlp: ArrayVec, + code_hash: B256, + storage_root: Option, + matching_overlay: OverlayState, + root_builder: &mut OverlayProver, + ) -> Result<(), Error> { + self.process_account_leaf_with_overlay( + &matching_overlay, + root_builder, + &path, + page, + nonce_rlp, + balance_rlp, + code_hash, + storage_root, + ) + } + + /// Handles a storage leaf node with potential overlay modifications. + fn handle_storage_leaf( + &self, + path: RawPath, + value_rlp: ArrayVec, + matching_overlay: OverlayState, + root_builder: &mut OverlayProver, + ) -> Result<(), Error> { + // Check if overlay has a value that replaces this storage slot + if let Some((overlay_path, _)) = matching_overlay.first() { + if overlay_path == path { + return self.add_overlay_to_root_builder(root_builder, &matching_overlay); + } + } + + // No overlay replacement - add the existing leaf + root_builder.add_leaf(path.try_into().unwrap(), &value_rlp); + Ok(()) + } + + fn process_branch_node_with_overlay( + &self, mut overlay: OverlayState, path: &RawPath, mut children: [Option; 16], @@ -293,24 +311,28 @@ impl StorageEngine { if child_pointer.is_some() && child_overlay.is_empty() { minimum_possible_child_count += 1; } else if let Some((_, Some(_))) = child_overlay.first() { - // we have a non-tombstone overlay, so there must be at least one descendant - // in this child index minimum_possible_child_count += 1; } child_data.push((child_path, child_pointer, child_overlay)); overlay = overlay_after_child; } - let can_add_by_hash = minimum_possible_child_count > 1; + let can_add_by_hash_globally = minimum_possible_child_count > 1; for (child_path, child_pointer, child_overlay) in child_data.into_iter().rev() { match child_pointer { Some(pointer) => { + let is_on_target_path = self.proof_targets.iter().any(|target| { + target.starts_with(&Nibbles::from_nibbles(child_path.nibbles())) + }); + let can_add_this_child_by_hash = + can_add_by_hash_globally && !is_on_target_path; + stack.push_pointer( child_path, pointer, current_page.clone(), - can_add_by_hash, + can_add_this_child_by_hash, child_overlay, ); } @@ -318,7 +340,6 @@ impl StorageEngine { if child_overlay.is_empty() { // nothing here to add } else { - // we have a nonconflicting overlay, add all of it to the hash builder stack.push_overlay(child_overlay); } } @@ -327,11 +348,104 @@ impl StorageEngine { Ok(()) } - fn process_account_leaf_with_overlay<'a>( - &'a self, - context: &TransactionContext, + fn compute_storage_tree( + &self, + storage_root: Option, + current_page: Option>>, + storage_overlay: OverlayState, + proof_targets: HashSet, + ) -> Result<(B256, HashMap, ProofNodes), Error> { + let mut storage_root_builder = if !proof_targets.is_empty() { + OverlayProver::default() + .with_proof_retainer(proof_targets.clone()) + .with_updates(true) + } else { + OverlayProver::default() + }; + let storage_builder = self + .engine + .overlay_trie(self.context, storage_overlay) + .with_proof_targets(proof_targets); + + match storage_root { + Some(pointer) => { + let mut storage_stack = TraversalStack::new(); + + let (root_storage_node, page) = if let Some(child_cell) = + pointer.location().cell_index() + { + let current_page = current_page.expect("Current page must be provided for cell index pointer"); + let node: Node = current_page.get_value(child_cell)?; + (node, current_page) + } else { + let storage_page = self + .engine + .get_page(self.context, pointer.location().page_id().unwrap())?; + let slotted_page = SlottedPage::try_from(storage_page)?; + let node: Node = slotted_page.get_value(0)?; + (node, Rc::new(slotted_page)) + }; + + storage_stack.push_node( + root_storage_node.prefix().into(), + root_storage_node, + page, + storage_builder.overlay.clone(), + ); + storage_builder.compute_traversal( + &mut storage_stack, + &mut storage_root_builder, + )?; + } + None => { + storage_builder.add_overlay_to_root_builder( + &mut storage_root_builder, + &storage_builder.overlay, + )?; + } + }; + + let (mut storage_hash_builder, updated_storage_branch_nodes) = + storage_root_builder.hash_builder.split(); + let root = storage_hash_builder.root(); + let proof_nodes = storage_hash_builder.take_proof_nodes(); + Ok(( + root, + updated_storage_branch_nodes, + proof_nodes, + )) + } + + fn compute_and_update_storage_root( + &self, + root_builder: &mut OverlayProver, + account_path_bytes: B256, + storage_root: Option, + current_page: Option>>, + storage_overlay: OverlayState, + proof_targets: HashSet, + ) -> Result { + let (new_root, updated_storage_branch_nodes, proof_nodes) = self.compute_storage_tree( + storage_root, + current_page, + storage_overlay, + proof_targets, + )?; + + root_builder + .add_storage_branch_updates(account_path_bytes, updated_storage_branch_nodes); + + if !proof_nodes.is_empty() { + root_builder.add_storage_proofs(account_path_bytes, proof_nodes); + } + + Ok(new_root) + } + + fn process_account_leaf_with_overlay( + &self, overlay: &OverlayState, - root_builder: &mut RootBuilder, + root_builder: &mut OverlayProver, path: &RawPath, current_page: Rc>, mut nonce_rlp: ArrayVec, @@ -342,22 +456,26 @@ impl StorageEngine { let overlayed_account = overlay.lookup(path); match overlayed_account { Some(None) => { - // The account is removed in the overlay return Ok(()); } Some(Some(OverlayValue::Account(overlayed_account))) => { - // The account is updated in the overlay nonce_rlp = alloy_rlp::encode_fixed_size(&overlayed_account.nonce); balance_rlp = alloy_rlp::encode_fixed_size(&overlayed_account.balance); code_hash = overlayed_account.code_hash; } _ => { - // The account is not updated in the overlay } }; let has_storage_overlays = overlay.iter().any(|(path, _)| path.len() > 64); - if !has_storage_overlays { + let account_path_nibbles = Nibbles::from_nibbles(path.nibbles()); + let storage_proof_targets = self + .storage_proof_targets + .get(&account_path_nibbles) + .cloned() + .unwrap_or_default(); + + if !has_storage_overlays && storage_proof_targets.is_empty() { let storage_root_hash = storage_root .as_ref() .map_or(EMPTY_ROOT_HASH, |p| p.rlp().as_hash().unwrap_or(EMPTY_ROOT_HASH)); @@ -373,60 +491,15 @@ impl StorageEngine { return Ok(()); } - let mut storage_root_builder = RootBuilder::new(); - - // We have storage overlays, need to compute a new storage root let storage_overlay = overlay.with_prefix_offset(64); - - match storage_root { - Some(pointer) => { - let mut storage_stack = TraversalStack::new(); - - // load the root storage node - if let Some(child_cell) = pointer.location().cell_index() { - let root_storage_node: Node = current_page.get_value(child_cell)?; - storage_stack.push_node( - root_storage_node.prefix().into(), - root_storage_node, - current_page, - storage_overlay, - ); - self.compute_root_with_overlay( - context, - &mut storage_stack, - &mut storage_root_builder, - )? - } else { - let storage_page = - self.get_page(context, pointer.location().page_id().unwrap())?; - let slotted_page = SlottedPage::try_from(storage_page)?; - let root_storage_node: Node = slotted_page.get_value(0)?; - storage_stack.push_node( - root_storage_node.prefix().into(), - root_storage_node, - Rc::new(slotted_page), - storage_overlay, - ); - self.compute_root_with_overlay( - context, - &mut storage_stack, - &mut storage_root_builder, - )?; - } - } - None => { - // No existing storage root, just add the overlay - self.add_overlay_to_root_builder(&mut storage_root_builder, &storage_overlay); - } - }; - let (mut storage_hash_builder, updated_storage_branch_nodes) = - storage_root_builder.hash_builder.split(); - let new_root = storage_hash_builder.root(); - - root_builder.add_storage_branch_updates( + let new_root = self.compute_and_update_storage_root( + root_builder, B256::from_slice(&path.pack::<32>()), - updated_storage_branch_nodes, - ); + storage_root, + Some(current_page), + storage_overlay, + storage_proof_targets, + )?; self.add_account_leaf_to_root_builder( root_builder, @@ -442,42 +515,34 @@ impl StorageEngine { fn add_account_leaf_to_root_builder( &self, - root_builder: &mut RootBuilder, + root_builder: &mut OverlayProver, path: RawPath, nonce_rlp: &ArrayVec, balance_rlp: &ArrayVec, code_hash: &B256, storage_root: &B256, ) { - let mut buf = [0u8; 110]; // max RLP length for an account: 2 bytes for list length, 9 for nonce, 33 for - // balance, 33 for storage root, 33 for code hash + let mut buf = [0u8; 110]; let mut value_rlp = buf.as_mut(); let account_rlp_length = encode_account_leaf(nonce_rlp, balance_rlp, code_hash, storage_root, &mut value_rlp); root_builder.add_leaf(path.try_into().unwrap(), &buf[..account_rlp_length]); } - fn process_overlayed_child<'a>( - &'a self, - context: &TransactionContext, + fn process_overlayed_child( + &self, overlay: OverlayState, - root_builder: &mut RootBuilder, + root_builder: &mut OverlayProver, mut child_path: RawPath, child: &Pointer, current_page: Rc>, stack: &mut TraversalStack<'a>, ) -> Result<(), Error> { - // First consider the overlay. All values in it must already contain the child_path prefix. - // If the overlay matches the child path, we can add it to the hash builder and skip - // actually reading the child node. - // Account values cannot be directly overlayed, as they may need to be merged with the - // existing storage trie. if let Some((overlay_path, overlay_value)) = overlay.first() { if child_path == overlay_path && !matches!(overlay_value, Some(OverlayValue::Account(_))) { - // the child path is directly overlayed, so only use the overlay state - self.add_overlay_to_root_builder(root_builder, &overlay); + self.add_overlay_to_root_builder(root_builder, &overlay)?; return Ok(()); } } @@ -488,7 +553,7 @@ impl StorageEngine { stack.push_node(child_path, child_node, current_page, overlay); } else { let child_page_id = child.location().page_id().unwrap(); - let child_page = self.get_page(context, child_page_id)?; + let child_page = self.engine.get_page(self.context, child_page_id)?; let child_slotted_page = SlottedPage::try_from(child_page).unwrap(); let child_node: Node = child_slotted_page.get_value(0)?; child_path.extend(&child_node.prefix().into()); @@ -499,35 +564,44 @@ impl StorageEngine { fn process_overlayed_account( &self, - root_builder: &mut RootBuilder, + root_builder: &mut OverlayProver, path: Nibbles, account: &Account, storage_overlay: OverlayState, ) -> Result<(), Error> { - if storage_overlay.is_empty() { + let storage_proof_targets = + self.storage_proof_targets.get(&path).cloned().unwrap_or_default(); + + if storage_overlay.is_empty() && storage_proof_targets.is_empty() { let encoded = self.encode_account(account); root_builder.add_leaf(path, &encoded); return Ok(()); } - let mut storage_root_builder = RootBuilder::new(); - self.add_overlay_to_root_builder(&mut storage_root_builder, &storage_overlay); - - let (mut storage_hash_builder, updated_storage_branch_nodes) = - storage_root_builder.hash_builder.split(); - let storage_root = storage_hash_builder.root(); - - root_builder.add_storage_branch_updates( + let storage_root = self.compute_and_update_storage_root( + root_builder, B256::from_slice(&path.pack()), - updated_storage_branch_nodes, - ); + None, + None, + storage_overlay, + storage_proof_targets, + )?; let encoded = self.encode_account_with_root(account, storage_root); root_builder.add_leaf(path, &encoded); Ok(()) } - fn add_overlay_to_root_builder(&self, root_builder: &mut RootBuilder, overlay: &OverlayState) { + /// Adds overlay entries to the root builder. + /// + /// This method processes overlay entries and adds them to the Merkle trie builder. + /// It handles accounts, storage slots, and hash nodes, skipping descendants of + /// already-processed paths for efficiency. + pub(crate) fn add_overlay_to_root_builder( + &self, + root_builder: &mut OverlayProver, + overlay: &OverlayState, + ) -> Result<(), Error> { let mut last_processed_path = None; for (path, value) in overlay.iter() { if let Some(last_processed_path) = last_processed_path { @@ -541,29 +615,31 @@ impl StorageEngine { Some(OverlayValue::Account(account)) => { let storage_overlay = overlay.sub_slice_for_prefix(&path).with_prefix_offset(64); + let nibbles = path.try_into().unwrap(); self.process_overlayed_account( root_builder, - path.try_into().unwrap(), + nibbles, account, storage_overlay, - ) - .unwrap(); + )?; last_processed_path = Some(path); } Some(OverlayValue::Storage(storage_value)) => { let encoded = self.encode_storage(storage_value); - root_builder.add_leaf(path.try_into().unwrap(), &encoded); + let nibbles = path.try_into().unwrap(); + root_builder.add_leaf(nibbles, &encoded); } Some(OverlayValue::Hash(hash)) => { - root_builder.add_branch(path.try_into().unwrap(), *hash, false); + let nibbles = path.try_into().unwrap(); + root_builder.add_branch(nibbles, *hash, false); last_processed_path = Some(path); } None => { - // Tombstone - skip last_processed_path = Some(path); } } } + Ok(()) } #[inline] @@ -594,6 +670,193 @@ impl StorageEngine { } } +#[derive(Debug)] +enum TriePosition<'a> { + Node(RawPath, Rc>, Node), + Pointer(RawPath, Rc>, Pointer, bool), + None, +} + +struct TraversalStack<'a> { + stack: Vec<(TriePosition<'a>, OverlayState)>, +} + +impl<'a> TraversalStack<'a> { + fn new() -> Self { + Self { stack: vec![] } + } + + fn push_node( + &mut self, + path: RawPath, + node: Node, + page: Rc>, + overlay: OverlayState, + ) { + self.push(TriePosition::Node(path, page, node), overlay); + } + + fn push_pointer( + &mut self, + path: RawPath, + pointer: Pointer, + page: Rc>, + can_add_by_hash: bool, + overlay: OverlayState, + ) { + self.push(TriePosition::Pointer(path, page, pointer, can_add_by_hash), overlay); + } + + fn push_overlay(&mut self, overlay: OverlayState) { + self.push(TriePosition::None, overlay); + } + + fn push(&mut self, position: TriePosition<'a>, overlay: OverlayState) { + self.stack.push((position, overlay)); + } + + fn pop(&mut self) -> Option<(TriePosition<'a>, OverlayState)> { + self.stack.pop() + } +} + +#[derive(Debug)] +#[derive(Default)] +pub struct OverlayProver { + pub(crate) hash_builder: HashBuilder, + pub(crate) storage_branch_updates: B256Map>, + pub(crate) retained_storage_proofs: B256Map, +} + + +impl OverlayProver { + pub(crate) fn with_proof_retainer(self, targets: HashSet) -> OverlayProver { + OverlayProver { + hash_builder: self.hash_builder + .with_proof_retainer(ProofRetainer::new( + targets.into_iter().collect(), + )) + .with_updates(true), + storage_branch_updates: self.storage_branch_updates, + retained_storage_proofs: self.retained_storage_proofs, + } + } + + pub(crate) fn with_updates(mut self, with_updates: bool) -> Self { + self.hash_builder = self.hash_builder.with_updates(with_updates); + self + } + + pub(crate) fn add_leaf(&mut self, key: Nibbles, value: &[u8]) { + self.hash_builder.add_leaf(key, value); + } + + pub(crate) fn add_branch(&mut self, key: Nibbles, value: B256, stored_in_database: bool) { + self.hash_builder.add_branch(key, value, stored_in_database); + } + + pub(crate) fn add_storage_branch_updates( + &mut self, + account: B256, + updates: HashMap, + ) { + self.storage_branch_updates.insert(account, updates); + } + + pub(crate) fn add_storage_proofs(&mut self, account: B256, proofs: ProofNodes) { + self.retained_storage_proofs.insert(account, proofs); + } + + pub fn account_proof(&mut self, address_path: AddressPath) -> Option { + let account_nibbles = Nibbles::from(address_path); + let proof_nodes = self.hash_builder.take_proof_nodes(); + let matching_nodes = proof_nodes.matching_nodes_sorted(&account_nibbles); + + let mut proof_map = BTreeMap::new(); + for (key, val) in matching_nodes.iter() { + proof_map.insert(RawPath::from(*key), val.clone()); + } + + let found_account = 'info: { + if let Some((_, last_node_bytes)) = matching_nodes.last() { + let mut data = last_node_bytes.as_ref(); + if let Ok(TrieNode::Leaf(leaf)) = TrieNode::decode(&mut data) { + let mut full_path = matching_nodes.last().unwrap().0; + full_path.extend(&leaf.key); + + if full_path == account_nibbles { + let mut value_data = leaf.value.as_ref(); + if let Ok(account) = TrieAccount::decode(&mut value_data) { + break 'info Some(Account::new( + account.nonce, + account.balance, + account.storage_root, + account.code_hash, + )); + } + } + } + } + None + }; + + found_account.map(|account| { + AccountProof { + hashed_address: account_nibbles, + account, + proof: proof_map, + storage_proofs: B256Map::default(), + } + }) + } + + pub fn storage_proof(&self, storage_path: StoragePath) -> Option { + let account_nibbles = Nibbles::from(storage_path.get_address().clone()); + let slot_nibbles = *storage_path.get_slot(); + + let account_key = B256::from_slice(&RawPath::from(account_nibbles).pack::<32>()); + if let Some(storage_proof_nodes) = self.retained_storage_proofs.get(&account_key) { + let matching_nodes = storage_proof_nodes.matching_nodes_sorted(&slot_nibbles); + + let mut proof_map = BTreeMap::new(); + for (key, val) in matching_nodes.iter() { + proof_map.insert(RawPath::from(*key), val.clone()); + } + + let found_value = 'info: { + if let Some((_, last_node_bytes)) = matching_nodes.last() { + let mut data = last_node_bytes.as_ref(); + if let Ok(TrieNode::Leaf(leaf)) = TrieNode::decode(&mut data) { + let mut full_path = matching_nodes.last().unwrap().0; + full_path.extend(&leaf.key); + if full_path == slot_nibbles { + let mut value_data = leaf.value.as_ref(); + if let Ok(value) = U256::decode(&mut value_data) { + break 'info Some(value); + } + } + } + } + + break 'info Some(U256::ZERO); + }; + + found_value.map(|value| StorageProof { + hashed_slot: slot_nibbles, + value, + proof: proof_map, + }) + } else { + None + } + } + + pub(crate) fn root(self) -> OverlayedRoot { + let (mut hash_builder, updated_branch_nodes) = self.hash_builder.split(); + OverlayedRoot::new(hash_builder.root(), updated_branch_nodes, self.storage_branch_updates) + } +} + #[cfg(test)] mod tests { use alloy_primitives::{address, Address, U256}; @@ -1754,4 +2017,4 @@ mod tests { (output.root, output.updated_branch_nodes, output.storage_branch_updates); assert_ne!(overlay_root, initial_root, "Overlay should not match initial root"); } -} +} \ No newline at end of file diff --git a/src/storage/proofs.rs b/src/storage/proofs.rs index 39627738..11c53dd3 100644 --- a/src/storage/proofs.rs +++ b/src/storage/proofs.rs @@ -324,36 +324,10 @@ impl StorageEngine { #[cfg(test)] mod tests { use alloy_primitives::{address, b256, hex, U256}; - use alloy_rlp::encode; - use alloy_trie::{proof::verify_proof, TrieAccount, KECCAK_EMPTY}; + use alloy_trie::KECCAK_EMPTY; use super::*; - use crate::storage::test_utils::{create_test_account, create_test_engine}; - - fn verify_account_proof(proof: &AccountProof, root: B256) { - let expected = Some(encode(TrieAccount { - nonce: proof.account.nonce, - balance: proof.account.balance, - storage_root: proof.account.storage_root, - code_hash: proof.account.code_hash, - })); - verify_proof(root, proof.hashed_address, expected, proof.proof.values()) - .expect("failed to verify account proof"); - - for storage_proof in proof.storage_proofs.values() { - verify_storage_proof(storage_proof, proof.account.storage_root); - } - } - - fn verify_storage_proof(proof: &StorageProof, root: B256) { - verify_proof( - root, - proof.hashed_slot, - Some(alloy_rlp::encode(proof.value)), - proof.proof.values(), - ) - .expect("failed to verify storage proof"); - } + use crate::storage::test_utils::{create_test_account, create_test_engine, verify_account_proof, verify_storage_proof}; #[test] fn test_get_nonexistent_proof() { diff --git a/src/storage/test_utils.rs b/src/storage/test_utils.rs index 9ccfdc5b..7afefb25 100644 --- a/src/storage/test_utils.rs +++ b/src/storage/test_utils.rs @@ -2,10 +2,12 @@ use crate::{ account::Account, context::TransactionContext, executor::threadpool, meta::MetadataManager, - storage::engine::StorageEngine, PageManager, + storage::{engine::StorageEngine, proofs::{AccountProof, StorageProof}}, + PageManager, }; -use alloy_primitives::U256; -use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY}; +use alloy_primitives::{B256, U256}; +use alloy_rlp::encode; +use alloy_trie::{proof::verify_proof, TrieAccount, EMPTY_ROOT_HASH, KECCAK_EMPTY}; use rand::{rngs::StdRng, RngCore}; pub(crate) fn create_test_engine(max_pages: u32) -> (StorageEngine, TransactionContext) { @@ -60,3 +62,43 @@ pub(crate) fn assert_metrics( "transaction metrics don't match:\n expected: {expected:?}\n actual: {actual:?}" ); } + +/// Verifies that an account proof is valid against the given root hash. +/// +/// This function validates both the account proof itself and any included storage proofs. +/// It will panic with a descriptive message if verification fails. +pub(crate) fn verify_account_proof(proof: &AccountProof, root: B256) { + let expected = Some(encode(TrieAccount { + nonce: proof.account.nonce, + balance: proof.account.balance, + storage_root: proof.account.storage_root, + code_hash: proof.account.code_hash, + })); + verify_proof(root, proof.hashed_address.clone(), expected, proof.proof.values()) + .expect("failed to verify account proof"); + + // Verify all storage proofs against the account's storage root + for storage_proof in proof.storage_proofs.values() { + verify_storage_proof(storage_proof, proof.account.storage_root); + } +} + +/// Verifies that a storage proof is valid against the given storage root hash. +/// +/// This function validates the storage proof, including handling zero values +/// (which are represented as None in the trie). It will panic with a descriptive +/// message if verification fails. +pub(crate) fn verify_storage_proof(proof: &StorageProof, root: B256) { + let expected_value = if proof.value.is_zero() { + None + } else { + Some(alloy_rlp::encode(proof.value)) + }; + verify_proof( + root, + proof.hashed_slot.clone(), + expected_value, + proof.proof.values(), + ) + .expect("failed to verify storage proof"); +} \ No newline at end of file diff --git a/src/transaction.rs b/src/transaction.rs index 17eaeb27..2cd8bdd0 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -8,7 +8,7 @@ use crate::{ node::TrieValue, overlay::OverlayState, path::{AddressPath, RawPath, StoragePath}, - storage::{overlay_root::OverlayedRoot, proofs::AccountProof}, + storage::{overlay::OverlayedRoot, proofs::AccountProof}, }; use alloy_primitives::{map::HashMap, StorageValue, B256}; pub use error::TransactionError;