From 1cf88dc1e3118dadefdc823526b9f1966e5b98b5 Mon Sep 17 00:00:00 2001
From: Karrq <francesco@moonsonglabs.com>
Date: Tue, 20 Aug 2024 19:00:31 +0200
Subject: [PATCH] fix(zk): add logs to `recorded_logs` (#522)

* test(zk): recordLogs cheatcode

* fix(zk): add logs to `recorded_logs`

* chore: avoid unnecessary low-level call
---
 crates/cheatcodes/src/inspector.rs | 14 ++++++++++++++
 crates/forge/tests/it/zk/cheats.rs |  8 ++++++++
 testdata/zk/Cheatcodes.t.sol       | 22 ++++++++++++++++++++++
 zk-tests/src/Cheatcodes.t.sol      | 24 +++++++++++++++++++++++-
 4 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs
index 6057d0b1a..f036e55ea 100644
--- a/crates/cheatcodes/src/inspector.rs
+++ b/crates/cheatcodes/src/inspector.rs
@@ -922,6 +922,13 @@ impl Cheatcodes {
                 ecx,
                 ccx,
             ) {
+                if let Some(recorded_logs) = &mut self.recorded_logs {
+                    recorded_logs.extend(result.logs.clone().into_iter().map(|log| Vm::Log {
+                        topics: log.data.topics().to_vec(),
+                        data: log.data.data.clone(),
+                        emitter: log.address,
+                    }));
+                }
                 self.combined_logs.extend(result.logs.clone().into_iter().map(Some));
 
                 // for each log in cloned logs call handle_expect_emit
@@ -1380,6 +1387,13 @@ impl Cheatcodes {
                 persisted_factory_deps: Some(&mut self.persisted_factory_deps),
             };
             if let Ok(result) = foundry_zksync_core::vm::call::<_, DatabaseError>(call, ecx, ccx) {
+                if let Some(recorded_logs) = &mut self.recorded_logs {
+                    recorded_logs.extend(result.logs.clone().into_iter().map(|log| Vm::Log {
+                        topics: log.data.topics().to_vec(),
+                        data: log.data.data.clone(),
+                        emitter: log.address,
+                    }));
+                }
                 self.combined_logs.extend(result.logs.clone().into_iter().map(Some));
                 //for each log in cloned logs call handle_expect_emit
                 if !self.expected_emits.is_empty() {
diff --git a/crates/forge/tests/it/zk/cheats.rs b/crates/forge/tests/it/zk/cheats.rs
index f50a6a6f1..68f40350e 100644
--- a/crates/forge/tests/it/zk/cheats.rs
+++ b/crates/forge/tests/it/zk/cheats.rs
@@ -105,3 +105,11 @@ async fn test_zk_can_mock_modifiers() {
 
     TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_zk_record_logs() {
+    let runner = TEST_DATA_DEFAULT.runner_zksync();
+    let filter = Filter::new("RecordLogs", "ZkCheatcodesTest", ".*");
+
+    TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
+}
diff --git a/testdata/zk/Cheatcodes.t.sol b/testdata/zk/Cheatcodes.t.sol
index 2e82a4a99..5d6333d63 100644
--- a/testdata/zk/Cheatcodes.t.sol
+++ b/testdata/zk/Cheatcodes.t.sol
@@ -202,4 +202,26 @@ contract ZkCheatcodesTest is DSTest {
         vm.deal(0x4e59b44847b379578588920cA78FbF26c0B4956C, 1 ether);
         assertEq(1 ether, address(0x4e59b44847b379578588920cA78FbF26c0B4956C).balance);
     }
+
+    function testRecordLogsInZkVm() public {
+        // ensure we are in zkvm
+        vm.zkVm(true);
+        vm.recordLogs();
+        Emitter emitter = new Emitter(); // +7 logs from system contracts
+        emitter.functionEmit(); // +3 from system contracts
+
+        Vm.Log[] memory entries = vm.getRecordedLogs();
+        assertEq(entries.length, 12);
+        // 0,1: EthToken, 2,3: L1 Messanger, 4: Known Code Storage
+        assertEq(entries[5].topics.length, 1);
+        assertEq(entries[5].topics[0], keccak256("EventConstructor(string)"));
+        assertEq(entries[5].data, abi.encode("constructor"));
+        // 6: L2 Deployer, 7: EthToken
+
+        // 8,9: EthToken
+        assertEq(entries[10].topics.length, 1);
+        assertEq(entries[10].topics[0], keccak256("EventFunction(string)"));
+        assertEq(entries[10].data, abi.encode("function"));
+        // 11: EthToken
+    }
 }
diff --git a/zk-tests/src/Cheatcodes.t.sol b/zk-tests/src/Cheatcodes.t.sol
index d2ff9e7c7..696d04898 100644
--- a/zk-tests/src/Cheatcodes.t.sol
+++ b/zk-tests/src/Cheatcodes.t.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity ^0.8.13;
 
-import {Test, console2 as console} from "forge-std/Test.sol";
+import "forge-std/Test.sol";
 
 contract FixedSlot {
     uint8 num; // slot index: 0
@@ -241,4 +241,26 @@ contract ZkCheatcodesTest is Test {
             address(0x4e59b44847b379578588920cA78FbF26c0B4956C).balance
         );
     }
+
+    function testRecordLogsInZkVm() public {
+        // ensure we are in zkvm
+        (bool _success, bytes memory _ret) = address(vm).call(abi.encodeWithSignature("zkVm(bool)", true));
+        vm.recordLogs();
+        Emitter emitter = new Emitter(); // +7 logs from system contracts
+        emitter.functionEmit(); // +3 from system contracts
+
+        Vm.Log[] memory entries = vm.getRecordedLogs();
+        assertEq(entries.length, 12);
+        // 0,1: EthToken, 2,3: L1 Messanger, 4: Known Code Storage
+        assertEq(entries[5].topics.length, 1);
+        assertEq(entries[5].topics[0], keccak256("EventConstructor(string)"));
+        assertEq(entries[5].data, abi.encode("constructor"));
+        // 6: L2 Deployer, 7: EthToken
+
+        // 8,9: EthToken
+        assertEq(entries[10].topics.length, 1);
+        assertEq(entries[10].topics[0], keccak256("EventFunction(string)"));
+        assertEq(entries[10].data, abi.encode("function"));
+        // 11: EthToken
+    }
 }