Skip to content

Commit

Permalink
refactor: Database initialization strategy (#20)
Browse files Browse the repository at this point in the history
* optimizations

* header prep strategy

* block build strategy

* db init strategy

* nit
  • Loading branch information
hashcashier authored Aug 22, 2023
1 parent ed8ec90 commit 3cf8f46
Show file tree
Hide file tree
Showing 12 changed files with 263 additions and 243 deletions.
7 changes: 3 additions & 4 deletions guests/eth-block/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,9 @@
use risc0_zkvm::guest::env;
use zeth_lib::{
block_builder::BlockBuilder, consts::ETH_MAINNET_CHAIN_SPEC, execution::EthTxExecStrategy,
mem_db::MemDb, validation::Input,
finalization::BuildFromMemDbStrategy, initialization::MemDbInitStrategy, input::Input,
mem_db::MemDb, preparation::EthHeaderPrepStrategy,
};
use zeth_lib::finalization::BuildFromMemDbStrategy;
use zeth_lib::preparation::EthHeaderPrepStrategy;

risc0_zkvm::guest::entry!(main);

Expand All @@ -29,7 +28,7 @@ pub fn main() {
let input: Input = env::read();
// Build the resulting block
let output = BlockBuilder::<MemDb>::new(&ETH_MAINNET_CHAIN_SPEC, input)
.initialize_db()
.initialize_database::<MemDbInitStrategy>()
.expect("Failed to create in-memory evm storage")
.prepare_header::<EthHeaderPrepStrategy>()
.expect("Failed to create the initial block header fields")
Expand Down
5 changes: 3 additions & 2 deletions host/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,10 @@ use zeth_lib::{
consts::{Network, ETH_MAINNET_CHAIN_SPEC},
execution::EthTxExecStrategy,
finalization::DebugBuildFromMemDbStrategy,
initialization::MemDbInitStrategy,
input::Input,
mem_db::MemDb,
preparation::EthHeaderPrepStrategy,
validation::Input,
};
use zeth_primitives::BlockHash;

Expand Down Expand Up @@ -108,7 +109,7 @@ async fn main() -> Result<()> {
info!("Running from memory ...");

let block_builder = BlockBuilder::<MemDb>::new(&ETH_MAINNET_CHAIN_SPEC, input)
.initialize_db()
.initialize_database::<MemDbInitStrategy>()
.expect("Error initializing MemDb from Input")
.prepare_header::<EthHeaderPrepStrategy>()
.expect("Error creating initial block header")
Expand Down
89 changes: 6 additions & 83 deletions lib/src/block_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,27 +12,15 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use core::mem;

use anyhow::Result;
use hashbrown::{hash_map, HashMap};
use revm::primitives::{Account, AccountInfo, Address, Bytecode, B160, B256, U256};
use zeth_primitives::{
block::Header,
keccak::{keccak, KECCAK_EMPTY},
revm::to_revm_b256,
trie::StateAccount,
Bytes,
};
use revm::primitives::{Account, Address, B160, B256, U256};
use zeth_primitives::block::Header;

use crate::{
consts::ChainSpec,
execution::TxExecStrategy,
finalization::BlockBuildStrategy,
guest_mem_forget,
mem_db::{AccountState, DbAccount},
consts::ChainSpec, execution::TxExecStrategy, finalization::BlockBuildStrategy,
initialization::DbInitStrategy, input::Input, mem_db::DbAccount,
preparation::HeaderPrepStrategy,
validation::{verify_parent_chain, verify_state_trie, verify_storage_trie, Input},
};

pub trait BlockBuilderDatabase: revm::Database + Sized {
Expand Down Expand Up @@ -76,73 +64,8 @@ where
}

/// Initializes the database from the input tries.
pub fn initialize_db(mut self) -> Result<Self> {
verify_state_trie(
&self.input.parent_state_trie,
&self.input.parent_header.state_root,
)?;

// hash all the contract code
let contracts: HashMap<B256, Bytes> = mem::take(&mut self.input.contracts)
.into_iter()
.map(|bytes| (keccak(&bytes).into(), bytes))
.collect();

// Load account data into db
let mut accounts = HashMap::with_capacity(self.input.parent_storage.len());
for (address, (storage_trie, slots)) in &mut self.input.parent_storage {
// consume the slots, as they are no longer needed afterwards
let slots = mem::take(slots);

// load the account from the state trie or empty if it does not exist
let state_account = self
.input
.parent_state_trie
.get_rlp::<StateAccount>(&keccak(address))?
.unwrap_or_default();
verify_storage_trie(address, storage_trie, &state_account.storage_root)?;

// load the corresponding code
let code_hash = to_revm_b256(state_account.code_hash);
let bytecode = if code_hash.0 == KECCAK_EMPTY.0 {
Bytecode::new()
} else {
let bytes = contracts.get(&code_hash).unwrap().clone();
unsafe { Bytecode::new_raw_with_hash(bytes.0, code_hash) }
};

// load storage reads
let mut storage = HashMap::with_capacity(slots.len());
for slot in slots {
let value: zeth_primitives::U256 = storage_trie
.get_rlp(&keccak(slot.to_be_bytes::<32>()))?
.unwrap_or_default();
storage.insert(slot, value);
}

let mem_account = DbAccount {
info: AccountInfo {
balance: state_account.balance,
nonce: state_account.nonce,
code_hash: to_revm_b256(state_account.code_hash),
code: Some(bytecode),
},
state: AccountState::None,
storage,
};

accounts.insert(*address, mem_account);
}
guest_mem_forget(contracts);

// prepare block hash history
let block_hashes =
verify_parent_chain(&self.input.parent_header, &self.input.ancestor_headers)?;

// Store database
self.db = Some(D::load(accounts, block_hashes));

Ok(self)
pub fn initialize_database<T: DbInitStrategy<Db = D>>(self) -> Result<Self> {
T::initialize_database(self)
}

/// Initializes the header. This must be called before executing transactions.
Expand Down
13 changes: 10 additions & 3 deletions lib/src/execution.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,8 @@ use zeth_primitives::{
use crate::{
block_builder::{BlockBuilder, BlockBuilderDatabase},
consts,
consts::GWEI_TO_WEI,
consts::{GWEI_TO_WEI, MIN_SPEC_ID},
guest_mem_forget,
validation::compute_spec_id,
};

pub trait TxExecStrategy {
Expand All @@ -56,7 +55,15 @@ impl TxExecStrategy for EthTxExecStrategy {
.header
.as_mut()
.expect("Header is not initialized");
let spec_id = compute_spec_id(block_builder.chain_spec, header.number)?;
// Compute the spec id
let spec_id = block_builder.chain_spec.spec_id(header.number);
if !SpecId::enabled(spec_id, MIN_SPEC_ID) {
bail!(
"Invalid protocol version: expected >= {:?}, got {:?}",
MIN_SPEC_ID,
spec_id,
)
}

#[cfg(not(target_os = "zkvm"))]
{
Expand Down
2 changes: 1 addition & 1 deletion lib/src/host/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,9 @@ use crate::{
mpt::{orphaned_digests, resolve_digests, shorten_key},
provider::{new_provider, BlockQuery},
},
input::{Input, StorageEntry},
mem_db::MemDb,
preparation::EthHeaderPrepStrategy,
validation::{Input, StorageEntry},
};

pub mod mpt;
Expand Down
165 changes: 165 additions & 0 deletions lib/src/initialization.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,165 @@
// Copyright 2023 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 core::{fmt::Debug, mem};

use anyhow::{bail, Result};
use hashbrown::HashMap;
use revm::{
primitives::{AccountInfo, Bytecode, B256},
Database,
};
use zeth_primitives::{
keccak::{keccak, KECCAK_EMPTY},
revm::to_revm_b256,
trie::StateAccount,
Bytes,
};

use crate::{
block_builder::{BlockBuilder, BlockBuilderDatabase},
consts::MAX_BLOCK_HASH_AGE,
guest_mem_forget,
mem_db::{AccountState, DbAccount, MemDb},
};

pub trait DbInitStrategy {
type Db;

fn initialize_database(block_builder: BlockBuilder<Self::Db>) -> Result<BlockBuilder<Self::Db>>
where
Self::Db: BlockBuilderDatabase,
<Self::Db as Database>::Error: Debug;
}

pub struct MemDbInitStrategy {}

impl DbInitStrategy for MemDbInitStrategy {
type Db = MemDb;

fn initialize_database(
mut block_builder: BlockBuilder<Self::Db>,
) -> Result<BlockBuilder<Self::Db>>
where
Self::Db: BlockBuilderDatabase,
<Self::Db as Database>::Error: Debug,
{
// Verify state trie root
if block_builder.input.parent_state_trie.hash()
!= block_builder.input.parent_header.state_root
{
bail!(
"Invalid state trie: expected {}, got {}",
block_builder.input.parent_header.state_root,
block_builder.input.parent_state_trie.hash()
);
}

// hash all the contract code
let contracts: HashMap<B256, Bytes> = mem::take(&mut block_builder.input.contracts)
.into_iter()
.map(|bytes| (keccak(&bytes).into(), bytes))
.collect();

// Load account data into db
let mut accounts = HashMap::with_capacity(block_builder.input.parent_storage.len());
for (address, (storage_trie, slots)) in &mut block_builder.input.parent_storage {
// consume the slots, as they are no longer needed afterwards
let slots = mem::take(slots);

// load the account from the state trie or empty if it does not exist
let state_account = block_builder
.input
.parent_state_trie
.get_rlp::<StateAccount>(&keccak(address))?
.unwrap_or_default();
// Verify storage trie root
if storage_trie.hash() != state_account.storage_root {
bail!(
"Invalid storage trie for {:?}: expected {}, got {}",
address,
state_account.storage_root,
storage_trie.hash()
);
}

// load the corresponding code
let code_hash = to_revm_b256(state_account.code_hash);
let bytecode = if code_hash.0 == KECCAK_EMPTY.0 {
Bytecode::new()
} else {
let bytes = contracts.get(&code_hash).unwrap().clone();
unsafe { Bytecode::new_raw_with_hash(bytes.0, code_hash) }
};

// load storage reads
let mut storage = HashMap::with_capacity(slots.len());
for slot in slots {
let value: zeth_primitives::U256 = storage_trie
.get_rlp(&keccak(slot.to_be_bytes::<32>()))?
.unwrap_or_default();
storage.insert(slot, value);
}

let mem_account = DbAccount {
info: AccountInfo {
balance: state_account.balance,
nonce: state_account.nonce,
code_hash: to_revm_b256(state_account.code_hash),
code: Some(bytecode),
},
state: AccountState::None,
storage,
};

accounts.insert(*address, mem_account);
}
guest_mem_forget(contracts);

// prepare block hash history
let mut block_hashes =
HashMap::with_capacity(block_builder.input.ancestor_headers.len() + 1);
block_hashes.insert(
block_builder.input.parent_header.number,
to_revm_b256(block_builder.input.parent_header.hash()),
);
let mut prev = &block_builder.input.parent_header;
for current in &block_builder.input.ancestor_headers {
let current_hash = current.hash();
if prev.parent_hash != current_hash {
bail!(
"Invalid chain: {} is not the parent of {}",
current.number,
prev.number
);
}
if block_builder.input.parent_header.number < current.number
|| block_builder.input.parent_header.number - current.number >= MAX_BLOCK_HASH_AGE
{
bail!(
"Invalid chain: {} is not one of the {} most recent blocks",
current.number,
MAX_BLOCK_HASH_AGE,
);
}
block_hashes.insert(current.number, to_revm_b256(current_hash));
prev = current;
}

// Store database
block_builder.db = Some(Self::Db::load(accounts, block_hashes));

Ok(block_builder)
}
}
Loading

0 comments on commit 3cf8f46

Please sign in to comment.