fix(zk): invariant testing

crates/cheatcodes/src/inspector.rs

@@ -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,17 @@ 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 ecx
+            .db
+            .get_test_contract_address()
+            .map(|addr| call.bytecode_address == addr)
+            .unwrap_or_default()
+        {
+            info!(
+                "running call in EVM, instead of zkEVM (Test Contract) {:#?}",
+                call.bytecode_address
+            );
+            return None
         }
 
         info!("running call in zkEVM {:#?}", call);

crates/evm/core/src/backend/mod.rs

@@ -24,7 +24,7 @@ use revm::{
     precompile::{PrecompileSpecId, Precompiles},
     primitives::{
         Account, AccountInfo, Bytecode, Env, EnvWithHandlerCfg, EvmState, EvmStorageSlot,
-        HashMap as Map, Log, ResultAndState, SpecId, TxKind, KECCAK_EMPTY,
+        HashMap as Map, Log, ResultAndState, SpecId, KECCAK_EMPTY,
     },
     Database, DatabaseCommit, JournaledState,
 };
@@ -771,18 +771,6 @@ impl Backend {
     pub(crate) fn initialize(&mut self, env: &EnvWithHandlerCfg) {
         self.set_caller(env.tx.caller);
         self.set_spec_id(env.handler_cfg.spec_id);
-
-        let test_contract = match env.tx.transact_to {
-            TxKind::Call(to) => to,
-            TxKind::Create => {
-                let nonce = self
-                    .basic_ref(env.tx.caller)
-                    .map(|b| b.unwrap_or_default().nonce)
-                    .unwrap_or_default();
-                env.tx.caller.create(nonce)
-            }
-        };
-        self.set_test_contract(test_contract);
     }
 
     /// Returns the `EnvWithHandlerCfg` with the current `spec_id` set.

crates/evm/evm/src/executors/mod.rs

@@ -328,7 +328,7 @@ impl Executor {
         trace!(?from, ?to, "setting up contract");
 
         let from = from.unwrap_or(CALLER);
-        self.backend_mut().set_test_contract(to).set_caller(from);
+        self.backend_mut().set_caller(from);
         let calldata = Bytes::from_static(&ITest::setUpCall::SELECTOR);
         let mut res = self.transact_raw(from, to, calldata, U256::ZERO)?;
         res = res.into_result(rd)?;

crates/forge/src/runner.rs

@@ -118,6 +118,7 @@ impl<'a> ContractRunner<'a> {
         }
 
         let address = self.sender.create(self.executor.get_nonce(self.sender)?);
+        self.executor.backend_mut().set_test_contract(address);
 
         // Set the contracts initial balance before deployment, so it is available during
         // construction

crates/forge/tests/cli/test_cmd.rs

@@ -1048,7 +1048,7 @@ contract CallEmptyCode is Test {
 "#,
     )
     .unwrap();
-    cmd.args(["test", "--zksync", "--evm-version", "shanghai"]);
+    cmd.args(["test", "--zksync", "--evm-version", "shanghai", "--mc", "CallEmptyCode"]);
 
     let output = cmd.stdout_lossy();
     assert!(output.contains("call may fail or behave unexpectedly due to empty code"));

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

@@ -6,8 +6,14 @@ 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 runner = TEST_DATA_DEFAULT.runner_zksync();
+
+    // FIXME: just use the inline config
+    runner.test_options.invariant.no_zksync_reserved_addresses = true;
+    runner.test_options.invariant.fail_on_revert = true;
+    runner.test_options.invariant.runs = 10;
+
+    let filter = Filter::new(".*", "ZkInvariantTest", ".*");
 
     TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
 }

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

@@ -36,3 +36,10 @@ async fn repro_config(
 
 // https://github.com/matter-labs/foundry-zksync/issues/497
 test_repro!(497);
+
+test_repro!(565; |cfg| {
+    // FIXME: just use the inline config
+    cfg.runner.test_options.invariant.no_zksync_reserved_addresses = true;
+    cfg.runner.test_options.invariant.fail_on_revert = true;
+    cfg.runner.test_options.invariant.runs = 2;
+});

crates/script/src/runner.rs

@@ -129,6 +129,7 @@ impl ScriptRunner {
         };
 
         let address = CALLER.create(self.executor.get_nonce(CALLER)?);
+        self.executor.backend_mut().set_test_contract(address);
 
         // Set the contracts initial balance before deployment, so it is available during the
         // construction
@@ -147,7 +148,6 @@ impl ScriptRunner {
 
         // Optionally call the `setUp` function
         let (success, gas_used, labeled_addresses, transactions) = if !setup {
-            self.executor.backend_mut().set_test_contract(address);
             (true, 0, Default::default(), Some(library_transactions))
         } else {
             match self.executor.setup(Some(self.evm_opts.sender), address, None) {

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,149 @@
+// 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;
+    }
+
+    address[] internal _targetedSenders;
+
+    function targetSenders() public view returns (address[] memory) {
+        return _targetedSenders;
+    }
+}
+
+// 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;
+
+    uint256 constant dealAmount = 1 ether;
+
+    function setUp() public {
+        cnt = new Counter();
+
+        // so we can fund them ahead of time 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);
+        }
+
+        // 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 {}
+}