fix(zk): invariant testing

crates/cheatcodes/src/inspector.rs

@@ -38,7 +38,7 @@ use foundry_zksync_compiler::{DualCompiledContract, DualCompiledContracts};
 use foundry_zksync_core::{
     convert::{ConvertH160, ConvertH256, ConvertRU256, ConvertU256},
     get_account_code_key, get_balance_key, get_nonce_key, Call, ZkTransactionMetadata,
-    DEFAULT_CREATE2_DEPLOYER_ZKSYNC,
+    DEFAULT_CREATE2_DEPLOYER_ZKSYNC, TEST_CONTRACT_ADDRESS_ZKSYNC,
 };
 use itertools::Itertools;
 use revm::{
@@ -48,7 +48,7 @@ use revm::{
     },
     primitives::{
         AccountInfo, BlockEnv, Bytecode, CreateScheme, EVMError, Env, EvmStorageSlot,
-        ExecutionResult, HashMap as rHashMap, Output, TransactTo, KECCAK_EMPTY,
+        ExecutionResult, HashMap as rHashMap, Output, KECCAK_EMPTY,
     },
     EvmContext, InnerEvmContext, Inspector,
 };
@@ -1530,11 +1530,12 @@ impl Cheatcodes {
             return None;
         }
 
-        if let TransactTo::Call(test_contract) = ecx.env.tx.transact_to {
-            if call.bytecode_address == test_contract {
-                info!("running call in EVM, instead of zkEVM (Test Contract) {:#?}", ecx.env.tx);
-                return None
-            }
+        if call.bytecode_address == TEST_CONTRACT_ADDRESS_ZKSYNC {
+            info!(
+                "running call in EVM, instead of zkEVM (Test Contract) {:#?}",
+                call.bytecode_address
+            );
+            return None
         }
 
         info!("running call in zkEVM {:#?}", call);

crates/forge/tests/it/zk/invariant.rs

@@ -1,13 +1,22 @@
 //! Invariant tests
 
-use crate::{config::*, test_helpers::TEST_DATA_DEFAULT};
+use crate::{
+    config::*,
+    test_helpers::{ForgeTestData, ForgeTestProfile},
+};
 use forge::revm::primitives::SpecId;
 use foundry_test_utils::Filter;
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_zk_invariant_deposit() {
-    let runner = TEST_DATA_DEFAULT.runner_zksync();
-    let filter = Filter::new("testZkInvariantDeposit", "ZkInvariantTest", ".*");
+    let mut test_data = ForgeTestData::new(ForgeTestProfile::Default);
+    // FIXME: just use the inline config
+    test_data.test_opts.invariant.no_zksync_reserved_addresses = true;
+    test_data.test_opts.invariant.fail_on_revert = true;
+    test_data.test_opts.invariant.runs = 10;
+
+    let runner = test_data.runner_zksync();
+    let filter = Filter::new(".*", "ZkInvariantTest", ".*");
 
     TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
 }

crates/forge/tests/it/zk/logs.rs

@@ -5,6 +5,7 @@ use std::collections::BTreeMap;
 use crate::{config::*, test_helpers::TEST_DATA_DEFAULT};
 use forge::revm::primitives::SpecId;
 use foundry_test_utils::Filter;
+use foundry_zksync_core::TEST_CONTRACT_ADDRESS_ZKSYNC;
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_zk_logs_work_in_call() {
@@ -23,7 +24,7 @@ async fn test_zk_logs_work_in_call() {
                 Some(vec![
                     "print".into(),
                     "outer print".into(),
-                    "0x7FA9385bE102ac3EAc297483Dd6233D62b3e1496".into(),
+                    TEST_CONTRACT_ADDRESS_ZKSYNC.to_string(),
                     "print".into(),
                     "0xff".into(),
                     "print".into(),

crates/forge/tests/it/zk/repros.rs

@@ -5,7 +5,7 @@
 use crate::{
     config::*,
     repros::test_repro,
-    test_helpers::{ForgeTestData, TEST_DATA_DEFAULT},
+    test_helpers::{ForgeTestData, ForgeTestProfile, TEST_DATA_DEFAULT},
 };
 use alloy_primitives::Address;
 use foundry_config::{fs_permissions::PathPermission, FsPermissions};
@@ -36,3 +36,14 @@ async fn repro_config(
 
 // https://github.com/matter-labs/foundry-zksync/issues/497
 test_repro!(497);
+
+#[tokio::test(flavor = "multi_thread")]
+async fn issue_565() {
+    let mut test_data = ForgeTestData::new(ForgeTestProfile::Default);
+    // FIXME: just use the inline config
+    test_data.test_opts.invariant.no_zksync_reserved_addresses = true;
+    test_data.test_opts.invariant.fail_on_revert = true;
+    test_data.test_opts.invariant.runs = 2;
+
+    repro_config(565, false, None, &test_data).await.run().await;
+}

crates/zksync/core/src/lib.rs

@@ -59,6 +59,10 @@ const MIN_VALID_ADDRESS: u32 = 2u32.pow(16);
 pub const DEFAULT_CREATE2_DEPLOYER_ZKSYNC: Address =
     address!("0000000000000000000000000000000000010000");
 
+/// The default test contract address in zkVM
+pub const TEST_CONTRACT_ADDRESS_ZKSYNC: Address =
+    address!("7fa9385be102ac3eac297483dd6233d62b3e1496");
+
 /// Returns the balance key for a provided account address.
 pub fn get_balance_key(address: Address) -> rU256 {
     storage_key_for_eth_balance(&address.to_h160()).key().to_ru256()

testdata/zk/Globals.sol

@@ -5,4 +5,6 @@ library Globals {
     string public constant ETHEREUM_MAINNET_URL =
         "https://eth-mainnet.alchemyapi.io/v2/Lc7oIGYeL_QvInzI0Wiu_pOZZDEKBrdf"; // trufflehog:ignore
     string public constant ZKSYNC_MAINNET_URL = "mainnet";
+
+    address public constant SYSTEM_CONTEXT_ADDR = address(0x000000000000000000000000000000000000800B);
 }

testdata/zk/InZkVm.sol

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.7 <0.9.0;
+
+import {Globals} from "./Globals.sol";
+
+// excerpt from system-contracts
+interface ISystemContext {
+    function chainId() external view returns (uint256);
+}
+
+library InZkVmLib {
+    function _inZkVm() internal returns (bool) {
+        (bool success, bytes memory retdata) =
+            Globals.SYSTEM_CONTEXT_ADDR.call(abi.encodeWithSelector(ISystemContext.chainId.selector));
+
+        return success;
+    }
+}
+
+abstract contract InZkVm {
+    modifier inZkVm() {
+        require(InZkVmLib._inZkVm(), "must be executed in zkVM");
+        _;
+    }
+}
+
+abstract contract DeployOnlyInZkVm is InZkVm {
+    constructor() {
+        require(InZkVmLib._inZkVm(), "must be deployed in zkVM");
+    }
+}

testdata/zk/InvariantDeposit.t.sol

@@ -5,25 +5,56 @@ import "ds-test/test.sol";
 import "../cheats/Vm.sol";
 import "./Deposit.sol";
 
-contract ZkInvariantTest is DSTest {
+// partial from forge-std/StdInvariant.sol
+abstract contract StdInvariant {
+    struct FuzzSelector {
+        address addr;
+        bytes4[] selectors;
+    }
+
+    address[] internal _targetedContracts;
+
+    function targetContracts() public view returns (address[] memory) {
+        return _targetedContracts;
+    }
+
+    FuzzSelector[] internal _targetedSelectors;
+
+    function targetSelectors() public view returns (FuzzSelector[] memory) {
+        return _targetedSelectors;
+    }
+
+    address[] internal _targetedSenders;
+
+    function targetSenders() public view returns (address[] memory) {
+        return _targetedSenders;
+    }
+}
+
+contract ZkInvariantTest is DSTest, StdInvariant {
     Vm constant vm = Vm(HEVM_ADDRESS);
-    // forge-config: default.invariant.runs = 2
     Deposit deposit;
 
+    uint256 constant dealAmount = 1 ether;
+
     function setUp() external {
+        // to fund for fees
+        _targetedSenders.push(address(65536 + 1));
+        _targetedSenders.push(address(65536 + 12));
+        _targetedSenders.push(address(65536 + 123));
+        _targetedSenders.push(address(65536 + 1234));
+
+        for (uint256 i = 0; i < _targetedSenders.length; i++) {
+            vm.deal(_targetedSenders[i], dealAmount); // to pay fees
+        }
+
         deposit = new Deposit();
-        vm.deal(address(deposit), 100 ether);
+        _targetedContracts.push(address(deposit));
     }
 
+    //FIXME: seems to not be detected, forcing values in test config
     // forge-config: default.invariant.runs = 2
-    function testZkInvariantDeposit() external payable {
-        deposit.deposit{value: 1 ether}();
-        uint256 balanceBefore = deposit.balance(address(this));
-        assertEq(balanceBefore, 1 ether);
-        deposit.withdraw();
-        uint256 balanceAfter = deposit.balance(address(this));
-        assertGt(balanceBefore, balanceAfter);
-    }
+    function invariant_itWorks() external payable {}
 
     receive() external payable {}
 }

testdata/zk/repros/Issue565.t.sol

@@ -0,0 +1,131 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity ^0.8.18;
+
+import "ds-test/test.sol";
+import "cheats/Vm.sol";
+import {Globals} from "../Globals.sol";
+import {DeployOnlyInZkVm} from "../InZkVm.sol";
+
+import "../../default/logs/console.sol";
+
+contract Counter is DeployOnlyInZkVm {
+    uint256 public number;
+
+    function inc() public inZkVm {
+        number += 1;
+    }
+
+    function reset() public inZkVm {
+        number = 0;
+    }
+}
+
+contract CounterHandler is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+
+    uint256 public incCounter;
+    uint256 public resetCounter;
+    bool public isResetLast;
+    Counter public counter;
+
+    constructor(Counter _counter) {
+        counter = _counter;
+    }
+
+    function inc() public {
+        console.log("inc");
+        incCounter += 1;
+        isResetLast = false;
+
+        vm.deal(tx.origin, 1 ether); // ensure caller has funds
+        counter.inc();
+    }
+
+    function reset() public {
+        console.log("reset");
+        resetCounter += 1;
+        isResetLast = true;
+
+        vm.deal(tx.origin, 1 ether); // ensure caller has funds
+        counter.reset();
+    }
+}
+
+// partial from forge-std/StdInvariant.sol
+abstract contract StdInvariant {
+    struct FuzzSelector {
+        address addr;
+        bytes4[] selectors;
+    }
+
+    address[] internal _targetedContracts;
+
+    function targetContracts() public view returns (address[] memory) {
+        return _targetedContracts;
+    }
+
+    FuzzSelector[] internal _targetedSelectors;
+
+    function targetSelectors() public view returns (FuzzSelector[] memory) {
+        return _targetedSelectors;
+    }
+}
+
+// https://github.com/matter-labs/foundry-zksync/issues/565
+contract Issue565 is DSTest, StdInvariant {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    Counter cnt;
+    CounterHandler handler;
+
+    function setUp() public {
+        cnt = new Counter();
+
+        vm.zkVmSkip();
+        handler = new CounterHandler(cnt);
+
+        // add the handler selectors to the fuzzing targets
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = CounterHandler.inc.selector;
+        selectors[1] = CounterHandler.reset.selector;
+
+        _targetedContracts.push(address(handler));
+        _targetedSelectors.push(FuzzSelector({addr: address(handler), selectors: selectors}));
+    }
+
+    //FIXME: seems to not be detected, forcing values in test config
+    /// forge-config: default.invariant.fail-on-revert = true
+    /// forge-config: default.invariant.no-zksync-reserved-addresses = true
+    function invariant_ghostVariables() external {
+        uint256 num = cnt.number();
+
+        if (handler.resetCounter() == 0) {
+            assert(handler.incCounter() == num);
+        } else if (handler.isResetLast()) {
+            assert(num == 0);
+        } else {
+            assert(num != 0);
+        }
+    }
+}
+
+contract Issue565WithoutHandler is DSTest, StdInvariant {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    Counter cnt;
+
+    function setUp() public {
+        cnt = new Counter();
+
+        // add the handler selectors to the fuzzing targets
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = Counter.inc.selector;
+        selectors[1] = Counter.reset.selector;
+
+        _targetedContracts.push(address(cnt));
+        _targetedSelectors.push(FuzzSelector({addr: address(cnt), selectors: selectors}));
+    }
+
+    //FIXME: seems to not be detected, forcing values in test config
+    /// forge-config: default.invariant.fail-on-revert = true
+    /// forge-config: default.invariant.no-zksync-reserved-addresses = true
+    function invariant_itWorks() external {}
+}