From e04104ab93e771441eab03fb76eda1402cb5927b Mon Sep 17 00:00:00 2001
From: grandizzy <38490174+grandizzy@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:53:08 +0300
Subject: [PATCH] feat: add `setArbitraryStorage`, `copyStorage`,
 `mockFunction` and `assumeNoRevert` cheatcodes (#602)

---
 src/Vm.sol    | 17 +++++++++++++++++
 test/Vm.t.sol |  4 ++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/Vm.sol b/src/Vm.sol
index ec495aaa..25b02789 100644
--- a/src/Vm.sol
+++ b/src/Vm.sol
@@ -1437,6 +1437,9 @@ interface VmSafe {
     /// If the condition is false, discard this run's fuzz inputs and generate new ones.
     function assume(bool condition) external pure;
 
+    /// Discard this run's fuzz inputs and generate new ones if next call reverted.
+    function assumeNoRevert() external pure;
+
     /// Writes a breakpoint to jump to in the debugger.
     function breakpoint(string calldata char) external;
 
@@ -1546,6 +1549,9 @@ interface VmSafe {
     /// Compute the address a contract will be deployed at for a given deployer address and nonce.
     function computeCreateAddress(address deployer, uint256 nonce) external pure returns (address);
 
+    /// Utility cheatcode to copy storage of `from` contract to another `to` contract.
+    function copyStorage(address from, address to) external;
+
     /// Returns ENS namehash for provided string.
     function ensNamehash(string calldata name) external pure returns (bytes32);
 
@@ -1571,6 +1577,9 @@ interface VmSafe {
     /// Unpauses collection of call traces.
     function resumeTracing() external view;
 
+    /// Utility cheatcode to set arbitrary storage for given target address.
+    function setArbitraryStorage(address target) external;
+
     /// Encodes a `bytes` value to a base64url string.
     function toBase64URL(bytes calldata data) external pure returns (string memory);
 
@@ -1699,6 +1708,14 @@ interface Vm is VmSafe {
     /// Calldata match takes precedence over `msg.value` in case of ambiguity.
     function mockCall(address callee, uint256 msgValue, bytes calldata data, bytes calldata returnData) 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;
+
     /// Sets the *next* call's `msg.sender` to be the input address.
     function prank(address msgSender) external;
 
diff --git a/test/Vm.t.sol b/test/Vm.t.sol
index 91cfcebf..5e772a92 100644
--- a/test/Vm.t.sol
+++ b/test/Vm.t.sol
@@ -9,7 +9,7 @@ contract VmTest is Test {
     // inadvertently moved between Vm and VmSafe. This test must be updated each time a function is
     // added to or removed from Vm or VmSafe.
     function test_interfaceId() public pure {
-        assertEq(type(VmSafe).interfaceId, bytes4(0x9dd1a1c8), "VmSafe");
-        assertEq(type(Vm).interfaceId, bytes4(0x02edefa2), "Vm");
+        assertEq(type(VmSafe).interfaceId, bytes4(0x74d42398), "VmSafe");
+        assertEq(type(Vm).interfaceId, bytes4(0xaf15a283), "Vm");
     }
 }