Mock function cheatcode (#680)

* Add k code for mockFunction cheatcode

* Fix k code

* Add test

* Set Version: 0.1.346

* Add tests for mock functions with args, provided and not. Fix K rule so it handles these cases

* Set Version: 0.1.348

* enable tests and update contract and foundry files

* Update expected files to show new cheatcode cell

* Set Version: 0.1.349

* Update kontrol cheatcodes ref

* Update kontrol-cheatcodes ref in workflow file

* Set Version: 0.1.350

* Update bytecode in circularity lemma

* Update expected output

* Set Version: 0.1.351

* Set Version: 0.1.353

* Improve formatting

* Update src/kontrol/kdist/cheatcodes.md

* Set Version: 0.1.356

* Set Version: 0.1.357

* Set Version: 0.1.358

---------

Co-authored-by: devops <devops@runtimeverification.com>
Co-authored-by: Palina Tolmach <polina.tolmach@gmail.com>

.github/workflows/test-pr.yml

@@ -223,7 +223,7 @@ jobs:
       - name: 'Run forge build'
         run: |
           docker exec --user ${DOCKER_USER} --workdir ${FOUNDRY_ROOT} ${CONTAINER_NAME} forge install --no-git foundry-rs/forge-std@75f1746
-          docker exec --user ${DOCKER_USER} --workdir ${FOUNDRY_ROOT} ${CONTAINER_NAME} forge install --no-git runtimeverification/kontrol-cheatcodes@0048278
+          docker exec --user ${DOCKER_USER} --workdir ${FOUNDRY_ROOT} ${CONTAINER_NAME} forge install --no-git runtimeverification/kontrol-cheatcodes@7baaf34
           docker exec --user ${DOCKER_USER} --workdir ${FOUNDRY_ROOT} ${CONTAINER_NAME} forge build
       - name: 'Run kontrol build'
         run: docker exec --user ${DOCKER_USER} --workdir ${FOUNDRY_ROOT} ${CONTAINER_NAME} kontrol build

package/version

@@ -1 +1 @@
-0.1.357
+0.1.358

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "kontrol"
-version = "0.1.357"
+version = "0.1.358"
 description = "Foundry integration for KEVM"
 authors = [
     "Runtime Verification, Inc. <contact@runtimeverification.com>",

src/kontrol/__init__.py

@@ -5,4 +5,4 @@
 if TYPE_CHECKING:
     from typing import Final
 
-VERSION: Final = '0.1.357'
+VERSION: Final = '0.1.358'

src/kontrol/kdist/cheatcodes.md

@@ -86,6 +86,12 @@ module FOUNDRY-CHEAT-CODES
                <mockValues>  .Map </mockValues>
             </mockCall>
          </mockCalls>
+        <mockFunctions>
+            <mockFunction multiplicity="*" type="Map">
+               <mockFunctionAddress> .Account </mockFunctionAddress>
+               <mockFunctionValues>  .Map </mockFunctionValues>
+            </mockFunction>
+         </mockFunctions>
       </cheatcodes>
 ```
 
@@ -962,6 +968,45 @@ We use `#next[OP]` to identify OpCodes that represent function calls. If there i
       [priority(30)]
 ```
 
+
+Mock functions
+--------------
+
+#### `mockFunction` - Treats all calls to callee strictly or loosely matching data as if they are instead called on calledContract. 
+
+
+```
+    function mockFunction(address callee, address calledContract, bytes calldata data) external;
+```
+
+```k
+    rule [cheatcode.call.mockFunction]:
+         <k> #cheatcode_call SELECTOR ARGS
+          => #loadAccount #asWord(#range(ARGS, 0, 32))
+          ~> #etchAccountIfEmpty #asWord(#range(ARGS, 0, 32))
+          ~> #setMockFunction #asWord(#range(ARGS, 0, 32)) #asWord(#range(ARGS, 32, 32)) #range(ARGS, 128, #asWord(#range(ARGS, 96, 32)))
+         ...
+         </k>
+      requires SELECTOR ==Int selector ( "mockFunction(address,address,bytes)" )
+
+    rule [foundry.set.mockFunction]:
+         <k> #call ACCTFROM ACCTTO ACCTCODE VALUE APPVALUE CALLDATA STATIC
+          => #callWithCode ACCTFROM ACCTTO ACCTCODE CODE VALUE APPVALUE CALLDATA STATIC
+         ...
+         </k>
+         <account>
+           <acctID> MOCKTARGET </acctID>
+           <code> CODE </code>
+           ...
+         </account>
+         <mockFunction>
+           <mockFunctionAddress> ACCTCODE </mockFunctionAddress>
+           <mockFunctionValues>...  CALLDATA_KEY |-> MOCKTARGET ...</mockFunctionValues>
+         </mockFunction>
+         requires #range(CALLDATA, 0, lengthBytes(CALLDATA_KEY)) ==K CALLDATA_KEY
+      [priority(30)]
+```
+
 Utils
 -----
 
@@ -1433,6 +1478,30 @@ If the flag is false, it skips comparison, assuming success; otherwise, it compa
          </mockCalls>
 ```
 
+- `#setMockFunction MOCKADDRESS MOCKTARGET MOCKCALLDATA` will update the `<mockFunctions>` mapping for the given account and calldata.
+
+
+```k
+    syntax KItem ::= "#setMockFunction" Account Account Bytes [symbol(foundry_setMockFunction)]
+ // ---------------------------------------------------------------------------------
+    rule <k> #setMockFunction MOCKADDRESS MOCKTARGET MOCKCALLDATA => .K ... </k>
+         <mockFunction>
+            <mockFunctionAddress> MOCKADDRESS </mockFunctionAddress>
+            <mockFunctionValues>  MOCKVALUES => MOCKVALUES [ MOCKCALLDATA <- MOCKTARGET ] </mockFunctionValues>
+         </mockFunction>
+
+   rule <k> #setMockFunction MOCKADDRESS MOCKTARGET MOCKCALLDATA => .K ... </k>
+         <mockFunctions>
+           ( .Bag
+            => <mockFunction>
+                  <mockFunctionAddress> MOCKADDRESS </mockFunctionAddress>
+                  <mockFunctionValues> .Map [ MOCKCALLDATA <- MOCKTARGET ] </mockFunctionValues>
+               </mockFunction>
+           )
+           ...
+         </mockFunctions>
+```
+
 - `#execMockCall` will update the output of the function call with `RETURNDATA` using `#setLocalMem`. In case the function did not end with `EVMC_SUCCESS` it will update the status code to `EVMC_SUCCESS`.
 
 ```k
@@ -1486,6 +1555,7 @@ If the flag is false, it skips comparison, assuming success; otherwise, it compa
     rule ( selector ( "infiniteGas()" )                            => 3986649939 )
     rule ( selector ( "setGas(uint256)" )                          => 3713137314 )
     rule ( selector ( "mockCall(address,bytes,bytes)" )            => 3110212580 )
+    rule ( selector ( "mockFunction(address,address,bytes)" )      => 2918731041 )
 ```
 
 - selectors for unimplemented cheat code functions.

src/kontrol/prove.py

@@ -903,6 +903,7 @@ def _init_cterm(
         'ADDRESSSET_CELL': set_empty(),
         'STORAGESLOTSET_CELL': set_empty(),
         'MOCKCALLS_CELL': KApply('.MockCallCellMap'),
+        'MOCKFUNCTIONS_CELL': KApply('.MockFunctionCellMap'),
         'ACTIVETRACING_CELL': TRUE if trace_options.active_tracing else FALSE,
         'TRACESTORAGE_CELL': TRUE if trace_options.trace_storage else FALSE,
         'TRACEWORDSTACK_CELL': TRUE if trace_options.trace_wordstack else FALSE,

src/tests/integration/conftest.py

@@ -25,7 +25,7 @@
 
 
 FORGE_STD_REF: Final = '75f1746'
-KONTROL_CHEATCODES_REF: Final = '0048278'
+KONTROL_CHEATCODES_REF: Final = '7baaf34'
 
 
 @pytest.fixture

src/tests/integration/test-data/foundry-prove-all

@@ -182,6 +182,9 @@ MockCallRevertTest.testMockCallResetsMockCallRevert()
 MockCallRevertTest.testMockCallRevertResetsMockCall()
 MockCallRevertTest.testMockCallRevertWithCall()
 MockCallRevertTest.testMockCallEmptyAccount()
+MockFunctionTest.test_mock_function()
+MockFunctionTest.test_mock_function_all_args()
+MockFunctionTest.test_mock_function_concrete_args()
 OwnerUpOnlyTest.testFailIncrementAsNotOwner()
 OwnerUpOnlyTest.testIncrementAsNotOwner()
 OwnerUpOnlyTest.testIncrementAsOwner()

src/tests/integration/test-data/foundry/test/MockFunction.t.sol

@@ -0,0 +1,75 @@
+pragma solidity =0.8.13;
+
+import "forge-std/Test.sol";
+
+import "kontrol-cheatcodes/KontrolCheats.sol";
+
+contract MockFunctionContract {
+    uint256 public a;
+
+    function mocked_function() public {
+        a = 321;
+    }
+
+    function mocked_args_function(uint256 x) public {
+        a = 321 + x;
+    }
+}
+
+contract ModelMockFunctionContract {
+    uint256 public a;
+
+    function mocked_function() public {
+        a = 123;
+    }
+
+    function mocked_args_function(uint256 x) public {
+        a = 123 + x;
+    }
+}
+
+contract MockFunctionTest is Test, KontrolCheats {
+    MockFunctionContract my_contract;
+    ModelMockFunctionContract model_contract;
+
+    function setUp() public {
+        my_contract = new MockFunctionContract();
+        model_contract = new ModelMockFunctionContract();
+    }
+
+    function test_mock_function() public {
+        kevm.mockFunction(
+            address(my_contract), 
+            address(model_contract),
+            abi.encodeWithSelector(MockFunctionContract.mocked_function.selector)
+        );
+        my_contract.mocked_function();
+        assertEq(my_contract.a(), 123);
+    }
+
+    function test_mock_function_concrete_args() public {
+        kevm.mockFunction(
+            address(my_contract), 
+            address(model_contract),
+            abi.encodeWithSelector(MockFunctionContract.mocked_args_function.selector, 456)
+        );
+        my_contract.mocked_args_function(456);
+        assertEq(my_contract.a(), 123 + 456);
+
+        my_contract.mocked_args_function(567);
+        assertEq(my_contract.a(), 321 + 567);
+    }
+
+    function test_mock_function_all_args() public {
+        kevm.mockFunction(
+            address(my_contract), 
+            address(model_contract),
+            abi.encodeWithSelector(MockFunctionContract.mocked_args_function.selector)
+        );
+        my_contract.mocked_args_function(678);
+        assertEq(my_contract.a(), 123 + 678);
+
+        my_contract.mocked_args_function(789);
+        assertEq(my_contract.a(), 123 + 789);
+    }
+}