feat(trie): poc implementation of rsp trie

Cargo.toml

@@ -173,7 +173,7 @@ members = [
     "testing/ef-tests/",
     "testing/testing-utils",
     "testing/runner",
-    "crates/tracing-otlp",
+    "crates/tracing-otlp", "crates/trie/linked",
 ]
 default-members = ["bin/reth"]
 exclude = ["docs/cli"]

crates/stateless/src/validation.rs

@@ -86,6 +86,19 @@ pub enum StatelessValidationError {
         expected: B256,
     },
 
+    /// Error when the computed pre-storage root does not match the expected one.
+    #[error(
+        "mismatched pre-storage root of account with hash {address_hash}: {got} \n {expected}"
+    )]
+    PreStorageRootMismatch {
+        /// The account address hash
+        address_hash: B256,
+        /// The computed pre-state root
+        got: B256,
+        /// The expected pre-state root from the previous block
+        expected: B256,
+    },
+
     /// Custom error.
     #[error("{0}")]
     Custom(&'static str),

crates/trie/linked/Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+name = "reth-trie-linked"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+homepage.workspace = true
+repository.workspace = true
+exclude.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+alloy-primitives.workspace = true
+alloy-rlp.workspace = true
+alloy-trie.workspace = true
+reth-stateless.workspace = true
+reth-storage-errors.workspace = true
+reth-trie-common.workspace = true
+revm-bytecode.workspace = true
+serde = { workspace = true, optional = true }
+thiserror.workspace = true
+itertools.workspace = true
+
+[features]
+serde = ["dep:serde", "serde/derive", "alloy-primitives/serde"]

crates/trie/linked/src/execution_witness.rs

@@ -0,0 +1,130 @@
+//! This is copied and modified from https://github.com/succinctlabs/rsp
+//! crates/mpt/src/execution_witness.rs rev@2a99f35a9b81452eb53af3848e50addfd481363c
+//! Under MIT license
+
+use crate::mpt::{resolve_nodes, MptNode, MptNodeData, MptNodeReference};
+use alloy_primitives::{
+    keccak256,
+    map::{B256Map, HashMap},
+    Bytes, B256,
+};
+use alloy_rlp::Decodable;
+use alloy_trie::TrieAccount;
+use reth_stateless::validation::StatelessValidationError;
+
+// Builds tries from the witness state.
+//
+// NOTE: This method should be called outside zkVM! In general, you construct tries, then
+// validate them inside zkVM.
+pub(crate) fn build_validated_tries<'a, I>(
+    pre_state_root: B256,
+    states: I,
+) -> Result<(MptNode, B256Map<MptNode>), StatelessValidationError>
+where
+    I: IntoIterator<Item = &'a Bytes>,
+{
+    let err_fn = |_e| StatelessValidationError::WitnessRevealFailed { pre_state_root };
+
+    // Step 1: Decode all RLP-encoded trie nodes and index by hash
+    // IMPORTANT: Witness state contains both *state trie* nodes and *storage tries* nodes!
+    let mut node_map = HashMap::<MptNodeReference, MptNode>::default();
+    let mut node_by_hash = B256Map::<MptNode>::default();
+    let mut root_node: Option<MptNode> = None;
+
+    for encoded in states.into_iter() {
+        let node = MptNode::decode(&mut encoded.as_ref()).map_err(err_fn)?;
+        let hash = keccak256(encoded);
+        if hash == pre_state_root {
+            root_node = Some(node.clone());
+        }
+        node_by_hash.insert(hash, node.clone());
+        node_map.insert(node.reference(), node);
+    }
+
+    // Step 2: Use root_node or fallback to Digest
+    let root = root_node.unwrap_or_else(|| MptNodeData::Digest(pre_state_root).into());
+
+    // Build state trie.
+    let mut raw_storage_tries = Vec::with_capacity(node_by_hash.len());
+    let state_trie = resolve_nodes(&root, &node_map);
+
+    state_trie.try_for_each_leaves(|key, mut value| {
+        let account = TrieAccount::decode(&mut value).map_err(err_fn)?;
+        let hashed_address = B256::from_slice(key);
+        raw_storage_tries.push((hashed_address, account.storage_root));
+        Ok::<_, StatelessValidationError>(())
+    })?;
+
+    // Step 3: Build storage tries per account efficiently
+    let mut storage_tries =
+        B256Map::<MptNode>::with_capacity_and_hasher(raw_storage_tries.len(), Default::default());
+
+    for (hashed_address, storage_root) in raw_storage_tries {
+        let root_node = match node_by_hash.get(&storage_root).cloned() {
+            Some(node) => node,
+            None => {
+                // An execution witness can include an account leaf (with non-empty storageRoot),
+                // but omit its entire storage trie when that account's storage was
+                // NOT touched during the block.
+                continue;
+            }
+        };
+        let storage_trie = resolve_nodes(&root_node, &node_map);
+
+        if storage_trie.is_digest() {
+            return Err(StatelessValidationError::WitnessRevealFailed { pre_state_root });
+        }
+
+        // Insert resolved storage trie.
+        storage_tries.insert(hashed_address, storage_trie);
+    }
+
+    // Step 3a: Verify that state_trie was built correctly - confirm tree hash with pre-state root.
+    validate_state_trie(&state_trie, pre_state_root)?;
+
+    // Step 3b: Verify that each storage trie matches the declared storage_root in the state trie.
+    validate_storage_tries(pre_state_root, &state_trie, &storage_tries)?;
+
+    Ok((state_trie, storage_tries))
+}
+
+// Validate that state_trie was built correctly - confirm tree hash with prev state root.
+fn validate_state_trie(
+    state_trie: &MptNode,
+    pre_state_root: B256,
+) -> Result<(), StatelessValidationError> {
+    if state_trie.hash() != pre_state_root {
+        return Err(StatelessValidationError::PreStateRootMismatch {
+            got: state_trie.hash(),
+            expected: pre_state_root,
+        });
+    }
+    Ok(())
+}
+
+// Validates that each storage trie matches the declared storage_root in the state trie.
+fn validate_storage_tries(
+    pre_state_root: B256,
+    state_trie: &MptNode,
+    storage_tries: &B256Map<MptNode>,
+) -> Result<(), StatelessValidationError> {
+    for (address_hash, storage_trie) in storage_tries.iter() {
+        let account = state_trie
+            .get_rlp::<TrieAccount>(address_hash.as_slice())
+            .map_err(|_e| StatelessValidationError::WitnessRevealFailed { pre_state_root })?
+            .ok_or(StatelessValidationError::WitnessRevealFailed { pre_state_root })?;
+
+        let expected = account.storage_root;
+        let got = storage_trie.hash();
+
+        if expected != got {
+            return Err(StatelessValidationError::PreStorageRootMismatch {
+                address_hash: *address_hash,
+                expected,
+                got,
+            });
+        }
+    }
+
+    Ok(())
+}

crates/trie/linked/src/lib.rs

@@ -0,0 +1,170 @@
+//! Linked Partial Merkle Patricia Trie optimized for stateless execution in zkvm
+
+use crate::mpt::MptNode;
+use alloy_primitives::{keccak256, map::B256Map, Address, Bytes, B256, U256};
+use alloy_trie::{TrieAccount, EMPTY_ROOT_HASH};
+use itertools::Itertools;
+use reth_stateless::{validation::StatelessValidationError, ExecutionWitness, StatelessTrie};
+use reth_storage_errors::ProviderError;
+use reth_trie_common::HashedPostState;
+use revm_bytecode::Bytecode;
+
+mod execution_witness;
+mod mpt;
+
+/// A partial trie that can be updated
+#[derive(Clone, Debug, Eq, PartialEq)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct LinkedSparseTrie {
+    /// The state trie
+    pub state_trie: MptNode,
+    /// The storage tries, keyed by the keccak256 hash of the account address
+    pub storage_tries: B256Map<MptNode>,
+}
+
+impl LinkedSparseTrie {
+    /// Create a partial state trie from a previous state root and a list of RLP-encoded MPT nodes
+    pub fn new<'a, I>(
+        prev_state_root: B256,
+        states: I,
+    ) -> Result<LinkedSparseTrie, StatelessValidationError>
+    where
+        I: IntoIterator<Item = &'a Bytes>,
+    {
+        let (state_trie, storage_tries) =
+            execution_witness::build_validated_tries(prev_state_root, states)?;
+
+        Ok(LinkedSparseTrie { state_trie, storage_tries })
+    }
+
+    /// Get account by address
+    pub fn account(&self, address: Address) -> Result<Option<TrieAccount>, ProviderError> {
+        let hashed_address = keccak256(address);
+        let account = self.state_trie.get_rlp::<TrieAccount>(&*hashed_address)?;
+
+        Ok(account)
+    }
+
+    /// Get storage value of an account at a specific slot.
+    pub fn storage(&self, address: Address, slot: U256) -> Result<U256, ProviderError> {
+        let hashed_address = keccak256(address);
+
+        // Usual case, where given storage slot is present.
+        if let Some(storage_trie) = self.storage_tries.get(&hashed_address) {
+            let key = keccak256(slot.to_be_bytes::<32>());
+            let value = storage_trie.get_rlp::<U256>(&*key)?.unwrap_or_default();
+            return Ok(value);
+        }
+
+        // Storage slot value is not present in the trie, validate that the witness is complete.
+        // FIXME: do we need to validate that the storage trie is empty? seems get_rlp has already
+        // proven that.
+        let account = self.state_trie.get_rlp::<TrieAccount>(&*hashed_address)?;
+        match account {
+            Some(account) => {
+                if account.storage_root != EMPTY_ROOT_HASH {
+                    // if an account has a non-empty storage root, then during the building of the
+                    // trie, we already failed if the storage trie was not present.
+                    unreachable!("pre-built storage trie shall be present");
+                }
+            }
+            None => {
+                todo!("Validate that account witness is valid");
+            }
+        }
+
+        // Account doesn't exist or has empty storage root.
+        Ok(U256::ZERO)
+    }
+
+    /// Mutates state based on diffs provided in [`HashedPostState`].
+    pub fn calculate_state_root(
+        &mut self,
+        state: HashedPostState,
+    ) -> Result<B256, StatelessValidationError> {
+        let err_fn = |_e| StatelessValidationError::StatelessStateRootCalculationFailed;
+        // 1. Apply storage‑slot updates and compute each contract’s storage root
+        //
+        //
+        // We walk over every (address, storage) pair in deterministic order
+        // and update the corresponding per‑account storage trie in‑place.
+        for (address, storage) in
+            state.storages.into_iter().sorted_unstable_by_key(|(addr, _)| *addr)
+        {
+            let storage_trie = self.storage_tries.entry(address).or_default();
+            if storage.wiped {
+                storage_trie.clear();
+            }
+
+            // Apply slot‑level changes
+            for (hashed_slot, value) in
+                storage.storage.into_iter().sorted_unstable_by_key(|(slot, _)| *slot)
+            {
+                if value.is_zero() {
+                    storage_trie.delete(&*hashed_slot).map_err(err_fn)?;
+                } else {
+                    storage_trie.insert_rlp(&*hashed_slot, value).map_err(err_fn)?;
+                }
+            }
+        }
+
+        for (hashed_address, account) in
+            state.accounts.into_iter().sorted_unstable_by_key(|(addr, _)| *addr)
+        {
+            let Some(account) = account else {
+                self.state_trie.delete(&*hashed_address).map_err(err_fn)?;
+                continue;
+            };
+
+            // Determine which storage root should be used for this account
+            let storage_root = if let Some(storage_trie) = self.storage_tries.get(&hashed_address) {
+                storage_trie.hash()
+            } else if let Some(value) =
+                self.state_trie.get_rlp::<TrieAccount>(&*hashed_address).map_err(err_fn)?
+            {
+                value.storage_root
+            } else {
+                EMPTY_ROOT_HASH
+            };
+
+            let account = account.into_trie_account(storage_root);
+            self.state_trie.insert_rlp(&*hashed_address, account).map_err(err_fn)?;
+        }
+
+        Ok(self.state_trie.hash())
+    }
+}
+
+impl StatelessTrie for LinkedSparseTrie {
+    fn new(
+        witness: &ExecutionWitness,
+        pre_state_root: B256,
+    ) -> Result<(Self, B256Map<Bytecode>), StatelessValidationError>
+    where
+        Self: Sized,
+    {
+        let mut bytecode = B256Map::default();
+        for code in witness.codes.iter() {
+            let hash = keccak256(code);
+            bytecode.insert(hash, Bytecode::new_raw(code.clone()));
+        }
+
+        let trie = Self::new(pre_state_root, &witness.state)?;
+        Ok((trie, bytecode))
+    }
+
+    fn account(&self, address: Address) -> Result<Option<TrieAccount>, ProviderError> {
+        LinkedSparseTrie::account(self, address)
+    }
+
+    fn storage(&self, address: Address, slot: U256) -> Result<U256, ProviderError> {
+        LinkedSparseTrie::storage(self, address, slot)
+    }
+
+    fn calculate_state_root(
+        &mut self,
+        state: HashedPostState,
+    ) -> Result<B256, StatelessValidationError> {
+        LinkedSparseTrie::calculate_state_root(self, state)
+    }
+}