feat: trie db (#71)

* feat: initial trie db

* fix: clippy + fmt

---------

Co-authored-by: leruaa <[email protected]>

.gitignore

@@ -23,3 +23,5 @@ latest_proof.json
 elf
 
 **.csv
+
+input/

crates/executor/client/src/io.rs

@@ -1,11 +1,13 @@
 use std::{collections::HashMap, iter::once};
 
 use itertools::Itertools;
+use reth_errors::ProviderError;
 use reth_primitives::{revm_primitives::AccountInfo, Address, Block, Header, B256, U256};
 use reth_trie::TrieAccount;
 use revm_primitives::{keccak256, Bytecode};
 use rsp_mpt::EthereumState;
-use rsp_witness_db::WitnessDb;
+//use rsp_witness_db::WitnessDb;
+use revm::DatabaseRef;
 use serde::{Deserialize, Serialize};
 
 use crate::error::ClientError;
@@ -38,7 +40,7 @@ impl ClientExecutorInput {
     }
 
     /// Creates a [`WitnessDb`].
-    pub fn witness_db(&self) -> Result<WitnessDb, ClientError> {
+    pub fn witness_db(&self) -> Result<TrieDB<'_>, ClientError> {
         <Self as WitnessInput>::witness_db(self)
     }
 }
@@ -70,6 +72,75 @@ impl WitnessInput for ClientExecutorInput {
     }
 }
 
+#[derive(Debug)]
+pub struct TrieDB<'a> {
+    inner: &'a EthereumState,
+    block_hashes: HashMap<u64, B256>,
+    bytecode_by_hash: HashMap<B256, &'a Bytecode>,
+}
+
+impl<'a> TrieDB<'a> {
+    pub fn new(
+        inner: &'a EthereumState,
+        block_hashes: HashMap<u64, B256>,
+        bytecode_by_hash: HashMap<B256, &'a Bytecode>,
+    ) -> Self {
+        Self { inner, block_hashes, bytecode_by_hash }
+    }
+}
+
+impl<'a> DatabaseRef for TrieDB<'a> {
+    /// 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 hashed_address = hashed_address.as_slice();
+
+        let account_in_trie = self.inner.state_trie.get_rlp::<TrieAccount>(hashed_address).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> {
+        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 hashed_address = hashed_address.as_slice();
+
+        let storage_trie = self
+            .inner
+            .storage_tries
+            .get(hashed_address)
+            .expect("A storage trie must be provided for each account");
+
+        Ok(storage_trie
+            .get_rlp::<U256>(keccak256(index.to_be_bytes::<32>()).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"))
+    }
+}
+
 /// A trait for constructing [`WitnessDb`].
 pub trait WitnessInput {
     /// Gets a reference to the state from which account info and storage slots are loaded.
@@ -99,7 +170,7 @@ 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, ClientError> {
+    fn witness_db(&self) -> Result<TrieDB<'_>, ClientError> {
         let state = self.state();
 
         if self.state_anchor() != state.state_root() {
@@ -109,52 +180,6 @@ pub trait WitnessInput {
         let bytecodes_by_hash =
             self.bytecodes().map(|code| (code.hash_slow(), code)).collect::<HashMap<_, _>>();
 
-        let mut accounts = HashMap::new();
-        let mut storage = HashMap::new();
-        for (&address, slots) in self.state_requests() {
-            let hashed_address = keccak256(address);
-            let hashed_address = hashed_address.as_slice();
-
-            let account_in_trie = state.state_trie.get_rlp::<TrieAccount>(hashed_address)?;
-
-            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(ClientError::MissingBytecode(address))?)
-                            // Cloning here is fine as `Bytes` is cheap to clone.
-                            .to_owned(),
-                        ),
-                    },
-                    None => Default::default(),
-                },
-            );
-
-            if !slots.is_empty() {
-                let mut address_storage = HashMap::new();
-
-                let storage_trie = state
-                    .storage_tries
-                    .get(hashed_address)
-                    .ok_or(ClientError::MissingTrie(address))?;
-
-                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 mut block_hashes: HashMap<u64, B256> = HashMap::new();
         for (child_header, parent_header) in self.headers().tuple_windows() {
@@ -176,6 +201,6 @@ pub trait WitnessInput {
             block_hashes.insert(parent_header.number, child_header.parent_hash);
         }
 
-        Ok(WitnessDb { accounts, storage, block_hashes })
+        Ok(TrieDB::new(state, block_hashes, bytecodes_by_hash))
     }
 }

crates/executor/client/src/lib.rs

@@ -21,7 +21,7 @@ use reth_evm_optimism::OpExecutorProvider;
 use reth_execution_types::ExecutionOutcome;
 use reth_optimism_consensus::validate_block_post_execution as validate_block_post_execution_optimism;
 use reth_primitives::{proofs, Block, BlockWithSenders, Bloom, Header, Receipt, Receipts, Request};
-use revm::{db::CacheDB, Database};
+use revm::{db::WrapDatabaseRef, Database};
 use revm_primitives::{address, U256};
 
 /// Chain ID for Ethereum Mainnet.
@@ -112,8 +112,10 @@ impl ClientExecutor {
         V: Variant,
     {
         // Initialize the witnessed database with verified storage proofs.
-        let witness_db = input.witness_db()?;
-        let cache_db = CacheDB::new(&witness_db);
+        let wrap_ref = profile!("initialize witness db", {
+            let trie_db = input.witness_db().unwrap();
+            WrapDatabaseRef(trie_db)
+        });
 
         // Execute the block.
         let spec = V::spec();
@@ -126,7 +128,7 @@ impl ClientExecutor {
         })?;
         let executor_difficulty = input.current_block.header.difficulty;
         let executor_output = profile!("execute", {
-            V::execute(&executor_block_input, executor_difficulty, cache_db)
+            V::execute(&executor_block_input, executor_difficulty, wrap_ref)
         })?;
 
         // Validate the block post execution.