prestate tracer

core/lib/multivm/src/versions/vm_latest/tests/prestate_tracer.rs

@@ -0,0 +1,43 @@
+use std::borrow::Borrow;
+
+use crate::vm_latest::ToTracerPointer;
+use crate::{
+    interface::{TxExecutionMode, VmExecutionMode, VmInterface},
+    vm_latest::{
+        constants::BLOCK_GAS_LIMIT, tests::tester::VmTesterBuilder, tracers::PrestateTracer,
+        HistoryEnabled,
+    },
+};
+use zksync_test_account::TxType;
+use zksync_types::{Address, Execute, U256};
+
+#[test]
+fn test_prestate_tracer() {
+    let mut vm = VmTesterBuilder::new(HistoryEnabled)
+        .with_empty_in_memory_storage()
+        .with_random_rich_accounts(1)
+        .with_deployer()
+        .with_gas_limit(BLOCK_GAS_LIMIT)
+        .with_execution_mode(TxExecutionMode::VerifyExecute)
+        .build();
+
+    vm.deploy_test_contract();
+    let account = &mut vm.rich_accounts[0];
+
+    let tx1 = account.get_test_contract_transaction(
+        vm.test_contract.unwrap(),
+        false,
+        Default::default(),
+        true,
+        TxType::L2,
+    );
+
+    println!("tx: {:?}", vm.test_contract.unwrap());
+    vm.vm.push_transaction(tx1);
+
+    println!("account: {:?}", account);
+    let prestate_tracer = PrestateTracer::new();
+    let tracer_ptr = prestate_tracer.into_tracer_pointer();
+    let res = vm.vm.inspect(tracer_ptr.into(), VmExecutionMode::Batch);
+    assert!(1 == 0);
+}

core/lib/multivm/src/versions/vm_latest/tracers/default_tracers.rs

@@ -33,7 +33,7 @@ use crate::{
                 computational_gas_price, gas_spent_on_bytecodes_and_long_messages_this_opcode,
                 print_debug_if_needed, VmHook,
             },
-            CircuitsTracer, RefundsTracer, ResultTracer,
+            CircuitsTracer, PrestateTracer, RefundsTracer, ResultTracer,
         },
         types::internals::ZkSyncVmState,
         VmTracer,
@@ -67,6 +67,7 @@ pub(crate) struct DefaultExecutionTracer<S: WriteStorage, H: HistoryMode> {
     // It only takes into account circuits that are generated for actual execution. It doesn't
     // take into account e.g circuits produced by the initial bootloader memory commitment.
     pub(crate) circuits_tracer: CircuitsTracer<S, H>,
+    pub(crate) prestate_tracer: PrestateTracer,
 
     storage: StoragePtr<S>,
     _phantom: PhantomData<H>,
@@ -81,6 +82,7 @@ impl<S: WriteStorage, H: HistoryMode> DefaultExecutionTracer<S, H> {
         refund_tracer: Option<RefundsTracer<S>>,
         pubdata_tracer: Option<PubdataTracer<S>>,
     ) -> Self {
+        println!("NEW DefaultExecutionTracer");
         Self {
             tx_has_been_processed: false,
             execution_mode,
@@ -95,6 +97,7 @@ impl<S: WriteStorage, H: HistoryMode> DefaultExecutionTracer<S, H> {
             pubdata_tracer,
             ret_from_the_bootloader: None,
             circuits_tracer: CircuitsTracer::new(),
+            prestate_tracer: PrestateTracer::new(),
             storage,
             _phantom: PhantomData,
         }

core/lib/multivm/src/versions/vm_latest/tracers/mod.rs

@@ -1,11 +1,13 @@
 pub(crate) use circuits_tracer::CircuitsTracer;
 pub(crate) use default_tracers::DefaultExecutionTracer;
+pub(crate) use prestate_tracer::PrestateTracer;
 pub(crate) use pubdata_tracer::PubdataTracer;
 pub(crate) use refunds::RefundsTracer;
 pub(crate) use result_tracer::ResultTracer;
 
 pub(crate) mod circuits_tracer;
 pub(crate) mod default_tracers;
+pub(crate) mod prestate_tracer;
 pub(crate) mod pubdata_tracer;
 pub(crate) mod refunds;
 pub(crate) mod result_tracer;

core/lib/multivm/src/versions/vm_latest/tracers/prestate_tracer.rs

@@ -0,0 +1,172 @@
+use std::fmt;
+use std::{collections::HashMap, marker::PhantomData};
+
+use crate::{
+    interface::{
+        dyn_tracers::vm_1_4_1::DynTracer,
+        tracer::{TracerExecutionStatus, TracerExecutionStopReason},
+        types::inputs::L1BatchEnv,
+        VmExecutionMode,
+    },
+    vm_latest::{
+        bootloader_state::{utils::apply_pubdata_to_memory, BootloaderState},
+        constants::BOOTLOADER_HEAP_PAGE,
+        old_vm::{history_recorder::HistoryMode, memory::SimpleMemory},
+        tracers::{traits::VmTracer, utils::VmHook},
+        types::internals::{PubdataInput, ZkSyncVmState},
+        utils::logs::collect_events_and_l1_system_logs_after_timestamp,
+        StorageOracle,
+    },
+};
+use zk_evm_1_4_1::{
+    aux_structures::Timestamp,
+    tracing::{BeforeExecutionData, VmLocalStateData},
+};
+use zksync_state::{StoragePtr, WriteStorage};
+use zksync_types::{
+    get_code_key, get_nonce_key, web3::signing::keccak256, AccountTreeId, Address, Bytes,
+    StorageKey, StorageValue, H160, H256, L2_ETH_TOKEN_ADDRESS, U256,
+};
+use zksync_utils::address_to_h256;
+
+#[derive(Debug, Clone)]
+pub struct Account {
+    balance: Option<U256>,
+    code: Option<U256>,
+    nonce: Option<U256>,
+    storage: Option<HashMap<H256, H256>>,
+}
+
+impl fmt::Display for Account {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        writeln!(f, "{{")?;
+        if let Some(balance) = self.balance {
+            writeln!(f, "  balance: \"0x{:x}\",", balance)?;
+        }
+        if let Some(code) = &self.code {
+            writeln!(f, "  code: \"{}\",", code)?;
+        }
+        if let Some(nonce) = self.nonce {
+            writeln!(f, "  nonce: {},", nonce)?;
+        }
+        if let Some(storage) = &self.storage {
+            writeln!(f, "  storage: {{")?;
+            for (key, value) in storage.iter() {
+                writeln!(f, "    {}: \"{}\",", key, value)?;
+            }
+            writeln!(f, "  }}")?;
+        }
+        writeln!(f, "}}")
+    }
+}
+
+type State = HashMap<Address, Account>;
+
+#[derive(Debug, Clone)]
+pub(crate) struct PrestateTracer {
+    pub pre: State,
+    pub post: State,
+    pub config: PrestateTracerConfig,
+}
+
+impl PrestateTracer {
+    pub(crate) fn new() -> Self {
+        println!("NEW");
+        Self {
+            pre: Default::default(),
+            post: Default::default(),
+            config: PrestateTracerConfig { diff_mode: false },
+        }
+    }
+
+    pub(crate) fn get_pre(&self) -> &State {
+        &self.pre
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct PrestateTracerConfig {
+    diff_mode: bool, // If true, this tracer will return state modifications
+}
+
+impl<S, H: HistoryMode> DynTracer<S, SimpleMemory<H>> for PrestateTracer {
+    fn before_execution(
+        &mut self,
+        state: VmLocalStateData<'_>,
+        data: BeforeExecutionData,
+        _memory: &SimpleMemory<H>,
+        _storage: StoragePtr<S>,
+    ) {
+    }
+}
+
+impl<S: WriteStorage, H: HistoryMode> VmTracer<S, H> for PrestateTracer {
+    fn initialize_tracer(&mut self, _state: &mut ZkSyncVmState<S, H>) {}
+
+    fn finish_cycle(
+        &mut self,
+        _state: &mut ZkSyncVmState<S, H>,
+        _bootloader_state: &mut BootloaderState,
+    ) -> TracerExecutionStatus {
+        TracerExecutionStatus::Continue
+    }
+
+    fn after_vm_execution(
+        &mut self,
+        state: &mut ZkSyncVmState<S, H>,
+        _bootloader_state: &BootloaderState,
+        _stop_reason: crate::interface::tracer::VmExecutionStopReason,
+    ) {
+        let modified_storage_keys = {
+            let binding = state.storage.storage.inner().storage_ptr.borrow_mut();
+            // Clone the keys or derive the needed information here
+            binding.modified_storage_keys().clone() // Assuming `clone` is possible and inexpensive
+        };
+
+        let read_keys = &state.storage.read_keys;
+        let map = read_keys.inner().clone();
+        let storage_keys_read = map.iter().map(|(k, _)| k.clone()).collect::<Vec<_>>();
+        let pre = storage_keys_read
+            .iter()
+            .map(|k| {
+                let nonce_key = get_nonce_key(k.account().address());
+                let code_key = get_code_key(k.account().address());
+                let nonce = state.storage.storage.read_from_storage(&nonce_key);
+                let address_h256 = address_to_h256(k.account().address());
+                let bytes = [address_h256.as_bytes(), &[0; 32]].concat();
+                let balance_key: H256 = keccak256(&bytes).into();
+                let balance_storage_key =
+                    StorageKey::new(AccountTreeId::new(L2_ETH_TOKEN_ADDRESS), balance_key);
+                let balance = state
+                    .storage
+                    .storage
+                    .read_from_storage(&balance_storage_key);
+
+                let code = state.storage.storage.read_from_storage(&code_key);
+                let storage = get_storage_if_present(k.account(), &modified_storage_keys);
+                (
+                    *(k.account().address()),
+                    Account {
+                        balance: Some(balance),
+                        code: Some(code),
+                        nonce: Some(nonce),
+                        storage: Some(storage),
+                    },
+                )
+            })
+            .collect::<State>();
+        self.pre = pre;
+    }
+}
+
+fn get_storage_if_present(
+    account: &AccountTreeId,
+    modified_storage_keys: &HashMap<StorageKey, StorageValue>,
+) -> HashMap<H256, H256> {
+    //check if there is a StorageKey struct wioth an account field that matches the account and return the key as the key and the StorageValue as the value
+    modified_storage_keys
+        .iter()
+        .filter(|(k, _)| k.account() == account)
+        .map(|(k, v)| (k.key().clone(), v.clone()))
+        .collect()
+}