feat: Implement log cheatcodes (#197)

crates/era-cheatcodes/src/cheatcodes.rs

@@ -17,14 +17,21 @@ use multivm::{
         },
     },
 };
-use std::{cell::RefMut, collections::HashMap, fmt::Debug, fs, process::Command};
+use serde::Serialize;
+use std::{
+    cell::RefMut,
+    collections::{HashMap, HashSet},
+    fmt::Debug,
+    fs,
+    process::Command,
+};
 use zksync_basic_types::{AccountTreeId, H160, H256, U256};
 use zksync_state::{ReadStorage, StoragePtr, StorageView, WriteStorage};
 use zksync_types::{
     block::{pack_block_info, unpack_block_info},
     get_code_key, get_nonce_key,
     utils::{decompose_full_nonce, nonces_to_full_nonce, storage_key_for_eth_balance},
-    LogQuery, StorageKey, Timestamp,
+    EventMessage, LogQuery, StorageKey, Timestamp,
 };
 use zksync_utils::{h256_to_u256, u256_to_h256};
 
@@ -66,6 +73,22 @@ pub struct CheatcodeTracer {
     return_ptr: Option<FatPointer>,
     near_calls: usize,
     serialized_objects: HashMap<String, String>,
+    recorded_logs: HashSet<LogEntry>,
+    recording_logs: bool,
+    recording_timestamp: u32,
+}
+
+#[derive(Debug, Clone, Serialize, Eq, Hash, PartialEq)]
+struct LogEntry {
+    topic: H256,
+    data: H256,
+    emitter: H160,
+}
+
+impl LogEntry {
+    fn new(topic: H256, data: H256, emitter: H160) -> Self {
+        LogEntry { topic, data, emitter }
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -155,6 +178,23 @@ impl<S: DatabaseExt + Send, H: HistoryMode> VmTracer<EraDb<S>, H> for CheatcodeT
         state: &mut ZkSyncVmState<EraDb<S>, H>,
         _bootloader_state: &mut BootloaderState,
     ) -> TracerExecutionStatus {
+        let emitter = state.local_state.callstack.current.this_address;
+        if self.recording_logs {
+            let logs = transform_to_logs(
+                state
+                    .event_sink
+                    .get_events_and_l2_l1_logs_after_timestamp(Timestamp(self.recording_timestamp))
+                    .0,
+                emitter,
+            );
+            if !logs.is_empty() {
+                let mut unique_set: HashSet<LogEntry> = HashSet::new();
+
+                // Filter out duplicates and extend the unique entries to the vector
+                self.recorded_logs
+                    .extend(logs.into_iter().filter(|log| unique_set.insert(log.clone())));
+            }
+        }
         while let Some(action) = self.one_time_actions.pop() {
             match action {
                 FinishCycleOneTimeActions::StorageWrite { key, read_value, write_value } => {
@@ -214,12 +254,15 @@ impl CheatcodeTracer {
             return_data: None,
             return_ptr: None,
             serialized_objects: HashMap::new(),
+            recorded_logs: HashSet::new(),
+            recording_logs: false,
+            recording_timestamp: 0,
         }
     }
 
     pub fn dispatch_cheatcode<S: DatabaseExt + Send, H: HistoryMode>(
         &mut self,
-        _state: VmLocalStateData<'_>,
+        state: VmLocalStateData<'_>,
         _data: AfterExecutionData,
         _memory: &SimpleMemory<H>,
         storage: StoragePtr<EraDb<S>>,
@@ -296,13 +339,47 @@ impl CheatcodeTracer {
                 tracing::info!("👷 Returndata is {:?}", account_nonce);
                 self.return_data = Some(vec![account_nonce]);
             }
+            getRecordedLogs(getRecordedLogsCall {}) => {
+                tracing::info!("👷 Getting recorded logs");
+                let logs: Vec<Log> = self
+                    .recorded_logs
+                    .iter()
+                    .map(|log| Log {
+                        topics: vec![log.topic.to_fixed_bytes().into()],
+                        data: log.data.to_fixed_bytes().into(),
+                        emitter: log.emitter.to_fixed_bytes().into(),
+                    })
+                    .collect_vec();
+
+                let result = getRecordedLogsReturn { logs };
+
+                let return_data: Vec<U256> =
+                    result.logs.abi_encode().chunks(32).map(|b| b.into()).collect_vec();
+
+                self.return_data = Some(return_data);
+
+                //clean up logs
+                self.recorded_logs = HashSet::new();
+                //disable flag of recording logs
+                self.recording_logs = false;
+            }
             load(loadCall { target, slot }) => {
                 tracing::info!("👷 Getting storage slot {:?} for account {:?}", slot, target);
                 let key = StorageKey::new(AccountTreeId::new(target.to_h160()), H256(*slot));
                 let mut storage = storage.borrow_mut();
                 let value = storage.read_value(&key);
                 self.return_data = Some(vec![h256_to_u256(value)]);
             }
+            recordLogs(recordLogsCall {}) => {
+                tracing::info!("👷 Recording logs");
+                tracing::info!(
+                    "👷 Logs will be with the timestamp {}",
+                    state.vm_local_state.timestamp
+                );
+
+                self.recording_timestamp = state.vm_local_state.timestamp;
+                self.recording_logs = true;
+            }
             readCallers(readCallersCall {}) => {
                 tracing::info!("👷 Reading callers");
 
@@ -533,6 +610,7 @@ impl CheatcodeTracer {
                 let int_value = value.to_string();
                 self.add_trimmed_return_data(int_value.as_bytes());
             }
+
             tryFfi(tryFfiCall { commandInput: command_input }) => {
                 tracing::info!("👷 Running try ffi: {command_input:?}");
                 let Some(first_arg) = command_input.get(0) else {
@@ -669,3 +747,15 @@ impl CheatcodeTracer {
         self.return_data = Some(data);
     }
 }
+fn transform_to_logs(events: Vec<EventMessage>, emitter: H160) -> Vec<LogEntry> {
+    events
+        .iter()
+        .filter_map(|event| {
+            if event.address == zksync_types::EVENT_WRITER_ADDRESS {
+                Some(LogEntry::new(u256_to_h256(event.key), u256_to_h256(event.value), emitter))
+            } else {
+                None
+            }
+        })
+        .collect()
+}

crates/era-cheatcodes/tests/src/cheatcodes/Logs.t.sol

@@ -0,0 +1,55 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import {Test, console2 as console} from "../../lib/forge-std/src/Test.sol";
+import {Constants} from "./Constants.sol";
+
+struct Log {
+    bytes32[] topics;
+    bytes data;
+    address emitter;
+}
+
+contract LogsTest is Test {
+    event LogTopic1(uint256 indexed topic1, bytes data);
+
+    function testRecordAndGetLogs() public {
+        bytes memory testData1 = "test";
+
+        (bool success, ) = Constants.CHEATCODE_ADDRESS.call(
+            abi.encodeWithSignature("recordLogs()")
+        );
+        require(success, "recordLogs failed");
+
+        emit LogTopic1(1, testData1);
+
+        (bool success2, bytes memory rawData) = Constants
+            .CHEATCODE_ADDRESS
+            .call(abi.encodeWithSignature("getRecordedLogs()"));
+        require(success2, "getRecordedLogs failed");
+
+        Log[] memory logs = abi.decode(rawData, (Log[]));
+        console.log("logs length: %d", logs.length);
+        require(logs.length == 8, "logs length should be 8");
+    }
+
+    function trimReturnBytes(
+        bytes memory rawData
+    ) internal pure returns (bytes memory) {
+        uint256 lengthStartingPos = rawData.length - 32;
+        bytes memory lengthSlice = new bytes(32);
+
+        for (uint256 i = 0; i < 32; i++) {
+            lengthSlice[i] = rawData[lengthStartingPos + i];
+        }
+
+        uint256 length = abi.decode(lengthSlice, (uint256));
+        bytes memory data = new bytes(length);
+
+        for (uint256 i = 0; i < length; i++) {
+            data[i] = rawData[i];
+        }
+
+        return data;
+    }
+}