direct trie db strategy

crates/core-ethereum/src/lib.rs

@@ -34,20 +34,30 @@ use std::fmt::Display;
 use std::mem::take;
 use std::sync::Arc;
 use zeth_core::db::memory::MemoryDB;
+use zeth_core::db::trie::TrieDB;
 use zeth_core::driver::CoreDriver;
 use zeth_core::stateless::client::StatelessClient;
 use zeth_core::stateless::execute::ExecutionStrategy;
-use zeth_core::stateless::finalize::RethFinalizationStrategy;
-use zeth_core::stateless::initialize::MemoryDbStrategy;
+use zeth_core::stateless::finalize::{MemoryDbFinalizationStrategy, TrieDbFinalizationStrategy};
+use zeth_core::stateless::initialize::{
+    MemoryDbInitializationStrategy, TrieDbInitializationStrategy,
+};
 use zeth_core::stateless::validate::ValidationStrategy;
 
 pub struct RethStatelessClient;
 
 impl StatelessClient<RethCoreDriver, MemoryDB> for RethStatelessClient {
-    type Initialization = MemoryDbStrategy;
+    type Initialization = MemoryDbInitializationStrategy;
+    type Validation = RethValidationStrategy;
+    type Execution = RethExecutionStrategy;
+    type Finalization = MemoryDbFinalizationStrategy;
+}
+
+impl StatelessClient<RethCoreDriver, TrieDB> for RethStatelessClient {
+    type Initialization = TrieDbInitializationStrategy;
     type Validation = RethValidationStrategy;
     type Execution = RethExecutionStrategy;
-    type Finalization = RethFinalizationStrategy;
+    type Finalization = TrieDbFinalizationStrategy;
 }
 
 pub struct RethValidationStrategy;

crates/core-optimism/src/lib.rs

@@ -35,20 +35,30 @@ use std::fmt::Display;
 use std::mem::take;
 use std::sync::Arc;
 use zeth_core::db::memory::MemoryDB;
+use zeth_core::db::trie::TrieDB;
 use zeth_core::driver::CoreDriver;
 use zeth_core::stateless::client::StatelessClient;
 use zeth_core::stateless::execute::ExecutionStrategy;
-use zeth_core::stateless::finalize::RethFinalizationStrategy;
-use zeth_core::stateless::initialize::MemoryDbStrategy;
+use zeth_core::stateless::finalize::{MemoryDbFinalizationStrategy, TrieDbFinalizationStrategy};
+use zeth_core::stateless::initialize::{
+    MemoryDbInitializationStrategy, TrieDbInitializationStrategy,
+};
 use zeth_core::stateless::validate::ValidationStrategy;
 
 pub struct OpRethStatelessClient;
 
 impl StatelessClient<OpRethCoreDriver, MemoryDB> for OpRethStatelessClient {
-    type Initialization = MemoryDbStrategy;
+    type Initialization = MemoryDbInitializationStrategy;
+    type Validation = OpRethValidationStrategy;
+    type Execution = OpRethExecutionStrategy;
+    type Finalization = MemoryDbFinalizationStrategy;
+}
+
+impl StatelessClient<OpRethCoreDriver, TrieDB> for OpRethStatelessClient {
+    type Initialization = TrieDbInitializationStrategy;
     type Validation = OpRethValidationStrategy;
     type Execution = OpRethExecutionStrategy;
-    type Finalization = RethFinalizationStrategy;
+    type Finalization = TrieDbFinalizationStrategy;
 }
 
 pub struct OpRethValidationStrategy;

crates/core/src/db/mod.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 pub mod memory;
+pub mod trie;
 pub mod unreachable;
 pub mod update;
 pub mod wrapper;

crates/core/src/db/trie.rs

@@ -0,0 +1,93 @@
+// Copyright 2024 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 crate::keccak::keccak;
+use crate::mpt::MptNode;
+use crate::rescue::Recoverable;
+use crate::stateless::data::StorageEntry;
+use alloy_consensus::Account;
+use alloy_primitives::map::HashMap;
+use alloy_primitives::{Address, B256, U256};
+use reth_primitives::revm_primitives::db::Database;
+use reth_primitives::revm_primitives::{AccountInfo, Bytecode};
+use reth_revm::DatabaseRef;
+use reth_storage_errors::provider::ProviderError;
+
+#[derive(Default)]
+pub struct TrieDB {
+    pub accounts: MptNode,
+    pub storage: HashMap<Address, StorageEntry>,
+    pub contracts: HashMap<B256, Bytecode>,
+    pub block_hashes: HashMap<u64, B256>,
+}
+
+impl Recoverable for TrieDB {
+    fn rescue(&mut self) -> Option<Self> {
+        Some(core::mem::take(self))
+    }
+}
+
+impl DatabaseRef for TrieDB {
+    type Error = ProviderError;
+
+    fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
+        Ok(self
+            .accounts
+            .get_rlp::<Account>(&keccak(address))
+            .unwrap()
+            .map(|acc| AccountInfo {
+                balance: acc.balance,
+                nonce: acc.nonce,
+                code_hash: acc.code_hash,
+                code: None,
+            }))
+    }
+
+    fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
+        Ok(self.contracts.get(&code_hash).unwrap().clone())
+    }
+
+    fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
+        let entry = self.storage.get(&address).unwrap();
+        Ok(entry
+            .storage_trie
+            .get_rlp(&keccak(index.to_be_bytes::<32>()))
+            .unwrap()
+            .unwrap_or_default())
+    }
+
+    fn block_hash_ref(&self, number: u64) -> Result<B256, Self::Error> {
+        Ok(*self.block_hashes.get(&number).unwrap())
+    }
+}
+
+impl Database for TrieDB {
+    type Error = ProviderError;
+
+    fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
+        self.basic_ref(address)
+    }
+
+    fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
+        self.code_by_hash_ref(code_hash)
+    }
+
+    fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
+        self.storage_ref(address, index)
+    }
+
+    fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
+        self.block_hash_ref(number)
+    }
+}

crates/core/src/stateless/finalize.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::db::memory::MemoryDB;
+use crate::db::trie::TrieDB;
 use crate::db::update::{into_plain_state, Update};
 use crate::driver::CoreDriver;
 use crate::keccak::keccak;
@@ -36,17 +38,60 @@ pub trait FinalizationStrategy<Driver: CoreDriver, Database> {
     ) -> anyhow::Result<()>;
 }
 
-pub struct RethFinalizationStrategy;
+pub struct TrieDbFinalizationStrategy;
 
-impl<Driver: CoreDriver, Database: Update> FinalizationStrategy<Driver, Database>
-    for RethFinalizationStrategy
-{
+impl<Driver: CoreDriver> FinalizationStrategy<Driver, TrieDB> for TrieDbFinalizationStrategy {
+    fn finalize_state(
+        block: &mut Driver::Block,
+        _state_trie: &mut MptNode,
+        _storage_tries: &mut HashMap<Address, StorageEntry>,
+        parent_header: &mut Driver::Header,
+        db: Option<&mut TrieDB>,
+        bundle_state: BundleState,
+        with_further_updates: bool,
+    ) -> anyhow::Result<()> {
+        let TrieDB {
+            accounts: state_trie,
+            storage: storage_tries,
+            block_hashes,
+            ..
+        } = db.expect("Missing TrieDB instance");
+
+        // Update the trie data
+        <MemoryDbFinalizationStrategy as FinalizationStrategy<Driver, MemoryDB>>::finalize_state(
+            block,
+            state_trie,
+            storage_tries,
+            parent_header,
+            None,
+            bundle_state,
+            false,
+        )?;
+
+        // Get the header
+        let header = Driver::block_header(block);
+
+        // Give back the tries
+        if !with_further_updates {
+            core::mem::swap(state_trie, _state_trie);
+            core::mem::swap(storage_tries, _storage_tries);
+        } else {
+            block_hashes.insert(Driver::block_number(header), Driver::header_hash(header));
+        }
+
+        Ok(())
+    }
+}
+
+pub struct MemoryDbFinalizationStrategy;
+
+impl<Driver: CoreDriver> FinalizationStrategy<Driver, MemoryDB> for MemoryDbFinalizationStrategy {
     fn finalize_state(
         block: &mut Driver::Block,
         state_trie: &mut MptNode,
         storage_tries: &mut HashMap<Address, StorageEntry>,
         parent_header: &mut Driver::Header,
-        db: Option<&mut Database>,
+        db: Option<&mut MemoryDB>,
         bundle_state: BundleState,
         with_further_updates: bool,
     ) -> anyhow::Result<()> {
@@ -58,6 +103,7 @@ impl<Driver: CoreDriver, Database: Update> FinalizationStrategy<Driver, Database
                 state_trie.hash()
             );
         }
+
         // Convert the state update bundle
         let state_changeset = into_plain_state(bundle_state);
         // Update the trie data

crates/core/src/stateless/initialize.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use crate::db::memory::MemoryDB;
+use crate::db::trie::TrieDB;
 use crate::driver::CoreDriver;
 use crate::keccak::keccak;
 use crate::mpt::MptNode;
@@ -38,9 +39,75 @@ pub trait InitializationStrategy<Driver: CoreDriver, Database> {
     ) -> anyhow::Result<Database>;
 }
 
-pub struct MemoryDbStrategy;
+pub struct TrieDbInitializationStrategy;
 
-impl<Driver: CoreDriver> InitializationStrategy<Driver, MemoryDB> for MemoryDbStrategy {
+impl<Driver: CoreDriver> InitializationStrategy<Driver, TrieDB> for TrieDbInitializationStrategy {
+    fn initialize_database(
+        state_trie: &mut MptNode,
+        storage_tries: &mut HashMap<Address, StorageEntry>,
+        contracts: &mut Vec<Bytes>,
+        parent_header: &mut Driver::Header,
+        ancestor_headers: &mut Vec<Driver::Header>,
+    ) -> anyhow::Result<TrieDB> {
+        // Verify starting state trie root
+        if Driver::state_root(parent_header) != state_trie.hash() {
+            bail!(
+                "Invalid initial state trie: expected {}, got {}",
+                Driver::state_root(parent_header),
+                state_trie.hash()
+            );
+        }
+
+        // hash all the contract code
+        let contracts = take(contracts)
+            .into_iter()
+            .map(|bytes| (keccak(&bytes).into(), Bytecode::new_raw(bytes)))
+            .collect();
+
+        // prepare block hash history
+        let mut block_hashes: HashMap<u64, B256> =
+            HashMap::with_capacity_and_hasher(ancestor_headers.len() + 1, Default::default());
+        block_hashes.insert(
+            Driver::block_number(parent_header),
+            Driver::header_hash(parent_header),
+        );
+        let mut prev = &*parent_header;
+        for current in ancestor_headers.iter() {
+            let current_hash = Driver::header_hash(current);
+            if Driver::parent_hash(prev) != current_hash {
+                bail!(
+                    "Invalid chain: {} is not the parent of {}",
+                    Driver::block_number(current),
+                    Driver::block_number(prev)
+                );
+            }
+            if Driver::block_number(parent_header) < Driver::block_number(current)
+                || Driver::block_number(parent_header) - Driver::block_number(current) >= 256
+            {
+                bail!(
+                    "Invalid chain: {} is not one of the {} most recent blocks",
+                    Driver::block_number(current),
+                    256,
+                );
+            }
+            block_hashes.insert(Driver::block_number(current), current_hash);
+            prev = current;
+        }
+
+        Ok(TrieDB {
+            accounts: take(state_trie),
+            storage: take(storage_tries),
+            contracts,
+            block_hashes,
+        })
+    }
+}
+
+pub struct MemoryDbInitializationStrategy;
+
+impl<Driver: CoreDriver> InitializationStrategy<Driver, MemoryDB>
+    for MemoryDbInitializationStrategy
+{
     fn initialize_database(
         state_trie: &mut MptNode,
         storage_tries: &mut HashMap<Address, StorageEntry>,

guests/reth-ethereum/src/main.rs

@@ -14,8 +14,9 @@
 
 // use c_kzg::KzgSettings;
 use risc0_zkvm::guest::env;
+use zeth_core::db::trie::TrieDB;
 use zeth_core::stateless::client::StatelessClient;
-use zeth_core_ethereum::RethStatelessClient;
+use zeth_core_ethereum::{RethCoreDriver, RethStatelessClient};
 // todo: use this instead of the alloy KzgEnv to save cycles
 // lazy_static::lazy_static! {
 //     /// KZG Ceremony data
@@ -36,11 +37,12 @@ fn main() {
     let stateless_client_data_rkyv = env::read_frame();
     let stateless_client_data_pot = env::read_frame();
     env::log("Deserializing input data");
-    let stateless_client_data = RethStatelessClient::data_from_parts(
-        &stateless_client_data_rkyv,
-        &stateless_client_data_pot,
-    )
-    .expect("Failed to load client data from stdin");
+    let stateless_client_data =
+        <RethStatelessClient as StatelessClient<RethCoreDriver, TrieDB>>::data_from_parts(
+            &stateless_client_data_rkyv,
+            &stateless_client_data_pot,
+        )
+        .expect("Failed to load client data from stdin");
     let validation_depth = stateless_client_data.blocks.len() as u64;
     assert!(
         stateless_client_data.chain.is_ethereum(),
@@ -49,8 +51,10 @@ fn main() {
     let chain_id = stateless_client_data.chain as u64;
     // Build the block
     env::log("Validating blocks");
-    let engine =
-        RethStatelessClient::validate(stateless_client_data).expect("block validation failed");
+    let engine = <RethStatelessClient as StatelessClient<RethCoreDriver, TrieDB>>::validate(
+        stateless_client_data,
+    )
+    .expect("block validation failed");
     // Build the journal (todo: make this a strategy)
     let block_hash = engine.data.parent_header.hash_slow();
     let total_difficulty = engine.data.total_difficulty;