feat: build state trie from proofs in host instead (#40)

Moves the trie building process (from proofs) from the client to the
host. The client now only verifies the root before and after the state
transition. This saves a lot of cycles as the trie building process
is expensive.

As a result, the account/storage loading logic has also been rewritten,
since the proofs are no longer passed into the client. Instead, the
client receive "state requests" which contains keys (i.e. addresses/
slots), with the actual values read from the trie.

Also switches to use the MPT implementation from `zeth` in `rsp-mpt`.

.github/workflows/pr.yml

@@ -55,7 +55,7 @@ jobs:
 
       - name: "Set up test fixture"
         run: |
-          git clone https://github.com/succinctlabs/rsp-tests --branch 2024-08-31 --depth 1 ../rsp-tests
+          git clone https://github.com/succinctlabs/rsp-tests --branch 2024-09-09 --depth 1 ../rsp-tests
           cd ../rsp-tests/
           docker compose up -d
 

Cargo.toml

@@ -38,6 +38,7 @@ url = "2.3"
 thiserror = "1.0.61"
 hex-literal = "0.4.1"
 rayon = "1.10.0"
+rlp = "0.5.2"
 
 # workspace
 rsp-rpc-db = { path = "./crates/storage/rpc-db" }

crates/executor/client/src/io.rs

@@ -1,10 +1,10 @@
 use std::collections::HashMap;
 
 use eyre::Result;
-use reth_primitives::{revm_primitives::AccountInfo, Address, Block, Bytes, Header, B256, U256};
-use reth_trie::AccountProof;
-use revm_primitives::keccak256;
-use rsp_primitives::account_proof::AccountProofWithBytecode;
+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 serde::{Deserialize, Serialize};
 
@@ -19,14 +19,14 @@ pub struct ClientExecutorInput {
     pub current_block: Block,
     /// The previous block header.
     pub previous_block: Header,
-    /// The dirty storage proofs for the storage slots that were modified.
-    pub dirty_storage_proofs: Vec<AccountProof>,
-    /// The storage proofs for the storage slots that were accessed.
-    pub used_storage_proofs: HashMap<Address, AccountProofWithBytecode>,
+    /// Network state as of the parent block.
+    pub parent_state: EthereumState,
+    /// Requests to account state and storage slots.
+    pub state_requests: HashMap<Address, Vec<U256>>,
+    /// Account bytecodes.
+    pub bytecodes: Vec<Bytecode>,
     /// The block hashes.
     pub block_hashes: HashMap<u64, B256>,
-    /// The trie node preimages.
-    pub trie_nodes: Vec<Bytes>,
 }
 
 impl ClientExecutorInput {
@@ -37,48 +37,62 @@ impl ClientExecutorInput {
     /// to avoid unnecessary cloning.
     pub fn witness_db(&mut self) -> Result<WitnessDb> {
         let state_root: B256 = self.previous_block.state_root;
+        if state_root != self.parent_state.state_root() {
+            eyre::bail!("parent state root mismatch");
+        }
+
+        let bytecodes_by_hash =
+            self.bytecodes.iter().map(|code| (code.hash_slow(), code)).collect::<HashMap<_, _>>();
 
         let mut accounts = HashMap::new();
         let mut storage = HashMap::new();
-        let used_storage_proofs = std::mem::take(&mut self.used_storage_proofs);
-        for (address, proof) in used_storage_proofs {
-            // Verify the storage proof.
-            proof.verify(state_root)?;
+        let state_requests = std::mem::take(&mut self.state_requests);
+        for (address, slots) in state_requests {
+            let hashed_address = keccak256(address);
+            let hashed_address = hashed_address.as_slice();
+
+            let account_in_trie =
+                self.parent_state.state_trie.get_rlp::<TrieAccount>(hashed_address)?;
 
-            // Update the accounts.
-            let account_info = match proof.proof.info {
-                Some(account_info) => AccountInfo {
-                    nonce: account_info.nonce,
-                    balance: account_info.balance,
-                    code_hash: account_info.bytecode_hash.unwrap(),
-                    code: Some(proof.code),
+            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_else(|| eyre::eyre!("missing bytecode"))?)
+                            // Cloning here is fine as `Bytes` is cheap to clone.
+                            .to_owned(),
+                        ),
+                    },
+                    None => Default::default(),
                 },
-                None => AccountInfo::default(),
-            };
-            accounts.insert(address, account_info);
+            );
 
-            // Update the storage.
-            let storage_values: HashMap<U256, U256> = proof
-                .proof
-                .storage_proofs
-                .into_iter()
-                .map(|storage_proof| (storage_proof.key.into(), storage_proof.value))
-                .collect();
-            storage.insert(address, storage_values);
-        }
+            if !slots.is_empty() {
+                let mut address_storage = HashMap::new();
+
+                let storage_trie = self
+                    .parent_state
+                    .storage_tries
+                    .get(hashed_address)
+                    .ok_or_else(|| eyre::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);
+                }
 
-        let mut trie_nodes = HashMap::new();
-        for preimage in self.trie_nodes.iter() {
-            // TODO: refactor witness db building to avoid cloning and `mem::take`.
-            trie_nodes.insert(keccak256(preimage), preimage.to_owned());
+                storage.insert(address, address_storage);
+            }
         }
 
-        Ok(WitnessDb {
-            accounts,
-            storage,
-            block_hashes: std::mem::take(&mut self.block_hashes),
-            state_root: self.current_block.state_root,
-            trie_nodes,
-        })
+        Ok(WitnessDb { accounts, storage, block_hashes: std::mem::take(&mut self.block_hashes) })
     }
 }

crates/executor/client/src/lib.rs

@@ -146,8 +146,10 @@ impl ClientExecutor {
 
         // Verify the state root.
         let state_root = profile!("compute state root", {
-            rsp_mpt::compute_state_root(&executor_outcome, &input.dirty_storage_proofs, &witness_db)
-        })?;
+            input.parent_state.update(&executor_outcome.hash_state_slow());
+            input.parent_state.state_root()
+        });
+
         if state_root != input.current_block.state_root {
             eyre::bail!("mismatched state root");
         }

crates/executor/host/src/lib.rs

@@ -1,15 +1,15 @@
-use std::marker::PhantomData;
+use std::{collections::BTreeSet, marker::PhantomData};
 
 use alloy_provider::{network::AnyNetwork, Provider};
 use alloy_transport::Transport;
 use eyre::{eyre, Ok};
-use itertools::Itertools;
 use reth_execution_types::ExecutionOutcome;
 use reth_primitives::{proofs, Block, Bloom, Receipts, B256};
 use revm::db::CacheDB;
 use rsp_client_executor::{
     io::ClientExecutorInput, ChainVariant, EthereumVariant, LineaVariant, OptimismVariant, Variant,
 };
+use rsp_mpt::EthereumState;
 use rsp_primitives::account_proof::eip1186_proof_to_account_proof;
 use rsp_rpc_db::RpcDb;
 
@@ -112,27 +112,61 @@ impl<T: Transport + Clone, P: Provider<T, AnyNetwork> + Clone> HostExecutor<T, P
             vec![executor_output.requests.into()],
         );
 
+        let state_requests = rpc_db.get_state_requests();
+
         // For every account we touched, fetch the storage proofs for all the slots we touched.
-        tracing::info!("fetching modified storage proofs");
-        let mut dirty_storage_proofs = Vec::new();
-        for (address, account) in executor_outcome.bundle_accounts_iter() {
-            let mut storage_keys = Vec::new();
-            for key in account.storage.keys().sorted() {
-                let slot = B256::new(key.to_be_bytes());
-                storage_keys.push(slot);
-            }
+        tracing::info!("fetching storage proofs");
+        let mut before_storage_proofs = Vec::new();
+        let mut after_storage_proofs = Vec::new();
+
+        for (address, used_keys) in state_requests.iter() {
+            let modified_keys = executor_outcome
+                .state()
+                .state
+                .get(address)
+                .map(|account| {
+                    account.storage.keys().map(|key| B256::from(*key)).collect::<BTreeSet<_>>()
+                })
+                .unwrap_or_default()
+                .into_iter()
+                .collect::<Vec<_>>();
+
+            let keys = used_keys
+                .iter()
+                .map(|key| B256::from(*key))
+                .chain(modified_keys.clone().into_iter())
+                .collect::<BTreeSet<_>>()
+                .into_iter()
+                .collect::<Vec<_>>();
+
             let storage_proof = self
                 .provider
-                .get_proof(address, storage_keys)
+                .get_proof(*address, keys.clone())
                 .block_id((block_number - 1).into())
                 .await?;
-            dirty_storage_proofs.push(eip1186_proof_to_account_proof(storage_proof));
+            before_storage_proofs.push(eip1186_proof_to_account_proof(storage_proof));
+
+            let storage_proof = self
+                .provider
+                .get_proof(*address, modified_keys)
+                .block_id((block_number).into())
+                .await?;
+            after_storage_proofs.push(eip1186_proof_to_account_proof(storage_proof));
         }
 
+        let state = EthereumState::from_proofs(
+            previous_block.state_root,
+            &before_storage_proofs.iter().map(|item| (item.address, item.clone())).collect(),
+            &after_storage_proofs.iter().map(|item| (item.address, item.clone())).collect(),
+        )?;
+
         // Verify the state root.
         tracing::info!("verifying the state root");
-        let state_root =
-            rsp_mpt::compute_state_root(&executor_outcome, &dirty_storage_proofs, &rpc_db)?;
+        let state_root = {
+            let mut mutated_state = state.clone();
+            mutated_state.update(&executor_outcome.hash_state_slow());
+            mutated_state.state_root()
+        };
         if state_root != current_block.state_root {
             eyre::bail!("mismatched state root");
         }
@@ -167,12 +201,12 @@ impl<T: Transport + Clone, P: Provider<T, AnyNetwork> + Clone> HostExecutor<T, P
 
         // Create the client input.
         let client_input = ClientExecutorInput {
-            previous_block: previous_block.header,
             current_block: V::pre_process_block(&current_block),
-            dirty_storage_proofs,
-            used_storage_proofs: rpc_db.fetch_used_accounts_and_proofs().await,
+            previous_block: previous_block.header,
+            parent_state: state,
+            state_requests,
+            bytecodes: rpc_db.get_bytecodes(),
             block_hashes: rpc_db.block_hashes.borrow().clone(),
-            trie_nodes: rpc_db.trie_nodes.borrow().values().cloned().collect(),
         };
         Ok(client_input)
     }

crates/mpt/Cargo.toml

@@ -11,7 +11,11 @@ repository.workspace = true
 workspace = true
 
 [dependencies]
+anyhow = { workspace = true, features = ["std"] }
 eyre.workspace = true
+rlp.workspace = true
+serde.workspace = true
+thiserror.workspace = true
 itertools = "0.13.0"
 
 # workspace
@@ -23,6 +27,7 @@ reth-trie.workspace = true
 reth-execution-types.workspace = true
 
 # revm
+revm.workspace = true
 revm-primitives.workspace = true
 
 # alloy