Support computing state root with overlayed state

benches/crud_benchmarks.rs

@@ -20,10 +20,17 @@ use alloy_primitives::{StorageKey, StorageValue, U256};
 use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY};
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use rand::prelude::*;
-use std::{fs, io, path::Path, time::Duration};
+use std::{
+    fs, io,
+    path::Path,
+    sync::{Arc, Barrier},
+    thread,
+    time::Duration,
+};
 use tempdir::TempDir;
 use triedb::{
     account::Account,
+    overlay::{OverlayStateMut, OverlayValue},
     path::{AddressPath, StoragePath},
     Database,
 };
@@ -82,6 +89,65 @@ fn bench_account_reads(c: &mut Criterion) {
             },
         );
     });
+    group.bench_function(BenchmarkId::new("eoa_reads_parallel", BATCH_SIZE), |b| {
+        b.iter_with_setup(
+            || {
+                let db_path = dir.path().join(&file_name);
+                let db = Arc::new(Database::open(db_path.clone()).unwrap());
+
+                // Spawn 4 reader threads
+                let thread_count = 4;
+                let setup_barrier = Arc::new(Barrier::new(thread_count + 1));
+                let test_barrier = Arc::new(Barrier::new(thread_count + 1));
+                let chunk_size = addresses.len() / thread_count;
+                let mut handles = Vec::new();
+
+                for i in 0..thread_count {
+                    let start_idx = i * chunk_size;
+                    let end_idx = if i == thread_count {
+                        // Last thread handles any remaining addresses
+                        addresses.len()
+                    } else {
+                        (i + 1) * chunk_size
+                    };
+
+                    let thread_addresses = addresses[start_idx..end_idx].to_vec();
+                    let db_clone = Arc::clone(&db);
+                    let setup_barrier = Arc::clone(&setup_barrier);
+                    let test_barrier = Arc::clone(&test_barrier);
+
+                    let handle = thread::spawn(move || {
+                        setup_barrier.wait();
+                        test_barrier.wait();
+                        // Each thread creates its own RO transaction
+                        let mut tx = db_clone.begin_ro().unwrap();
+
+                        // Read its chunk of addresses
+                        for addr in thread_addresses {
+                            let a = tx.get_account(&addr).unwrap();
+                            assert!(a.is_some());
+                        }
+
+                        // Commit the transaction
+                        tx.commit().unwrap();
+                    });
+
+                    handles.push(handle);
+                }
+
+                setup_barrier.wait();
+
+                (handles, test_barrier)
+            },
+            |(handles, test_barrier)| {
+                test_barrier.wait();
+                // Wait for all threads to complete
+                for handle in handles {
+                    handle.join().unwrap();
+                }
+            },
+        );
+    });
 
     group.finish();
 }
@@ -191,7 +257,7 @@ fn bench_account_updates(c: &mut Criterion) {
         b.iter_with_setup(
             || {
                 let db_path = dir.path().join(&file_name);
-                Database::open(db_path.clone()).unwrap()
+                Database::open(db_path).unwrap()
             },
             |db| {
                 let mut tx = db.begin_rw().unwrap();
@@ -538,6 +604,51 @@ fn bench_storage_deletes(c: &mut Criterion) {
     group.finish();
 }
 
+fn bench_state_root_with_overlay(c: &mut Criterion) {
+    let mut group = c.benchmark_group("state_root_with_overlay");
+    let base_dir = get_base_database(
+        DEFAULT_SETUP_DB_EOA_SIZE,
+        DEFAULT_SETUP_DB_CONTRACT_SIZE,
+        DEFAULT_SETUP_DB_STORAGE_PER_CONTRACT,
+    );
+    let dir = TempDir::new("triedb_bench_state_root_with_overlay").unwrap();
+    let file_name = base_dir.main_file_name.clone();
+    copy_files(&base_dir, dir.path()).unwrap();
+
+    let mut rng = StdRng::seed_from_u64(SEED_CONTRACT);
+    let addresses: Vec<AddressPath> =
+        (0..BATCH_SIZE).map(|_| generate_random_address(&mut rng)).collect();
+
+    let mut account_overlay_mut = OverlayStateMut::new();
+    addresses.iter().enumerate().for_each(|(i, addr)| {
+        let new_account =
+            Account::new(i as u64, U256::from(i as u64), EMPTY_ROOT_HASH, KECCAK_EMPTY);
+        account_overlay_mut.insert(addr.clone().into(), Some(OverlayValue::Account(new_account)));
+    });
+    let account_overlay = account_overlay_mut.freeze();
+
+    group.throughput(criterion::Throughput::Elements(BATCH_SIZE as u64));
+    group.measurement_time(Duration::from_secs(30));
+    group.bench_function(BenchmarkId::new("state_root_with_account_overlay", BATCH_SIZE), |b| {
+        b.iter_with_setup(
+            || {
+                let db_path = dir.path().join(&file_name);
+                Database::open(db_path).unwrap()
+            },
+            |db| {
+                let tx = db.begin_ro().unwrap();
+
+                // Compute the root hash with the overlay
+                let _root_result = tx.compute_root_with_overlay(account_overlay.clone()).unwrap();
+
+                tx.commit().unwrap();
+            },
+        );
+    });
+
+    group.finish();
+}
+
 criterion_group!(
     benches,
     bench_account_reads,
@@ -550,5 +661,6 @@ criterion_group!(
     bench_storage_inserts,
     bench_storage_updates,
     bench_storage_deletes,
+    bench_state_root_with_overlay,
 );
 criterion_main!(benches);

src/account.rs

@@ -1,9 +1,10 @@
 use alloy_primitives::{B256, U256};
+use alloy_rlp::{MaxEncodedLen, RlpEncodable};
 use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY};
 use proptest::prelude::*;
 use proptest_derive::Arbitrary;
 
-#[derive(Debug, Clone, PartialEq, Eq, Default, Arbitrary)]
+#[derive(Debug, Clone, PartialEq, Eq, Default, Arbitrary, RlpEncodable)]
 pub struct Account {
     pub nonce: u64,
     pub balance: U256,
@@ -19,6 +20,16 @@ impl Account {
     }
 }
 
+/// This is the maximum possible RLP-encoded length of an account.
+///
+/// This value is derived from the maximum possible length of an account, which is the largest
+/// case. An account is encoded as a list of 4 elements, with 3 of these represnting 32 byte values
+/// and the nonce being an 8 byte value. Each element has 1 extra byte of encoding overhead.
+/// The list also has 2 bytes of encoding overhead. The total length is `2 + 3*33 + 9 = 110`.
+const MAX_RLP_ENCODED_LEN: usize = 2 + 3 * 33 + 9;
+
+unsafe impl MaxEncodedLen<MAX_RLP_ENCODED_LEN> for Account {}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/lib.rs

@@ -17,6 +17,7 @@ pub mod location;
 pub mod meta;
 pub mod metrics;
 pub mod node;
+pub mod overlay;
 pub mod page;
 pub mod path;
 pub mod pointer;

src/node.rs

@@ -3,7 +3,7 @@ use crate::{
     pointer::Pointer,
     storage::value::{self, Value},
 };
-use alloy_primitives::{hex, StorageValue, B256, U256};
+use alloy_primitives::{StorageValue, B256, U256};
 use alloy_rlp::{
     decode_exact, encode_fixed_size, length_of_length, BufMut, Encodable, Header, MaxEncodedLen,
     EMPTY_STRING_CODE,
@@ -17,11 +17,6 @@ use proptest::{arbitrary, strategy, strategy::Strategy};
 use proptest_derive::Arbitrary;
 use std::cmp::{max, min};
 
-// This is equivalent to RlpNode::word_rlp(&EMPTY_ROOT_HASH), and is used to encode the storage root
-// of an account with no storage.
-const EMPTY_ROOT_RLP: [u8; 33] =
-    hex!("0xa056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421");
-
 const MAX_PREFIX_LENGTH: usize = 64;
 
 /// A node in the trie.
@@ -152,6 +147,10 @@ impl Node {
         &self.0.kind
     }
 
+    pub fn into_kind(self) -> NodeKind {
+        self.0.kind
+    }
+
     /// Returns whether the [Node] type supports children.
     pub const fn has_children(&self) -> bool {
         matches!(self.kind(), NodeKind::Branch { .. } | NodeKind::AccountLeaf { .. })
@@ -567,17 +566,17 @@ impl Encodable for Node {
                 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 value_rlp = buf.as_mut();
-                let storage_root_rlp =
-                    storage_root.as_ref().map(|p| p.rlp().as_slice()).unwrap_or(&EMPTY_ROOT_RLP);
-                let len = 2 + nonce_rlp.len() + balance_rlp.len() + storage_root_rlp.len() + 33;
-                value_rlp.put_u8(0xf8);
-                value_rlp.put_u8((len - 2) as u8);
-                value_rlp.put_slice(nonce_rlp);
-                value_rlp.put_slice(balance_rlp);
-                value_rlp.put_slice(storage_root_rlp);
-                value_rlp.put_u8(0xa0);
-                value_rlp.put_slice(code_hash.as_slice());
-                LeafNodeRef { key: prefix, value: &buf[..len] }.encode(out);
+                let storage_root_hash = storage_root
+                    .as_ref()
+                    .map_or(EMPTY_ROOT_HASH, |p| p.rlp().as_hash().unwrap_or(EMPTY_ROOT_HASH));
+                let account_rlp_length = encode_account_leaf(
+                    nonce_rlp,
+                    balance_rlp,
+                    code_hash,
+                    &storage_root_hash,
+                    &mut value_rlp,
+                );
+                LeafNodeRef { key: prefix, value: &buf[..account_rlp_length] }.encode(out);
             }
             NodeKind::Branch { ref children } => {
                 if prefix.is_empty() {
@@ -634,6 +633,28 @@ impl Encodable for Node {
     }
 }
 
+#[inline]
+pub fn encode_account_leaf(
+    nonce_rlp: &ArrayVec<u8, 9>,
+    balance_rlp: &ArrayVec<u8, 33>,
+    code_hash: &B256,
+    storage_root: &B256,
+    out: &mut dyn BufMut,
+) -> usize {
+    let len = 2 + nonce_rlp.len() + balance_rlp.len() + 33 * 2;
+    out.put_u8(0xf8);
+    out.put_u8((len - 2) as u8);
+    out.put_slice(nonce_rlp);
+    out.put_slice(balance_rlp);
+    out.put_u8(0xa0);
+    out.put_slice(storage_root.as_slice());
+    out.put_u8(0xa0);
+    out.put_slice(code_hash.as_slice());
+
+    len
+}
+
+#[inline]
 pub fn encode_branch(children: &[Option<Pointer>], out: &mut dyn BufMut) -> usize {
     // first encode the header
     let mut payload_length = 1;