Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: trie db #71

Merged
merged 3 commits into from
Jan 21, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -23,3 +23,5 @@ latest_proof.json
elf

**.csv

input/
125 changes: 75 additions & 50 deletions crates/executor/client/src/io.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
use std::{collections::HashMap, iter::once};

use itertools::Itertools;
use reth_errors::ProviderError;
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 rsp_witness_db::WitnessDb;
use revm::DatabaseRef;
use serde::{Deserialize, Serialize};

use crate::error::ClientError;
@@ -38,7 +40,7 @@ impl ClientExecutorInput {
}

/// Creates a [`WitnessDb`].
pub fn witness_db(&self) -> Result<WitnessDb, ClientError> {
pub fn witness_db(&self) -> Result<TrieDB<'_>, ClientError> {
<Self as WitnessInput>::witness_db(self)
}
}
@@ -70,6 +72,75 @@ impl WitnessInput for ClientExecutorInput {
}
}

#[derive(Debug)]
pub struct TrieDB<'a> {
inner: &'a EthereumState,
block_hashes: HashMap<u64, B256>,
bytecode_by_hash: HashMap<B256, &'a Bytecode>,
}

impl<'a> TrieDB<'a> {
pub fn new(
inner: &'a EthereumState,
block_hashes: HashMap<u64, B256>,
bytecode_by_hash: HashMap<B256, &'a Bytecode>,
) -> Self {
Self { inner, block_hashes, bytecode_by_hash }
}
}

impl<'a> DatabaseRef for TrieDB<'a> {
/// The database error type.
type Error = ProviderError;

/// Get basic account information.
fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
let hashed_address = keccak256(address);
let hashed_address = hashed_address.as_slice();

let account_in_trie = self.inner.state_trie.get_rlp::<TrieAccount>(hashed_address).unwrap();

let account = account_in_trie.map(|account_in_trie| AccountInfo {
balance: account_in_trie.balance,
nonce: account_in_trie.nonce,
code_hash: account_in_trie.code_hash,
code: None,
});

Ok(account)
}

/// Get account code by its hash.
fn code_by_hash_ref(&self, hash: B256) -> Result<Bytecode, Self::Error> {
Ok(self.bytecode_by_hash.get(&hash).map(|code| (*code).clone()).unwrap())
}

/// Get storage value of address at index.
fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
let hashed_address = keccak256(address);
let hashed_address = hashed_address.as_slice();

let storage_trie = self
.inner
.storage_tries
.get(hashed_address)
.expect("A storage trie must be provided for each account");

Ok(storage_trie
.get_rlp::<U256>(keccak256(index.to_be_bytes::<32>()).as_slice())
.expect("Can get from MPT")
.unwrap_or_default())
}

/// Get block hash by block number.
fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
Ok(*self
.block_hashes
.get(&number)
.expect("A block hash must be provided for each block number"))
}
}

/// A trait for constructing [`WitnessDb`].
pub trait WitnessInput {
/// Gets a reference to the state from which account info and storage slots are loaded.
@@ -99,7 +170,7 @@ pub trait WitnessInput {
/// implementing this trait causes a zkVM run to cost over 5M cycles more. To avoid this, define
/// a method inside the type that calls this trait method instead.
#[inline(always)]
fn witness_db(&self) -> Result<WitnessDb, ClientError> {
fn witness_db(&self) -> Result<TrieDB<'_>, ClientError> {
let state = self.state();

if self.state_anchor() != state.state_root() {
@@ -109,52 +180,6 @@ pub trait WitnessInput {
let bytecodes_by_hash =
self.bytecodes().map(|code| (code.hash_slow(), code)).collect::<HashMap<_, _>>();

let mut accounts = HashMap::new();
let mut storage = HashMap::new();
for (&address, slots) in self.state_requests() {
let hashed_address = keccak256(address);
let hashed_address = hashed_address.as_slice();

let account_in_trie = state.state_trie.get_rlp::<TrieAccount>(hashed_address)?;

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(ClientError::MissingBytecode(address))?)
// Cloning here is fine as `Bytes` is cheap to clone.
.to_owned(),
),
},
None => Default::default(),
},
);

if !slots.is_empty() {
let mut address_storage = HashMap::new();

let storage_trie = state
.storage_tries
.get(hashed_address)
.ok_or(ClientError::MissingTrie(address))?;

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);
}

storage.insert(address, address_storage);
}
}

// Verify and build block hashes
let mut block_hashes: HashMap<u64, B256> = HashMap::new();
for (child_header, parent_header) in self.headers().tuple_windows() {
@@ -176,6 +201,6 @@ pub trait WitnessInput {
block_hashes.insert(parent_header.number, child_header.parent_hash);
}

Ok(WitnessDb { accounts, storage, block_hashes })
Ok(TrieDB::new(state, block_hashes, bytecodes_by_hash))
}
}
10 changes: 6 additions & 4 deletions crates/executor/client/src/lib.rs
Original file line number Diff line number Diff line change
@@ -21,7 +21,7 @@ use reth_evm_optimism::OpExecutorProvider;
use reth_execution_types::ExecutionOutcome;
use reth_optimism_consensus::validate_block_post_execution as validate_block_post_execution_optimism;
use reth_primitives::{proofs, Block, BlockWithSenders, Bloom, Header, Receipt, Receipts, Request};
use revm::{db::CacheDB, Database};
use revm::{db::WrapDatabaseRef, Database};
use revm_primitives::{address, U256};

/// Chain ID for Ethereum Mainnet.
@@ -112,8 +112,10 @@ impl ClientExecutor {
V: Variant,
{
// Initialize the witnessed database with verified storage proofs.
let witness_db = input.witness_db()?;
let cache_db = CacheDB::new(&witness_db);
let wrap_ref = profile!("initialize witness db", {
let trie_db = input.witness_db().unwrap();
WrapDatabaseRef(trie_db)
});

// Execute the block.
let spec = V::spec();
@@ -126,7 +128,7 @@ impl ClientExecutor {
})?;
let executor_difficulty = input.current_block.header.difficulty;
let executor_output = profile!("execute", {
V::execute(&executor_block_input, executor_difficulty, cache_db)
V::execute(&executor_block_input, executor_difficulty, wrap_ref)
})?;

// Validate the block post execution.