feat: Reduce number of debug_storageRangeAt calls by precomputing some preimages

Cargo.toml

@@ -7,6 +7,9 @@ version = "0.3.0"
 edition = "2024"
 rust-version = "1.86"
 
+[profile.dev.build-override]
+opt-level = 3
+
 [workspace.dependencies]
 # internal crates
 guests = { path = "guests" }

crates/rpc-proxy/Cargo.toml

@@ -4,6 +4,10 @@ version = { workspace = true }
 edition = { workspace = true }
 rust-version = { workspace = true }
 
+[build-dependencies]
+alloy-primitives = { workspace = true }
+alloy-trie = { workspace = true }
+
 [dependencies]
 actix-web = "4"
 alloy = { workspace = true }

crates/rpc-proxy/build.rs

@@ -0,0 +1,74 @@
+// Copyright 2025 RISC Zero, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use alloy_primitives::{B256, U256, keccak256};
+use alloy_trie::Nibbles;
+use std::{
+    env,
+    fs::File,
+    io::Write,
+    path::Path,
+    sync::{
+        Arc, OnceLock,
+        atomic::{AtomicUsize, Ordering},
+    },
+    thread,
+};
+
+/// The number of nibbles to use for the prefix.
+const NIBBLES: usize = 5;
+/// The total number of unique prefixes for the given number of nibbles (16^N).
+const PREFIX_COUNT: usize = 16usize.pow(NIBBLES as u32);
+
+fn main() {
+    let out_dir = env::var("OUT_DIR").unwrap();
+    let dest_path = Path::new(&out_dir).join("preimages.bin");
+    println!("cargo:rerun-if-changed=build.rs");
+
+    let table: Arc<Vec<OnceLock<u64>>> =
+        Arc::new((0..PREFIX_COUNT).map(|_| OnceLock::new()).collect());
+    let found = Arc::new(AtomicUsize::new(0));
+
+    thread::scope(|s| {
+        let threads = thread::available_parallelism().unwrap().get();
+        for tid in 0..threads {
+            let table = table.clone();
+            let found = found.clone();
+
+            s.spawn(move || {
+                let mut nonce = tid as u64;
+                while found.load(Ordering::Relaxed) < PREFIX_COUNT {
+                    let hash = keccak256(B256::from(U256::from(nonce)));
+
+                    let nibbles = Nibbles::unpack(hash);
+                    // Calculate the little-endian index from the first N nibbles of the hash.
+                    let idx = (0..NIBBLES)
+                        .map(|i| nibbles.get(i).unwrap() as usize)
+                        .rfold(0, |a, n| (a << 4) | n);
+
+                    if table[idx].set(nonce).is_ok() {
+                        found.fetch_add(1, Ordering::Relaxed);
+                    }
+                    nonce += threads as u64;
+                }
+            });
+        }
+    });
+
+    let mut file = File::create(&dest_path).expect("Could not create file");
+    for cell in table.iter() {
+        let nonce_bytes = cell.get().unwrap().to_le_bytes();
+        file.write_all(&nonce_bytes).expect("Failed to write to file");
+    }
+}

crates/rpc-proxy/src/trie.rs

@@ -18,12 +18,15 @@ use alloy::{
     primitives::{Address, B256, keccak256, map::B256Set},
     providers::Provider,
 };
+use alloy_primitives::U256;
 use anyhow::{Context, Result, bail};
 use revm::database::StorageWithOriginalValues;
 use risc0_ethereum_trie::{Nibbles, Trie, orphan};
 use std::collections::HashSet;
 use tracing::{debug, trace};
 
+static LOOKUP: PreimageLookup = PreimageLookup::new();
+
 pub(crate) async fn handle_removed_account<P, N>(
     provider: &P,
     block_hash: B256,
@@ -117,11 +120,15 @@ where
         return Ok(());
     }
 
-    debug!(%address, "Using debug_storageRangeAt to find preimages for orphan nodes");
-
     let mut missing_storage_keys = B256Set::default();
     for prefix in unresolvable {
-        let storage_key = provider.get_next_storage_key(block_hash, address, prefix).await?;
+        let storage_key = match LOOKUP.find(&prefix) {
+            Some(preimage) => B256::from(preimage),
+            None => {
+                debug!(%address, ?prefix, "Using debug_storageRangeAt to find preimage");
+                provider.get_next_storage_key(block_hash, address, prefix).await?
+            }
+        };
         missing_storage_keys.insert(storage_key);
     }
 
@@ -138,3 +145,44 @@ where
 
     Ok(())
 }
+
+/// A zero-cost wrapper for a precomputed table of Keccak256 pre-images.
+///
+/// This struct holds a static reference to the binary data generated by the `build.rs`
+/// script. It provides a fast method to find a `U256` pre-image for a given hash prefix.
+pub struct PreimageLookup(&'static [u8]);
+
+impl PreimageLookup {
+    /// Creates a new lookup table instance at compile time.
+    pub const fn new() -> Self {
+        Self(include_bytes!(concat!(env!("OUT_DIR"), "/preimages.bin")))
+    }
+
+    /// Returns the number of nibbles the table was precomputed for.
+    pub const fn precomputed_nibbles(&self) -> usize {
+        // The table size in bytes is `16^N * 8`.
+        self.0.len().ilog2().saturating_sub(3) as usize / 4
+    }
+
+    /// Finds a pre-image for a given nibble prefix.
+    ///
+    /// This function looks up a pre-image from the precomputed table. It correctly handles
+    /// input prefixes that are shorter than the table's precomputed depth by finding
+    /// the first corresponding entry.
+    pub fn find(&self, nibbles: &Nibbles) -> Option<U256> {
+        if nibbles.len() > self.precomputed_nibbles() {
+            return None;
+        }
+
+        // Calculate the little-endian index from the input nibbles.
+        // E.g., for [A, B, C], the index will be 0x...CBA.
+        let idx = nibbles.to_vec().iter().rfold(0, |a, n| (a << 4) | *n as usize);
+
+        // Read the 8-byte nonce from the table at the calculated index.
+        let start = idx * 8;
+        let nonce_bytes: [u8; 8] = self.0[start..start + 8].try_into().unwrap();
+        let nonce = u64::from_le_bytes(nonce_bytes);
+
+        Some(U256::from(nonce))
+    }
+}