risc0 · Wollac · Mar 13, 2025 · Mar 13, 2025 · Mar 14, 2025 · Mar 27, 2025
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -36,7 +36,7 @@ features = ["orphan", "rkyv", "rlp_serialize", "serde"]
 # External
 [workspace.dependencies]
 # Alloy
-alloy = { version = "0.4.2", features = ["full"] }
+alloy = { version = "0.4.2", features = ["full", "serde"] }
 alloy-chains = "0.1.38"
 alloy-consensus = "0.4.2"
 alloy-genesis = "0.4.2"

diff --git a/bin/ethereum/Cargo.toml b/bin/ethereum/Cargo.toml
@@ -23,3 +23,4 @@ disable-dev-mode = ["zeth/disable-dev-mode"]
 cuda = ["zeth/cuda"]
 metal = ["zeth/metal"]
 prove = ["zeth/prove"]
+erigon-range-queries = ["zeth-preflight-ethereum/erigon-range-queries"]
diff --git a/bin/ethereum/data/ethereum/1-12965000.json.gz b/bin/ethereum/data/ethereum/1-12965000.json.gz
diff --git a/bin/optimism/Cargo.toml b/bin/optimism/Cargo.toml
@@ -24,3 +24,4 @@ disable-dev-mode = ["zeth/disable-dev-mode"]
 cuda = ["zeth/cuda"]
 metal = ["zeth/metal"]
 prove = ["zeth/prove"]
+erigon-range-queries = ["zeth-preflight-optimism/erigon-range-queries"]
diff --git a/crates/core/src/mpt.rs b/crates/core/src/mpt.rs
@@ -1,7 +1,6 @@
-use alloy_primitives::map::B256Set;
 use alloy_primitives::B256;
 use alloy_rlp::{Decodable, Encodable};
-use risc0_ethereum_trie::{orphan, CachedTrie};
+use risc0_ethereum_trie::{orphan, CachedTrie, Nibbles};
 use serde::{Deserialize, Serialize};
 use std::borrow::Borrow;
 use std::marker::PhantomData;
@@ -49,29 +48,18 @@ impl<T: Decodable + Encodable> MptNode<T> {
     }
 
     /// Tries to resolve the potential removal orphan corresponding to `key` from the given
-    /// post-removal proof. If the orphan cannot be resolved from the proof alone, the `key`
-    /// corresponding to the unresolved path is added to `unresolvable`.
-    pub fn resolve_orphan<K, N>(
+    /// post-removal proof. If the orphan cannot be resolved from the proof alone, the
+    /// prefix of the missing MPT key is returned.
+    pub fn resolve_orphan<K: AsRef<[u8]>, N: AsRef<[u8]>>(
         &mut self,
         key: K,
         post_state_proof: impl IntoIterator<Item = N>,
-        unresolvable: &mut B256Set,
-    ) -> anyhow::Result<()>
-    where
-        K: AsRef<[u8]>,
-        N: AsRef<[u8]>,
-    {
-        match self.inner.resolve_orphan(key, post_state_proof) {
-            Ok(_) => {}
-            Err(orphan::Error::Unresolvable(prefix)) => {
-                // convert the unresolvable prefix nibbles into a B256 key with zero padding
-                let key = B256::right_padding_from(&prefix.pack());
-                unresolvable.insert(key);
-            }
-            Err(err) => return Err(err.into()),
-        };
-
-        Ok(())
+    ) -> anyhow::Result<Option<Nibbles>> {
+        match self.inner.resolve_orphan(&key, post_state_proof) {
+            Ok(_) => Ok(None),
+            Err(orphan::Error::Unresolvable(prefix)) => Ok(Some(prefix)),
+            Err(err) => Err(err.into()),
+        }
     }
 
     #[inline]

diff --git a/crates/preflight-ethereum/Cargo.toml b/crates/preflight-ethereum/Cargo.toml
@@ -16,3 +16,6 @@ workspace = true
 alloy.workspace = true
 reth-chainspec.workspace = true
 reth-primitives.workspace = true
+
+[features]
+erigon-range-queries = ["zeth-preflight/erigon-range-queries"]
diff --git a/crates/preflight-optimism/Cargo.toml b/crates/preflight-optimism/Cargo.toml
@@ -29,4 +29,7 @@ reth-primitives = { workspace = true, features = ["optimism"] }
 reth-revm = { workspace = true, features = ["optimism"] }
 reth-storage-errors.workspace = true
 
-serde_json.workspace = true
+serde_json.workspace = true
+
+[features]
+erigon-range-queries = ["zeth-preflight/erigon-range-queries"]
diff --git a/crates/preflight/Cargo.toml b/crates/preflight/Cargo.toml
@@ -15,6 +15,7 @@ itertools.workspace = true
 k256.workspace = true
 log.workspace = true
 pot.workspace = true
+risc0-ethereum-trie.workspace = true
 rkyv.workspace = true
 serde.workspace = true
 serde_json.workspace = true
@@ -24,3 +25,6 @@ reth-chainspec.workspace = true
 reth-revm.workspace = true
 reth-primitives.workspace = true
 reth-storage-errors.workspace = true
+
+[features]
+erigon-range-queries = []
diff --git a/crates/preflight/src/client.rs b/crates/preflight/src/client.rs
@@ -17,13 +17,15 @@ use crate::driver::PreflightDriver;
 use crate::provider::db::ProviderDB;
 use crate::provider::query::{BlockQuery, UncleQuery};
 use crate::provider::{new_provider, Provider};
+use alloy::hex;
 use alloy::network::Network;
-use alloy::primitives::map::{AddressHashMap, B256Set, HashSet};
+use alloy::primitives::map::{AddressHashMap, HashSet};
 use alloy::primitives::{keccak256, Bytes, U256};
 use alloy::rpc::types::EIP1186StorageProof;
 use anyhow::Context;
 use itertools::Itertools;
 use log::{debug, info, warn};
+use risc0_ethereum_trie::Nibbles;
 use std::cell::RefCell;
 use std::iter::zip;
 use std::path::PathBuf;
@@ -258,8 +260,10 @@ where
 
             info!("Extending tries from post-state proofs...");
 
-            let mut unresolvable_state_keys = B256Set::default();
+            let mut resolved_state_keys = 0usize;
+            let mut resolved_storage_keys = 0usize;
 
+            let mut unresolved_state_keys: HashSet<Nibbles> = HashSet::default();
             for (address, account_proof) in latest_proofs {
                 let db_key = keccak256(address);
 
@@ -272,15 +276,18 @@ where
                 }
 
                 // otherwise, prepare trie for the removal of that key
-                state_trie
-                    .resolve_orphan(
-                        db_key,
-                        account_proof.account_proof,
-                        &mut unresolvable_state_keys,
-                    )
-                    .with_context(|| format!("failed to resolve orphan for {}", address))?;
+                if let Some(unresolved_prefix) = state_trie
+                    .resolve_orphan(db_key, account_proof.account_proof)
+                    .with_context(|| format!("failed to resolve orphan for {}", address))?
+                {
+                    debug!(
+                        "unresolved key in state trie: {}",
+                        hex::encode(unresolved_prefix.pack())
+                    );
+                    unresolved_state_keys.insert(unresolved_prefix);
+                }
 
-                let mut unresolvable_storage_keys = B256Set::default();
+                let mut unresolved_storage_keys: HashSet<Nibbles> = HashSet::default();
 
                 let storage_trie = &mut storage_tries.get_mut(&address).unwrap().storage_trie;
                 for EIP1186StorageProof { key, proof, .. } in account_proof.storage_proof {
@@ -290,35 +297,52 @@ where
                         storage_trie.hydrate_from_rlp(proof)?;
                     } else {
                         // otherwise, prepare trie for the removal of that key
-                        storage_trie
-                            .resolve_orphan(db_key, proof, &mut unresolvable_storage_keys)
+                        if let Some(unresolved_prefix) = storage_trie
+                            .resolve_orphan(db_key, proof)
                             .with_context(|| {
                                 format!("failed to resolve orphan for {}@{}", key.0, address)
-                            })?;
+                            })?
+                        {
+                            debug!(
+                                "unresolved key in storage trie for {}: {}",
+                                address,
+                                hex::encode(unresolved_prefix.pack())
+                            );
+                            unresolved_storage_keys.insert(unresolved_prefix);
+                        }
                     }
                 }
 
                 // if orphans could not be resolved, use a range query to get that missing info
-                if !unresolvable_storage_keys.is_empty() {
+                if !unresolved_storage_keys.is_empty() {
+                    resolved_storage_keys += unresolved_storage_keys.len();
                     let proof = preflight_db
-                        .get_next_slot_proofs(block_count, address, unresolvable_storage_keys)
-                        .with_context(|| format!("failed to get next slot for {}", address))?;
+                        .get_storage_proofs_with_prefix(address, unresolved_storage_keys)
+                        .with_context(|| {
+                            format!("failed to get proof for unresolved slots for {}", address)
+                        })?;
                     storage_trie
                         .hydrate_from_rlp(proof.storage_proof.iter().flat_map(|p| &p.proof))
                         .with_context(|| format!("invalid storage proof for {}", address))?;
                 }
             }
 
-            for state_key in unresolvable_state_keys {
+            resolved_state_keys += unresolved_state_keys.len();
+            for prefix in unresolved_state_keys {
                 let proof = preflight_db
-                    .get_next_account_proof(block_count, state_key)
-                    .context("failed to get next account")?;
+                    .get_account_proof_with_prefix(prefix)
+                    .context("failed to get proof for unresolved account")?;
                 state_trie
                     .hydrate_from_rlp(proof.account_proof)
                     .with_context(|| format!("invalid account proof for {}", proof.address))?;
             }
 
-            info!("Saving provider cache ...");
+            info!(
+                "Resolved {} addresses and {} storage keys",
+                resolved_state_keys, resolved_storage_keys
+            );
+
+            info!("Saving provider cache...");
             preflight_db.save_provider()?;
 
             // Increment block number counter

diff --git a/crates/preflight/src/db.rs b/crates/preflight/src/db.rs
@@ -15,17 +15,18 @@
 use crate::driver::PreflightDriver;
 use crate::provider::db::ProviderDB;
 use crate::provider::get_proofs;
-use crate::provider::query::{AccountRangeQuery, BlockQuery, ProofQuery, StorageRangeQuery};
+use crate::provider::query::{BlockQuery, NextAccountQuery, NextSlotQuery, ProofQuery};
 use alloy::network::Network;
 use alloy::primitives::map::HashMap;
-use alloy::primitives::{Address, B256, U256};
+use alloy::primitives::{keccak256, Address, B256, U256};
 use alloy::rpc::types::EIP1186AccountProofResponse;
-use anyhow::Context;
-use log::{debug, error};
+use anyhow::{ensure, Context};
+use log::error;
 use reth_primitives::revm_primitives::{Account, AccountInfo, Bytecode};
 use reth_revm::db::states::StateChangeset;
 use reth_revm::db::CacheDB;
 use reth_revm::{Database, DatabaseCommit, DatabaseRef};
+use risc0_ethereum_trie::Nibbles;
 use std::cell::{Ref, RefCell};
 use std::collections::BTreeSet;
 use std::marker::PhantomData;
@@ -252,25 +253,30 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
         Ok(headers)
     }
 
-    /// Fetches the EIP-1186 proof for the next account after a given key.
+    /// Retrieves an EIP-1186 account proof for an account whose hashed address starts with the
+    /// specified prefix.
     ///
-    /// This method retrieves an [EIP1186AccountProofResponse] for the account whose address, when
-    /// hashed, lexicographically follows the provided `start` key. The proof is generated for the
-    /// block `block_count` after the currently configured block in the provider.
-    pub fn get_next_account_proof(
+    /// The method finds the first account whose address when hashed with keccak256 starts with the
+    /// given prefix. Then, it retrieves and returns an [EIP1186AccountProofResponse] of that
+    /// account.
+    pub fn get_account_proof_with_prefix(
         &mut self,
-        block_count: u64,
-        start: B256,
+        prefix: Nibbles,
     ) -> anyhow::Result<EIP1186AccountProofResponse> {
         let initial_db = self.inner.db.db.borrow_mut();
         let provider_db = initial_db.db.borrow_db();
         let mut provider = provider_db.provider.borrow_mut();
-        let block_no = initial_db.db.borrow_db().block_no + block_count - 1;
+        let block_no = initial_db.db.borrow_db().block_no;
+        // convert the prefix nibbles into a B256 key with zero padding
+        let start = B256::right_padding_from(&prefix.pack());
 
-        debug!("getting next account: start={}", start);
         let address = provider
-            .get_next_account(&AccountRangeQuery::new(block_no, start))
-            .context("debug_accountRange call failed")?;
+            .get_next_account(&NextAccountQuery { block_no, start })
+            .context("getting next account failed")?;
+        ensure!(
+            Nibbles::unpack(keccak256(address)).starts_with(&prefix),
+            "invalid provider response: address doesn't match prefix"
+        );
 
         provider
             .get_proof(&ProofQuery {
@@ -281,33 +287,40 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
             .context("eth_getProof call failed")
     }
 
-    /// Fetches EIP-1186 proofs for the next storage slots of a given account.
+    /// Retrieves an EIP-1186 account proof for storage slots whose key hashes match the given
+    /// prefixes.
     ///
-    /// This method retrieves an [EIP1186AccountProofResponse] for multiple storage slots of a given
-    /// account. For each `B256` key provided in the `starts` iterator, the method finds the next
-    /// storage slot whose hashed index lexicographically follows the given key. The proofs are
-    /// generated for the block `block_count` after the currently configured block in the provider.
-    pub fn get_next_slot_proofs(
+    /// For each prefix provided, the method finds the first storage slot whose key, when hashed
+    /// with keccak256, starts with the prefix. Then, it retrieves and returns an
+    /// [EIP1186AccountProofResponse] of all those storage slots in the given account.
+    pub fn get_storage_proofs_with_prefix(
         &mut self,
-        block_count: u64,
         address: Address,
-        starts: impl IntoIterator<Item = B256>,
+        prefixes: impl IntoIterator<Item = Nibbles>,
     ) -> anyhow::Result<EIP1186AccountProofResponse> {
         let initial_db = self.inner.db.db.borrow_mut();
         let provider_db = initial_db.db.borrow_db();
         let mut provider = provider_db.provider.borrow_mut();
-        let block_no = initial_db.db.borrow_db().block_no + block_count - 1;
+        let block_no = initial_db.db.borrow_db().block_no;
 
         let mut indices = BTreeSet::new();
-        for start in starts {
-            debug!(
-                "getting next storage key: address={},start={}",
-                address, start
+        for prefix in prefixes {
+            // convert the prefix nibbles into a B256 key with zero padding
+            let start = B256::right_padding_from(&prefix.pack());
+            let key: B256 = provider
+                .get_next_slot(&NextSlotQuery {
+                    block_no,
+                    address,
+                    start,
+                })
+                .context("getting next storage key failed")?
+                .into();
+            ensure!(
+                Nibbles::unpack(keccak256(key)).starts_with(&prefix),
+                "invalid provider response: storage key doesn't match prefix"
             );
-            let slot = provider
-                .get_next_slot(&StorageRangeQuery::new(block_no, address, start))
-                .context("debug_storageRangeAt call failed")?;
-            indices.insert(B256::from(slot));
+
+            indices.insert(B256::from(key));
         }
 
         provider

diff --git a/crates/preflight/src/provider/cache_provider.rs b/crates/preflight/src/provider/cache_provider.rs
@@ -189,7 +189,7 @@ impl<N: Network> Provider<N> for CachedRpcProvider<N> {
         Ok(out)
     }
 
-    fn get_next_account(&mut self, query: &AccountRangeQuery) -> anyhow::Result<Address> {
+    fn get_next_account(&mut self, query: &NextAccountQuery) -> anyhow::Result<Address> {
         let cache_out = self.cache.get_next_account(query);
         if cache_out.is_ok() {
             return cache_out;
@@ -201,7 +201,7 @@ impl<N: Network> Provider<N> for CachedRpcProvider<N> {
         Ok(out)
     }
 
-    fn get_next_slot(&mut self, query: &StorageRangeQuery) -> anyhow::Result<U256> {
+    fn get_next_slot(&mut self, query: &NextSlotQuery) -> anyhow::Result<U256> {
         let cache_out = self.cache.get_next_slot(query);
         if cache_out.is_ok() {
             return cache_out;