From 1d9b8acd53da3dcd692c0bbf979deae4ec18af04 Mon Sep 17 00:00:00 2001
From: grandizzy <38490174+grandizzy@users.noreply.github.com>
Date: Fri, 13 Sep 2024 15:40:24 +0300
Subject: [PATCH] feat(cheatcodes): docs for newly added cheatcodes (#1291)

---
 src/SUMMARY.md                          |  4 ++
 src/cheatcodes/README.md                | 17 ++++++
 src/cheatcodes/assume-no-revert.md      | 33 +++++++++++
 src/cheatcodes/copy-storage.md          | 41 ++++++++++++++
 src/cheatcodes/environment.md           |  1 +
 src/cheatcodes/fuzzer.md                |  1 +
 src/cheatcodes/mock-function.md         | 73 +++++++++++++++++++++++++
 src/cheatcodes/set-arbitrary-storage.md | 46 ++++++++++++++++
 src/cheatcodes/utilities.md             |  2 +
 9 files changed, 218 insertions(+)
 create mode 100644 src/cheatcodes/assume-no-revert.md
 create mode 100644 src/cheatcodes/copy-storage.md
 create mode 100644 src/cheatcodes/mock-function.md
 create mode 100644 src/cheatcodes/set-arbitrary-storage.md

diff --git a/src/SUMMARY.md b/src/SUMMARY.md
index dc9b986b9..64dc61512 100644
--- a/src/SUMMARY.md
+++ b/src/SUMMARY.md
@@ -415,6 +415,7 @@
       - [`getNonce`](./cheatcodes/get-nonce.md)
       - [`mockCall`](./cheatcodes/mock-call.md)
       - [`mockCallRevert`](./cheatcodes/mock-call-revert.md)
+      - [`mockFunction`](./cheatcodes/mock-function.md)
       - [`clearMockedCalls`](./cheatcodes/clear-mocked-calls.md)
       - [`coinbase`](./cheatcodes/coinbase.md)
       - [`broadcast`](./cheatcodes/broadcast.md)
@@ -432,6 +433,7 @@
       - [`expectCall`](./cheatcodes/expect-call.md)
     - [Fuzzer](./cheatcodes/fuzzer.md)
       - [`assume`](./cheatcodes/assume.md)
+      - [`assumeNoRevert`](./cheatcodes/assume-no-revert.md)
     - [Forking](./cheatcodes/forking.md)
       - [`createFork`](./cheatcodes/create-fork.md)
       - [`selectFork`](./cheatcodes/select-fork.md)
@@ -487,6 +489,8 @@
       - [`toString`](./cheatcodes/to-string.md)
       - [`breakpoint`](./cheatcodes/breakpoint.md)
       - [`createWallet`](./cheatcodes/create-wallet.md)
+      - [`copyStorage`](./cheatcodes/copy-storage.md)
+      - [`setArbitraryStorage`](./cheatcodes/set-arbitrary-storage.md)
     - [Snapshots](./cheatcodes/snapshots.md)
     - [RPC](./cheatcodes/rpc.md)
     - [Files](./cheatcodes/fs.md)
diff --git a/src/cheatcodes/README.md b/src/cheatcodes/README.md
index 663e537d5..342100d11 100644
--- a/src/cheatcodes/README.md
+++ b/src/cheatcodes/README.md
@@ -314,6 +314,14 @@ interface CheatCodes {
     // function will be mocked.
     function mockCallRevert(address where, bytes calldata data, bytes calldata retdata) external;
 
+    /// Whenever a call is made to `callee` with calldata `data`, this cheatcode instead calls
+    /// `target` with the same calldata. This functionality is similar to a delegate call made to
+    /// `target` contract from `callee`.
+    /// Can be used to substitute a call to a function with another implementation that captures
+    /// the primary logic of the original function but is easier to reason about.
+    /// If calldata is not a strict match then partial match by selector is attempted.
+    function mockFunction(address callee, address target, bytes calldata data) external;
+
     // Clears all mocked and reverted mocked calls
     function clearMockedCalls() external;
 
@@ -339,6 +347,9 @@ interface CheatCodes {
     // When fuzzing, generate new inputs if conditional not met
     function assume(bool) external;
 
+    /// Discard this run's fuzz inputs and generate new ones if next call reverted.
+    function assumeNoRevert() external;
+
     // Set block.coinbase (who)
     function coinbase(address) external;
 
@@ -459,5 +470,11 @@ interface CheatCodes {
     function rpcUrl(string calldata) external returns (string memory);
     /// Returns all rpc urls and their aliases `[alias, url][]`
     function rpcUrls() external returns (string[2][] memory);
+
+    /// Utility cheatcode to copy storage of `from` contract to another `to` contract.
+    function copyStorage(address from, address to) external;
+
+    /// Utility cheatcode to set arbitrary storage for given target address.
+    function setArbitraryStorage(address target) external;
 }
 ```
diff --git a/src/cheatcodes/assume-no-revert.md b/src/cheatcodes/assume-no-revert.md
new file mode 100644
index 000000000..0ac4febe9
--- /dev/null
+++ b/src/cheatcodes/assume-no-revert.md
@@ -0,0 +1,33 @@
+## `assumeNoRevert`
+
+### Signature
+
+```solidity
+function assumeNoRevert() external;
+```
+
+### Description
+
+The fuzzer will discard the current fuzz inputs and start a new fuzz run if next call reverted.
+
+The test may fail if you hit the max number of rejects.
+
+You can configure the rejection thresholds by setting [`fuzz.max_test_rejects`][max-test-rejects] in your `foundry.toml` file.
+
+### Examples
+
+For a function that requires an amount in certain range:
+```solidity
+function doSomething(uint256 amount) public {
+    require(amount > 100 ether && amount < 1_000 ether);
+}
+```
+reverts are discarded, resulting in test pass (or fail if max number of rejects hit):
+```solidity
+function testSomething(uint256 amount) public {
+    vm.assumeNoRevert();
+    target.doSomething(amount);
+    // [PASS]
+}
+```
+
diff --git a/src/cheatcodes/copy-storage.md b/src/cheatcodes/copy-storage.md
new file mode 100644
index 000000000..c0184ab2c
--- /dev/null
+++ b/src/cheatcodes/copy-storage.md
@@ -0,0 +1,41 @@
+## `copyStorage`
+
+### Signature
+
+```solidity
+function copyStorage(address from, address to) external;
+```
+
+### Description
+
+Utility cheatcode to copy storage of `from` contract to another `to` contract.
+Cheatcode is not allowed if the target address has arbitrary storage set.
+
+### Examples
+
+Given a contract
+```solidity
+contract Counter {
+    uint256 public count;
+
+    function setCount(uint256 x) public {
+        count = x;
+    }
+}
+```
+using `copyStorage` cheatcode copies the storage set on an instance to another address:
+```solidity
+function testCopyStorage() public {
+    Counter original = new Counter();
+    original.setCount(1000);
+    Counter copy = new Counter();
+    copy.setCount(1);
+    // Check initial count on copy.
+    assertEq(copy.count(), 1);
+
+    vm.copyStorage(address(original), address(copy));
+    // Value is copied from first contract to copy.
+    assertEq(copy.count(), 1000);
+}
+```
+
diff --git a/src/cheatcodes/environment.md b/src/cheatcodes/environment.md
index 7d044e2b2..a5caa4ad7 100644
--- a/src/cheatcodes/environment.md
+++ b/src/cheatcodes/environment.md
@@ -22,6 +22,7 @@
 - [`getNonce`](./get-nonce.md)
 - [`mockCall`](./mock-call.md)
 - [`mockCallRevert`](./mock-call-revert.md)
+- [`mockFunction`](./mock-function.md)
 - [`clearMockedCalls`](./clear-mocked-calls.md)
 - [`coinbase`](./coinbase.md)
 - [`broadcast`](./broadcast.md)
diff --git a/src/cheatcodes/fuzzer.md b/src/cheatcodes/fuzzer.md
index 799923c40..e56af83fc 100644
--- a/src/cheatcodes/fuzzer.md
+++ b/src/cheatcodes/fuzzer.md
@@ -1,3 +1,4 @@
 ## Fuzzer
 
 - [`assume`](./assume.md)
+- [`assumeNoRevert`](./assume-no-revert.md)
diff --git a/src/cheatcodes/mock-function.md b/src/cheatcodes/mock-function.md
new file mode 100644
index 000000000..034433fcd
--- /dev/null
+++ b/src/cheatcodes/mock-function.md
@@ -0,0 +1,73 @@
+## `mockFunction`
+
+### Signature
+
+```solidity
+function mockFunction(address callee, address target, bytes calldata data) external;
+```
+
+### Description
+
+Executes calls to an address `callee` with bytecode of address `target` if the call data either strictly or loosely matches `data`.
+
+When a call is made to `callee` the call data is first checked to see if it matches in its entirety with `data`.
+If not, the call data is checked to see if there is a partial match on function selector.
+
+If a match is found, then call is executed using the bytecode of `target` address.
+
+> ℹ️ **Isolated tests**
+>
+> This cheatcode does not currently work if using isolated test mode.
+
+### Examples
+
+For two contracts (with same storage layout):
+```solidity
+contract Counter {
+    uint256 public a;
+
+    function count(uint256 x) public {
+        a = 321 + x;
+    }
+}
+
+contract ModelCounter {
+    uint256 public a;
+
+    function count(uint256 x) public {
+        a = 123 + x;
+    }
+}
+```
+Mocking an exact call to `count` function:
+
+```solidity
+function testMockFunction() public {
+    vm.mockFunction(
+        address(counter),
+        address(model),
+        abi.encodeWithSelector(Counter.count.selector, 456)
+    );
+    counter.count(456);
+    assertEq(counter.a(), 123 + 456);
+    counter.count(567);
+    assertEq(counter.a(), 321 + 567);
+}
+```
+
+Mocking all calls to `count` function:
+
+```solidity
+function testMockCall() public {
+    vm.mockFunction(
+        address(counter),
+        address(model),
+        abi.encodeWithSelector(Counter.count.selector)
+    );
+    counter.count(678);
+    assertEq(counter.a(), 123 + 678);
+    counter.count(789);
+    assertEq(counter.a(), 123 + 789);
+}
+```
+
diff --git a/src/cheatcodes/set-arbitrary-storage.md b/src/cheatcodes/set-arbitrary-storage.md
new file mode 100644
index 000000000..73b0ad8b1
--- /dev/null
+++ b/src/cheatcodes/set-arbitrary-storage.md
@@ -0,0 +1,46 @@
+## `setArbitraryStorage`
+
+### Signature
+
+```solidity
+function setArbitraryStorage(address target) external;
+```
+
+### Description
+
+Utility cheatcode to make the storage of the given address fully symbolic.
+Any subsequent `SLOAD` to target storage reads an arbitrary value which is memorized and returned if the same slot is loaded again.
+If the storage slot is explicitly written (before or after first load), then the written value is returned.
+
+### Examples
+
+For a contract with following storage layout:
+```solidity
+contract Counter {
+    address[] public owners;
+
+    function getOwner(uint256 pos) public view returns (address) {
+        return owners[pos];
+    }
+
+    function setOwner(uint256 pos, address owner) public {
+        owners[pos] = owner;
+    }
+}
+```
+using `setArbitraryStorage` cheatcode ensures that arbitrary values are returned: 
+```solidity
+contract ArbitraryStorageTest is Test {
+    function testArbitraryStorage() public {
+        Counter counter = new Counter();
+        vm.setArbitraryStorage(address(counter));
+        // Next call would fail with array out of bounds without arbitrary storage
+        address owner = counter.getOwner(55);
+        // Subsequent calls to same slot returns same value
+        assertEq(counter.getOwner(55), owner);
+        // The new value is returned if explicitly written
+        counter.setOwner(55, address(111));
+        assertEq(counter.getOwner(55), address(111));
+    }
+}
+```
diff --git a/src/cheatcodes/utilities.md b/src/cheatcodes/utilities.md
index 75c3304e5..da8e16a41 100644
--- a/src/cheatcodes/utilities.md
+++ b/src/cheatcodes/utilities.md
@@ -9,3 +9,5 @@
 - [`toString`](./to-string.md)
 - [`breakpoint`](./breakpoint.md)
 - [`createWallet`](./create-wallet.md)
+- [`copyStorage`](./copy-storage.md)
+- [`setArbitraryStorage`](./set-arbitrary-storage.md)