Skip to content

Commit

Permalink
feat: change cheatcode tracer to have access to InMemoryNode
Browse files Browse the repository at this point in the history
  • Loading branch information
@aon
aon committed Nov 13, 2023
1 parent c87aa4d commit 24f462f
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 21 deletions.
41 changes: 24 additions & 17 deletions src/cheatcodes.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use crate::{fork::ForkStorage, utils::bytecode_to_factory_dep};
use crate::{node::InMemoryNodeInner, utils::bytecode_to_factory_dep};
use anyhow::{anyhow, Result};
use ethers::{abi::AbiDecode, prelude::abigen};
use multivm::{
vm_1_3_2::zk_evm_1_3_3::{
Expand All @@ -10,6 +11,7 @@ use multivm::{
DynTracer, ExecutionEndTracer, ExecutionProcessing, HistoryMode, SimpleMemory, VmTracer,
},
};
use std::sync::{Arc, RwLock};
use zksync_basic_types::{H160, H256};
use zksync_state::{StoragePtr, WriteStorage};
use zksync_types::{
Expand All @@ -24,12 +26,12 @@ const CHEATCODE_ADDRESS: H160 = H160([
]);

#[derive(Clone, Debug, Default)]
pub struct CheatcodeTracer<F: Send + Sync> {
pub struct CheatcodeTracer<F> {
factory_deps: F,
}

pub trait FactoryDeps {
fn store_factory_dep(&mut self, hash: H256, bytecode: Vec<u8>);
fn store_factory_dep(&mut self, hash: H256, bytecode: Vec<u8>) -> Result<()>;
}

abigen!(
Expand All @@ -41,9 +43,7 @@ abigen!(
]"#
);

impl<F: FactoryDeps + Send + Sync, S: WriteStorage, H: HistoryMode> DynTracer<S, H>
for CheatcodeTracer<F>
{
impl<F: FactoryDeps, S: WriteStorage, H: HistoryMode> DynTracer<S, H> for CheatcodeTracer<F> {
fn before_execution(
&mut self,
state: VmLocalStateData<'_>,
Expand Down Expand Up @@ -94,17 +94,14 @@ impl<F: FactoryDeps + Send + Sync, S: WriteStorage, H: HistoryMode> DynTracer<S,
}
}

impl<F: FactoryDeps + Send + Sync, S: WriteStorage, H: HistoryMode> VmTracer<S, H>
for CheatcodeTracer<F>
{
}
impl<F: FactoryDeps + Send + Sync, H: HistoryMode> ExecutionEndTracer<H> for CheatcodeTracer<F> {}
impl<F: FactoryDeps + Send + Sync, S: WriteStorage, H: HistoryMode> ExecutionProcessing<S, H>
impl<F: FactoryDeps + Send, S: WriteStorage, H: HistoryMode> VmTracer<S, H> for CheatcodeTracer<F> {}
impl<F: FactoryDeps, H: HistoryMode> ExecutionEndTracer<H> for CheatcodeTracer<F> {}
impl<F: FactoryDeps, S: WriteStorage, H: HistoryMode> ExecutionProcessing<S, H>
for CheatcodeTracer<F>
{
}

impl<F: FactoryDeps + Send + Sync> CheatcodeTracer<F> {
impl<F: FactoryDeps> CheatcodeTracer<F> {
pub fn new(factory_deps: F) -> Self {
Self { factory_deps }
}
Expand All @@ -130,7 +127,7 @@ impl<F: FactoryDeps + Send + Sync> CheatcodeTracer<F> {
let code_key = get_code_key(&who);
let (hash, code) = bytecode_to_factory_dep(code.0.into());
let hash = u256_to_h256(hash);
self.factory_deps.store_factory_dep(
if let Err(err) = self.factory_deps.store_factory_dep(
hash,
code.iter()
.flat_map(|entry| {
Expand All @@ -139,7 +136,13 @@ impl<F: FactoryDeps + Send + Sync> CheatcodeTracer<F> {
bytes.to_vec()
})
.collect(),
) {
tracing::error!(
"Etch cheatcode failed, failed to store factory dep: {:?}",
err
);
return;
}
storage.borrow_mut().set_value(code_key, hash);
}
SetNonce(SetNonceCall { account, nonce }) => {
Expand Down Expand Up @@ -179,9 +182,13 @@ impl<F: FactoryDeps + Send + Sync> CheatcodeTracer<F> {
}
}

impl<T> FactoryDeps for ForkStorage<T> {
fn store_factory_dep(&mut self, hash: H256, bytecode: Vec<u8>) {
self.store_factory_dep(hash, bytecode)
impl<T> FactoryDeps for Arc<RwLock<InMemoryNodeInner<T>>> {
fn store_factory_dep(&mut self, hash: H256, bytecode: Vec<u8>) -> Result<()> {
self.try_write()
.map_err(|e| anyhow!(format!("Failed to grab write lock: {e}")))?
.fork_storage
.store_factory_dep(hash, bytecode);
Ok(())
}
}

Expand Down
26 changes: 22 additions & 4 deletions src/node/in_memory.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -999,7 +999,8 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
.write()
.map_err(|e| format!("Failed to acquire write lock: {}", e))?;

let storage = StorageView::new(&inner.fork_storage).to_rc_ptr();
let fork_storage = inner.fork_storage.clone();
let storage = StorageView::new(&fork_storage).to_rc_ptr();

let bootloader_code = inner.system_contracts.contracts_for_l2_call();

Expand All @@ -1020,16 +1021,24 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo

let call_tracer_result = Arc::new(OnceCell::default());

let cheatcode_tracer = CheatcodeTracer::new(inner.fork_storage.clone());
let cheatcode_tracer = CheatcodeTracer::new(self.get_inner());
let custom_tracers = vec![
Box::new(cheatcode_tracer)
as Box<dyn VmTracer<StorageView<&ForkStorage<S>>, HistoryDisabled>>,
Box::new(CallTracer::new(call_tracer_result.clone(), HistoryDisabled))
as Box<dyn VmTracer<StorageView<&ForkStorage<S>>, HistoryDisabled>>,
];

// Drop inner to allow `CheatcodeTracer` to write to `InMemoryNode`
drop(inner);

let tx_result = vm.inspect(custom_tracers, VmExecutionMode::OneTx);

let inner = self
.inner
.write()
.map_err(|e| format!("Failed to acquire write lock: {}", e))?;

let call_traces = Arc::try_unwrap(call_tracer_result)
.unwrap()
.take()
Expand Down Expand Up @@ -1266,7 +1275,8 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
.write()
.map_err(|e| format!("Failed to acquire write lock: {}", e))?;

let storage = StorageView::new(&inner.fork_storage).to_rc_ptr();
let fork_storage = inner.fork_storage.clone();
let storage = StorageView::new(&fork_storage).to_rc_ptr();

let (batch_env, block_ctx) = inner.create_l1_batch_env(storage.clone());

Expand Down Expand Up @@ -1303,7 +1313,7 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo

let call_tracer_result = Arc::new(OnceCell::default());
let bootloader_debug_result = Arc::new(OnceCell::default());
let cheatcode_tracer = CheatcodeTracer::new(inner.fork_storage.clone());
let cheatcode_tracer = CheatcodeTracer::new(self.get_inner());

let custom_tracers = vec![
Box::new(cheatcode_tracer)
Expand All @@ -1315,8 +1325,16 @@ impl<S: ForkSource + std::fmt::Debug + Clone + Send + Sync + 'static> InMemoryNo
}) as Box<dyn VmTracer<StorageView<&ForkStorage<S>>, HistoryDisabled>>,
];

// Drop inner to allow `CheatcodeTracer` to write to `InMemoryNode`
drop(inner);

let tx_result = vm.inspect(custom_tracers, VmExecutionMode::OneTx);

let inner = self
.inner
.write()
.map_err(|e| format!("Failed to acquire write lock: {}", e))?;

let call_traces = call_tracer_result.get().unwrap();

let spent_on_pubdata =
Expand Down

0 comments on commit 24f462f

Please sign in to comment.