Skip to content

Commit

Permalink
Perf: load code by hash (#126)
Browse files Browse the repository at this point in the history 
* Make MutDB generic

* Use code_by_hash to load contract code

* improve logging
  • Loading branch information
@Wollac
Wollac authored Dec 12, 2024
1 parent 2bb79b0 commit 80e3c5d
Show file tree
Hide file tree
Showing 5 changed files with 103 additions and 146 deletions.
2 changes: 1 addition & 1 deletion crates/core/src/stateless/data.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ pub struct StatelessClientData<Block, Header> {
/// Maps each address with its storage trie and the used storage slots.
pub storage_tries: HashMap<Address, StorageEntry>,
/// The code for each account
pub contracts: HashMap<Address, Bytes>,
pub contracts: Vec<Bytes>,
/// Immediate parent header
pub parent_header: Header,
/// List of at most 256 previous block headers
Expand Down
22 changes: 6 additions & 16 deletions crates/core/src/stateless/initialize.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ use alloy_primitives::{Address, Bytes, B256, U256};
use anyhow::bail;
use core::mem::take;
use reth_primitives::revm_primitives::Bytecode;
use reth_primitives::KECCAK_EMPTY;
use reth_revm::db::{AccountState, DbAccount};
use reth_revm::primitives::AccountInfo;
use std::default::Default;
Expand All @@ -32,7 +31,7 @@ pub trait InitializationStrategy<Driver: CoreDriver, Database> {
fn initialize_database(
state_trie: &mut MptNode,
storage_tries: &mut HashMap<Address, StorageEntry>,
contracts: &mut HashMap<Address, Bytes>,
contracts: &mut Vec<Bytes>,
parent_header: &mut Driver::Header,
ancestor_headers: &mut Vec<Driver::Header>,
) -> anyhow::Result<Database>;
Expand All @@ -44,7 +43,7 @@ impl<Driver: CoreDriver> InitializationStrategy<Driver, MemoryDB> for MemoryDbSt
fn initialize_database(
state_trie: &mut MptNode,
storage_tries: &mut HashMap<Address, StorageEntry>,
contracts: &mut HashMap<Address, Bytes>,
contracts: &mut Vec<Bytes>,
parent_header: &mut Driver::Header,
ancestor_headers: &mut Vec<Driver::Header>,
) -> anyhow::Result<MemoryDB> {
Expand All @@ -58,9 +57,9 @@ impl<Driver: CoreDriver> InitializationStrategy<Driver, MemoryDB> for MemoryDbSt
}

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

// Load account data into db
Expand All @@ -84,16 +83,6 @@ impl<Driver: CoreDriver> InitializationStrategy<Driver, MemoryDB> for MemoryDbSt
);
}

// load the corresponding code
let code_hash = state_account.code_hash;
let bytecode = if code_hash.0 == KECCAK_EMPTY.0 {
Bytecode::new()
} else {
let (bytecode_hash, bytes) = contracts.remove(address).unwrap();
assert_eq!(bytecode_hash, code_hash);
Bytecode::new_raw(bytes)
};

// load storage reads
let mut storage = HashMap::with_capacity_and_hasher(slots.len(), Default::default());
for slot in slots {
Expand All @@ -108,7 +97,7 @@ impl<Driver: CoreDriver> InitializationStrategy<Driver, MemoryDB> for MemoryDbSt
balance: state_account.balance,
nonce: state_account.nonce,
code_hash: state_account.code_hash,
code: Some(bytecode),
code: None,
},
account_state: AccountState::None,
storage,
Expand Down Expand Up @@ -150,6 +139,7 @@ impl<Driver: CoreDriver> InitializationStrategy<Driver, MemoryDB> for MemoryDbSt
// Initialize database
Ok(MemoryDB {
accounts,
contracts,
block_hashes,
..Default::default()
})
Expand Down
14 changes: 6 additions & 8 deletions crates/preflight/src/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use crate::provider::{new_provider, Provider};
use crate::trie::extend_proof_tries;
use alloy::network::Network;
use alloy::primitives::map::HashMap;
use alloy::primitives::Bytes;
use anyhow::Context;
use log::{debug, info, warn};
use std::cell::RefCell;
Expand Down Expand Up @@ -167,7 +168,7 @@ where
let core_parent_header = P::derive_header(data.parent_header.clone());
let mut state_trie = MptNode::from(R::state_root(&core_parent_header));
let mut storage_tries = Default::default();
let mut contracts = data.contracts.clone();
let mut contracts: Vec<Bytes> = Default::default();
let mut ancestor_headers: Vec<R::Header> = Default::default();

for num_blocks in 1..=block_count {
Expand Down Expand Up @@ -218,16 +219,13 @@ where
info!("Saving provider cache ...");
preflight_db.save_provider()?;

// collect the code from each account
info!("Collecting contracts ...");
// collect the code of the used contracts
let initial_db = preflight_db.inner.db.db.borrow();
for (address, account) in initial_db.accounts.iter() {
let code = account.info.code.clone().context("missing code")?;
if !code.is_empty() && !contracts.contains_key(address) {
contracts.insert(*address, code.bytes());
}
for code in initial_db.contracts.values() {
contracts.push(code.bytes().clone());
}
drop(initial_db);
info!("Collected contracts: {}", contracts.len());

// construct the sparse MPTs from the inclusion proofs
info!(
Expand Down
101 changes: 45 additions & 56 deletions crates/preflight/src/db.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,48 +26,33 @@ use reth_primitives::revm_primitives::{Account, AccountInfo, Bytecode};
use reth_revm::db::states::StateChangeset;
use reth_revm::db::CacheDB;
use reth_revm::{Database, DatabaseCommit, DatabaseRef};
use std::cell::RefCell;
use std::cell::{Ref, RefCell};
use std::marker::PhantomData;
use std::ops::DerefMut;
use zeth_core::db::apply_changeset;
use zeth_core::driver::CoreDriver;
use zeth_core::rescue::{Recoverable, Rescued};

#[derive(Clone)]
pub struct MutCacheDB<T: DatabaseRef> {
pub db: RefCell<CacheDB<T>>,
/// Wraps a [`Database`] to provide a [`DatabaseRef`] implementation.
#[derive(Clone, Debug, Default)]
pub struct MutDB<T> {
pub db: RefCell<T>,
}

impl<T: DatabaseRef> MutCacheDB<T> {
pub fn new(db: CacheDB<T>) -> Self {
impl<T: Database> MutDB<T> {
pub fn new(db: T) -> Self {
Self {
db: RefCell::new(db),
}
}
}

impl<T: DatabaseRef> Database for MutCacheDB<T> {
type Error = <CacheDB<T> as Database>::Error;

fn basic(&mut self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
self.db.borrow_mut().basic(address)
}

fn code_by_hash(&mut self, code_hash: B256) -> Result<Bytecode, Self::Error> {
self.db.borrow_mut().code_by_hash(code_hash)
}

fn storage(&mut self, address: Address, index: U256) -> Result<U256, Self::Error> {
self.db.borrow_mut().storage(address, index)
}

fn block_hash(&mut self, number: u64) -> Result<B256, Self::Error> {
self.db.borrow_mut().block_hash(number)
pub fn borrow_db(&self) -> Ref<T> {
self.db.borrow()
}
}

impl<T: DatabaseRef> DatabaseRef for MutCacheDB<T> {
type Error = <CacheDB<T> as DatabaseRef>::Error;
impl<T: Database> DatabaseRef for MutDB<T> {
type Error = <T as Database>::Error;

fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
self.db.borrow_mut().basic(address)
Expand All @@ -86,7 +71,7 @@ impl<T: DatabaseRef> DatabaseRef for MutCacheDB<T> {
}
}

pub type PrePostDB<N, R, P> = CacheDB<MutCacheDB<ProviderDB<N, R, P>>>;
pub type PrePostDB<N, R, P> = CacheDB<MutDB<CacheDB<MutDB<ProviderDB<N, R, P>>>>>;

#[derive(Clone)]
pub struct PreflightDB<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> {
Expand All @@ -109,7 +94,7 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> From<ProviderDB<N, R,
{
fn from(value: ProviderDB<N, R, P>) -> Self {
Self {
inner: CacheDB::new(MutCacheDB::new(CacheDB::new(value))),
inner: CacheDB::new(MutDB::new(CacheDB::new(MutDB::new(value)))),
driver: PhantomData,
}
}
Expand All @@ -136,17 +121,24 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> From<Rescued<PrePostDB

impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
pub fn clear(&mut self) -> anyhow::Result<()> {
let cleared = Self::from(self.inner.db.db.borrow().db.clone());
let cleared = Self::from(self.inner.db.borrow_db().db.borrow_db().clone());
drop(core::mem::replace(self, cleared));
Ok(())
}

pub fn save_provider(&mut self) -> anyhow::Result<()> {
self.inner.db.db.borrow_mut().db.save_provider()
self.inner.db.db.borrow_mut().db.db.borrow().save_provider()
}

pub fn advance_provider_block(&mut self) -> anyhow::Result<()> {
self.inner.db.db.borrow_mut().db.advance_provider_block()
self.inner
.db
.db
.borrow_mut()
.db
.db
.borrow_mut()
.advance_provider_block()
}

pub fn apply_changeset(&mut self, state_changeset: StateChangeset) -> anyhow::Result<()> {
Expand All @@ -156,11 +148,11 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
pub fn sanity_check(&mut self, state_changeset: StateChangeset) -> anyhow::Result<()> {
// storage sanity check
let initial_db = &self.inner.db;
let mut provider_db = initial_db.db.borrow().db.clone();
let mut provider_db = initial_db.db.borrow().db.db.borrow().clone();
provider_db.block_no += 1;
for (address, db_account) in &self.inner.accounts {
use reth_revm::Database;
let provider_info = provider_db.basic(*address)?.unwrap();
let provider_info = provider_db.basic(*address)?.unwrap_or_default();
if db_account.info != provider_info {
error!("State difference for account {address}:");
if db_account.info.balance != provider_info.balance {
Expand Down Expand Up @@ -198,13 +190,11 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
pub fn get_initial_proofs(
&mut self,
) -> anyhow::Result<HashMap<Address, EIP1186AccountProofResponse>> {
let initial_db = &self.inner.db;
let storage_keys = enumerate_storage_keys(&initial_db.db.borrow());

let initial_db = self.inner.db.db.borrow_mut();
let block_no = initial_db.db.block_no;
let initial_db = self.inner.db.borrow_db();
let storage_keys = enumerate_storage_keys(&initial_db);
let block_no = initial_db.db.borrow_db().block_no;
let res = get_proofs(
initial_db.db.provider.borrow_mut().deref_mut(),
initial_db.db.borrow_db().provider.borrow_mut().deref_mut(),
block_no,
storage_keys,
)?;
Expand All @@ -215,8 +205,8 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
&mut self,
) -> anyhow::Result<HashMap<Address, EIP1186AccountProofResponse>> {
// get initial keys
let initial_db = &self.inner.db;
let mut storage_keys = enumerate_storage_keys(&initial_db.db.borrow());
let initial_db = self.inner.db.borrow_db();
let mut storage_keys = enumerate_storage_keys(&initial_db);
// merge initial keys with latest db storage keys
for (address, mut indices) in enumerate_storage_keys(&self.inner) {
match storage_keys.get_mut(&address) {
Expand All @@ -227,27 +217,27 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
}
}
// return proofs as of next block
let initial_db = self.inner.db.db.borrow_mut();
let block_no = initial_db.db.block_no + 1;
let block_no = initial_db.db.borrow_db().block_no + 1;
let res = get_proofs(
initial_db.db.provider.borrow_mut().deref_mut(),
initial_db.db.borrow_db().provider.borrow_mut().deref_mut(),
block_no,
storage_keys,
)?;
Ok(res)
}

pub fn get_ancestor_headers(&mut self) -> anyhow::Result<Vec<N::HeaderResponse>> {
let initial_db = &self.inner.db.db.borrow_mut();
let db_block_number = initial_db.db.block_no;
let initial_db = self.inner.db.db.borrow_mut();
let db_block_number = initial_db.db.borrow_db().block_no;
let earliest_block = initial_db
.block_hashes
.keys()
.min()
.copied()
.map(|v| v.to())
.unwrap_or(db_block_number);
let mut provider = initial_db.db.provider.borrow_mut();
let provider_db = initial_db.db.borrow_db();
let mut provider = provider_db.provider.borrow_mut();
let headers = (earliest_block..db_block_number)
.rev()
.map(|block_no| {
Expand Down Expand Up @@ -277,14 +267,12 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
state_orphans: &[TrieOrphan],
block_count: u64,
) -> Vec<EIP1186AccountProofResponse> {
let initial_db = &self.inner.db.db.borrow_mut();
let mut provider = initial_db.db.provider.borrow_mut();
let initial_db = self.inner.db.db.borrow_mut();
let provider_db = initial_db.db.borrow_db();
let mut provider = provider_db.provider.borrow_mut();
let mut result = Vec::new();
let block_no = initial_db.db.block_no + block_count - 1;
let block_no = initial_db.db.borrow_db().block_no + block_count - 1;
for (start, digest) in state_orphans {
// if let Ok(val) = provider.get_preimage(&PreimageQuery { digest: *digest }) {
// continue;
// }
if let Ok(next_account) = provider.get_next_account(&AccountRangeQuery {
block_no,
start: *start,
Expand Down Expand Up @@ -313,10 +301,11 @@ impl<N: Network, R: CoreDriver, P: PreflightDriver<R, N>> PreflightDB<N, R, P> {
storage_orphans: &[(Address, TrieOrphan)],
block_count: u64,
) -> Vec<EIP1186AccountProofResponse> {
let initial_db = &self.inner.db.db.borrow_mut();
let mut provider = initial_db.db.provider.borrow_mut();
let initial_db = self.inner.db.db.borrow_mut();
let provider_db = initial_db.db.borrow_db();
let mut provider = provider_db.provider.borrow_mut();
let mut result = Vec::new();
let block_no = initial_db.db.block_no + block_count - 1;
let block_no = initial_db.db.borrow_db().block_no + block_count - 1;
for (address, (start, digest)) in storage_orphans {
// if let Ok(val) = provider.get_preimage(&PreimageQuery { digest: *digest }) {
// continue;
Expand Down
Loading

0 comments on commit 80e3c5d

Please sign in to comment.