diff --git a/Cargo.toml b/Cargo.toml index 89887e13..3bbd1f33 100644 --- a/Cargo.toml +++ b/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" } diff --git a/crates/rpc-proxy/Cargo.toml b/crates/rpc-proxy/Cargo.toml index 6499de56..a9f2903f 100644 --- a/crates/rpc-proxy/Cargo.toml +++ b/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 } diff --git a/crates/rpc-proxy/build.rs b/crates/rpc-proxy/build.rs new file mode 100644 index 00000000..5a37f8ba --- /dev/null +++ b/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>> = + 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"); + } +} diff --git a/crates/rpc-proxy/src/trie.rs b/crates/rpc-proxy/src/trie.rs index 18771e74..e1ab75cd 100644 --- a/crates/rpc-proxy/src/trie.rs +++ b/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( 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 { + 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)) + } +}