feat: build state trie from proofs in host instead

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