feat: small performance optimizations (#493)

* perf: improve allocs

- reserve a bigger vector for trie nodes from the beginning. remove capacity increase in inserts.
- use an uninit array for branch childs during decoding

* perf: decode keccak alignment

Encode RLP representation of each node with extra padding to make them word-aligned. This let us to avoid memcpy operations when calculating `keccak(rlp_encoded)` during decoding.

* feat: improve witness db

- Revm's `CacheDB` already caches state's keys and values. We remove the duplicate caching in witness db by simplifying it.
- Add a keccak cache of address and storage slots. With this we avoid hashing some addresses and storage slots more than once.

* fix: fmt

* perf: remove keccak cache

* chore: add `SAFETY` comments

Author: shayanh

Cargo.lock

@@ -7,7 +7,6 @@ members = [
     "crates/mpt-tools",
     "crates/primitives",
     "crates/storage/rpc-db",
-    "crates/storage/witness-db",
     "crates/host-bench",
     "crates/mptnew",
     "bin/verifier-bench",
@@ -56,7 +55,6 @@ bitcode = { version = "0.6.5", default-features = false, features = ["serde"] }
 
 # workspace
 openvm-rpc-db = { path = "./crates/storage/rpc-db" }
-openvm-witness-db = { path = "./crates/storage/witness-db" }
 openvm-client-executor = { path = "./crates/executor/client" }
 openvm-host-executor = { path = "./crates/executor/host" }
 openvm-mpt = { path = "./crates/mpt" }

Cargo.toml

@@ -17,7 +17,6 @@ bumpalo.workspace = true
 itertools = "0.14"
 
 # workspace
-openvm-witness-db.workspace = true
 openvm-mpt.workspace = true
 openvm-primitives.workspace = true
 mptnew.workspace = true

bin/client-eth/Cargo.lock

@@ -1,12 +1,15 @@
 use std::iter::once;
 
 use bumpalo::Bump;
-use eyre::{bail, OptionExt, Result};
+use eyre::{bail, Result};
 use itertools::Itertools;
-use openvm_witness_db::WitnessDb;
+use reth_evm::execute::ProviderError;
 use reth_primitives::{Block, Header, TransactionSigned};
 use reth_trie::TrieAccount;
-use revm::state::{AccountInfo, Bytecode};
+use revm::{
+    state::{AccountInfo, Bytecode},
+    DatabaseRef,
+};
 use revm_primitives::{keccak256, map::DefaultHashBuilder, Address, HashMap, B256, U256};
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
@@ -30,8 +33,6 @@ pub struct ClientExecutorInput {
     pub ancestor_headers: Vec<Header>,
     /// Network state as of the parent block.
     pub parent_state_bytes: mptnew::EthereumStateBytes,
-    /// Requests to account state and storage slots.
-    pub state_requests: HashMap<Address, Vec<U256>>,
     /// Account bytecodes.
     pub bytecodes: Vec<Bytecode>,
 }
@@ -65,7 +66,7 @@ impl ClientExecutorInputWithState {
                 &input.parent_state_bytes.storage_tries
             {
                 let account_in_trie =
-                    state_trie.get_rlp::<TrieAccount>(hashed_address.as_slice()).unwrap();
+                    state_trie.get_rlp::<TrieAccount>(hashed_address.as_slice())?;
                 let expected_storage_root =
                     account_in_trie.map_or(reth_trie::EMPTY_ROOT_HASH, |a| a.storage_root);
 
@@ -80,6 +81,7 @@ impl ClientExecutorInputWithState {
 
                 storage_tries.insert(*hashed_address, storage_trie);
             }
+
             mptnew::EthereumState { state_trie, storage_tries, bump }
         };
 
@@ -95,7 +97,7 @@ impl ClientExecutorInputWithState {
     }
 
     /// Creates a [`WitnessDb`].
-    pub fn witness_db(&self) -> Result<WitnessDb> {
+    pub fn witness_db(&self) -> Result<WitnessDb<'_>> {
         <Self as WitnessInput>::witness_db(self)
     }
 }
@@ -111,11 +113,6 @@ impl WitnessInput for ClientExecutorInputWithState {
         self.parent_header().state_root
     }
 
-    #[inline(always)]
-    fn state_requests(&self) -> impl Iterator<Item = (&Address, &Vec<U256>)> {
-        self.input.state_requests.iter()
-    }
-
     #[inline(always)]
     fn bytecodes(&self) -> impl Iterator<Item = &Bytecode> {
         self.input.bytecodes.iter()
@@ -125,6 +122,11 @@ impl WitnessInput for ClientExecutorInputWithState {
     fn headers(&self) -> impl Iterator<Item = &Header> {
         once(&self.input.current_block.header).chain(self.input.ancestor_headers.iter())
     }
+
+    #[inline(always)]
+    fn headers_len(&self) -> usize {
+        1 + self.input.ancestor_headers.len()
+    }
 }
 
 /// A trait for constructing [`WitnessDb`].
@@ -136,18 +138,16 @@ pub trait WitnessInput {
     /// [state()](trait.WitnessInput#tymethod.state) must conform to.
     fn state_anchor(&self) -> B256;
 
-    /// Gets an iterator over address state requests. For each request, the account info and storage
-    /// slots are loaded from the relevant tries in the state returned by
-    /// [state()](trait.WitnessInput#tymethod.state).
-    fn state_requests(&self) -> impl Iterator<Item = (&Address, &Vec<U256>)>;
-
     /// Gets an iterator over account bytecodes.
     fn bytecodes(&self) -> impl Iterator<Item = &Bytecode>;
 
     /// Gets an iterator over references to a consecutive, reverse-chronological block headers
     /// starting from the current block header.
     fn headers(&self) -> impl Iterator<Item = &Header>;
 
+    /// Gets the number of headers.
+    fn headers_len(&self) -> usize;
+
     /// Creates a [`WitnessDb`] from a [`WitnessInput`] implementation. To do so, it verifies the
     /// state root, ancestor headers and account bytecodes, and constructs the account and
     /// storage values by reading against state tries.
@@ -156,68 +156,16 @@ pub trait WitnessInput {
     /// implementing this trait causes a zkVM run to cost over 5M cycles more. To avoid this, define
     /// a method inside the type that calls this trait method instead.
     #[inline(always)]
-    fn witness_db(&self) -> Result<WitnessDb> {
+    fn witness_db(&self) -> Result<WitnessDb<'_>> {
         let state = self.state();
 
-        let bytecodes_by_hash =
+        let bytecode_by_hash =
             self.bytecodes().map(|code| (code.hash_slow(), code)).collect::<HashMap<_, _>>();
 
-        let state_requests_iter = self.state_requests();
-        let (lower, _) = state_requests_iter.size_hint();
-        let mut accounts = HashMap::with_capacity_and_hasher(lower, DefaultHashBuilder::default());
-        let mut storage = HashMap::with_capacity_and_hasher(lower, DefaultHashBuilder::default());
-
-        for (&address, slots) in state_requests_iter {
-            let hashed_address = keccak256(address);
-
-            let account_in_trie =
-                state.state_trie.get_rlp::<TrieAccount>(hashed_address.as_slice())?;
-
-            accounts.insert(
-                address,
-                match account_in_trie {
-                    Some(account_in_trie) => AccountInfo {
-                        balance: account_in_trie.balance,
-                        nonce: account_in_trie.nonce,
-                        code_hash: account_in_trie.code_hash,
-                        code: Some(
-                            (*bytecodes_by_hash
-                                .get(&account_in_trie.code_hash)
-                                .ok_or_eyre("missing bytecode")?)
-                            // Cloning here is fine as `Bytes` is cheap to clone.
-                            .to_owned(),
-                        ),
-                    },
-                    _ => Default::default(),
-                },
-            );
-
-            if !slots.is_empty() {
-                let mut address_storage =
-                    HashMap::with_capacity_and_hasher(slots.len(), DefaultHashBuilder::default());
-
-                let storage_trie = state
-                    .storage_tries
-                    .get(&hashed_address)
-                    .ok_or_eyre("parent state does not contain storage trie")?;
-
-                for &slot in slots {
-                    let slot_value = storage_trie
-                        .get_rlp::<U256>(keccak256(slot.to_be_bytes::<32>()).as_slice())?
-                        .unwrap_or_default();
-                    address_storage.insert(slot, slot_value);
-                }
-
-                storage.insert(address, address_storage);
-            }
-        }
-
         // Verify and build block hashes
-        let headers_iter = self.headers();
-        let (lower, _) = headers_iter.size_hint();
         let mut block_hashes: HashMap<u64, B256, _> =
-            HashMap::with_capacity_and_hasher(lower, DefaultHashBuilder::default());
-        for (child_header, parent_header) in headers_iter.tuple_windows() {
+            HashMap::with_capacity_and_hasher(self.headers_len(), DefaultHashBuilder::default());
+        for (child_header, parent_header) in self.headers().tuple_windows() {
             if parent_header.number != child_header.number - 1 {
                 eyre::bail!("non-consecutive blocks");
             }
@@ -229,6 +177,76 @@ pub trait WitnessInput {
             block_hashes.insert(parent_header.number, child_header.parent_hash);
         }
 
-        Ok(WitnessDb { accounts, storage, block_hashes })
+        Ok(WitnessDb { inner: state, block_hashes, bytecode_by_hash })
+    }
+}
+
+#[derive(Debug)]
+pub struct WitnessDb<'a> {
+    inner: &'a mptnew::EthereumState,
+    block_hashes: HashMap<u64, B256>,
+    bytecode_by_hash: HashMap<B256, &'a Bytecode>,
+}
+
+impl<'a> WitnessDb<'a> {
+    pub fn new(
+        inner: &'a mptnew::EthereumState,
+        block_hashes: HashMap<u64, B256>,
+        bytecode_by_hash: HashMap<B256, &'a Bytecode>,
+    ) -> Self {
+        Self { inner, block_hashes, bytecode_by_hash }
+    }
+}
+
+impl DatabaseRef for WitnessDb<'_> {
+    /// The database error type.
+    type Error = ProviderError;
+
+    /// Get basic account information.
+    fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
+        let hashed_address = keccak256(address);
+
+        let account_in_trie =
+            self.inner.state_trie.get_rlp::<TrieAccount>(hashed_address.as_slice()).unwrap();
+
+        let account = account_in_trie.map(|account_in_trie| AccountInfo {
+            balance: account_in_trie.balance,
+            nonce: account_in_trie.nonce,
+            code_hash: account_in_trie.code_hash,
+            code: None,
+        });
+
+        Ok(account)
+    }
+
+    /// Get account code by its hash.
+    fn code_by_hash_ref(&self, hash: B256) -> Result<Bytecode, Self::Error> {
+        // Cloning here is fine as `Bytes` is cheap to clone.
+        Ok(self.bytecode_by_hash.get(&hash).map(|code| (*code).clone()).unwrap())
+    }
+
+    /// Get storage value of address at index.
+    fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
+        let hashed_address = keccak256(address);
+
+        let storage_trie = self
+            .inner
+            .storage_tries
+            .get(&hashed_address)
+            .expect("A storage trie must be provided for each account");
+
+        let hashed_slot = keccak256(index.to_be_bytes::<32>());
+        Ok(storage_trie
+            .get_rlp::<U256>(hashed_slot.as_slice())
+            .expect("Can get from MPT")
+            .unwrap_or_default())
+    }
+
+    /// Get block hash by block number.
+    fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
+        Ok(*self
+            .block_hashes
+            .get(&number)
+            .expect("A block hash must be provided for each block number"))
     }
 }

crates/executor/client/Cargo.toml

@@ -103,6 +103,8 @@ impl ClientExecutor {
             vec![executor_output.result.requests],
         );
 
+        drop(witness_db);
+
         // Verify the state root.
         let state_root = profile!("compute state root", {
             input.state.update_from_bundle_state(&executor_outcome.bundle).unwrap();

crates/executor/client/src/io.rs

@@ -215,7 +215,6 @@ impl<P: Provider<Ethereum> + Clone> HostExecutor<P> {
             current_block,
             ancestor_headers,
             parent_state_bytes: state_bytes,
-            state_requests,
             bytecodes: rpc_db.get_bytecodes(),
         };
         tracing::info!("successfully generated client input");

crates/executor/client/src/lib.rs

@@ -865,25 +865,26 @@ impl MptTrie<'_> {
     }
 
     pub fn encode_trie(&self) -> Vec<u8> {
-        let mut payload = Vec::new();
-        self.encode_trie_internal(self.root_id, &mut payload);
-
         let mut encoded = Vec::new();
-        alloy_rlp::Header { list: true, payload_length: payload.len() }.encode(&mut encoded);
-        encoded.append(&mut payload);
+        self.encode_trie_internal(self.root_id, &mut encoded);
         encoded
     }
 
     fn encode_trie_internal(&self, id: NodeId, out: &mut dyn alloy_rlp::BufMut) {
         let payload_length = self.payload_length_id(id);
         self.encode_id_with_payload_len(id, payload_length, out);
 
+        const MIN_ALIGN: usize = 4;
+        let rlp_length = payload_length + alloy_rlp::length_of_length(payload_length);
+        let padding_len = (MIN_ALIGN - (rlp_length % MIN_ALIGN)) % MIN_ALIGN;
+        for _ in 0..padding_len {
+            out.put_u8(0);
+        }
+
         match self.nodes[id as usize] {
             NodeData::Branch(childs) => childs.iter().for_each(|c| {
                 if let Some(node) = c {
                     self.encode_trie_internal(*node, out);
-                } else {
-                    out.put_u8(alloy_rlp::EMPTY_STRING_CODE)
                 }
             }),
             NodeData::Extension(_, node) => {