Add basic forking cheatcodes

CHANGELOG.md

@@ -28,6 +28,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - New concrete fuzzer that can be controlled via `--num-cex-fuzz`
 - Partial support for dynamic jumps when the jump destination can be computed
   given already available information
+- Added two forking cheatcodes: `createFork` and `selectFork`
 
 ## Fixed
 

doc/src/controlling-the-unit-testing-environment.md

@@ -38,3 +38,9 @@ These can be accessed by calling into a contract (typically called `Vm`) at addr
 
 - `function prank(address sender) public`
   Sets `msg.sender` to the specified `sender` for the next call.
+
+- `function createFork(string calldata urlOrAlias) external returns (uint256)`
+  Creates a new fork with the given endpoint and the _latest_ block and returns the identifier of the fork.
+
+- `function selectFork(uint256 forkId) external`
+  Takes a fork identifier created by `createFork` and sets the corresponding forked state as active.

src/EVM.hs

@@ -122,16 +122,7 @@ makeVm o = do
       }
     , logs = []
     , traces = Zipper.fromForest []
-    , block = Block
-      { coinbase = o.coinbase
-      , timestamp = o.timestamp
-      , number = o.number
-      , prevRandao = o.prevRandao
-      , maxCodeSize = o.maxCodeSize
-      , gaslimit = o.blockGaslimit
-      , baseFee = o.baseFee
-      , schedule = o.schedule
-      }
+    , block = block
     , state = FrameState
       { pc = 0
       , stack = mempty
@@ -147,13 +138,8 @@ makeVm o = do
       , returndata = mempty
       , static = False
       }
-    , env = Env
-      { chainId = o.chainId
-      , contracts = Map.fromList ((o.address,o.contract):o.otherContracts)
-      , freshAddresses = 0
-      , freshGasVals = 0
-      }
-    , cache = Cache mempty mempty
+    , env = env
+    , cache = cache
     , burned = initialGas
     , constraints = snd o.calldata
     , iterations = mempty
@@ -162,7 +148,27 @@ makeVm o = do
       , overrideCaller = Nothing
       , baseState = o.baseState
       }
+    , forks = Seq.singleton (ForkState env block cache "")
+    , currentFork = 0
     }
+    where
+    env = Env
+      { chainId = o.chainId
+      , contracts = Map.fromList ((o.address,o.contract):o.otherContracts)
+      , freshAddresses = 0
+      , freshGasVals = 0
+      }
+    block = Block
+      { coinbase = o.coinbase
+      , timestamp = o.timestamp
+      , number = o.number
+      , prevRandao = o.prevRandao
+      , maxCodeSize = o.maxCodeSize
+      , gaslimit = o.blockGaslimit
+      , baseFee = o.baseFee
+      , schedule = o.schedule
+      }
+    cache = Cache mempty mempty
 
 -- | Initialize an abstract contract with unknown code
 unknownContract :: Expr EAddr -> Contract
@@ -1650,8 +1656,50 @@ cheatActions =
           [addr]  -> case wordToAddr addr of
             Just a -> assign (#config % #overrideCaller) (Just a)
             Nothing -> vmError (BadCheatCode sig)
-          _ -> vmError (BadCheatCode sig)
+          _ -> vmError (BadCheatCode sig),
 
+      action "createFork(string)" $
+        \sig outOffset _ input -> case decodeBuf [AbiStringType] input of
+          CAbi valsArr -> case valsArr of
+            [AbiString bytes] -> do
+              forkId <- length <$> gets (.forks)
+              let urlOrAlias = Char8.unpack bytes
+              modify' $ \vm -> vm { forks = vm.forks Seq.|> ForkState vm.env vm.block vm.cache urlOrAlias }
+              let encoded = encodeAbiValue $ AbiUInt 256 (fromIntegral forkId)
+              assign (#state % #returndata) (ConcreteBuf encoded)
+              copyBytesToMemory (ConcreteBuf encoded) (Lit . unsafeInto . BS.length $ encoded) (Lit 0) outOffset
+            _ -> vmError (BadCheatCode sig)
+          _ -> vmError (BadCheatCode sig),
+
+      action "selectFork(uint256)" $
+        \sig _ _ input -> case decodeStaticArgs 0 1 input of
+          [forkId] ->
+            forceConcrete forkId "forkId must be concrete" $ \(fromIntegral -> forkId') -> do
+              saved <- Seq.lookup forkId' <$> gets (.forks)
+              case saved of
+                Just forkState -> do
+                  vm <- get
+                  let contractAddr = vm.state.contract
+                  let callerAddr = vm.state.caller
+                  fetchAccount contractAddr $ \contractAcct -> fetchAccount callerAddr $ \callerAcct -> do
+                    let
+                      -- the current contract is persited across forks
+                      newContracts = Map.insert callerAddr callerAcct $
+                                     Map.insert contractAddr contractAcct forkState.env.contracts
+                      newEnv = (forkState.env :: Env) { contracts = newContracts }
+
+                    when (vm.currentFork /= forkId') $ do
+                      modify' $ \vm' -> vm'
+                        { env = newEnv
+                        , block = forkState.block
+                        , forks = Seq.adjust' (\state -> (state :: ForkState)
+                            { env = vm.env, block = vm.block, cache = vm.cache }
+                          ) vm.currentFork  vm.forks
+                        , currentFork = forkId'
+                        }
+                Nothing ->
+                  vmError (NonexistentFork forkId')
+          _ -> vmError (BadCheatCode sig)
     ]
   where
     action s f = (abiKeccak s, f (abiKeccak s))

src/EVM/Format.hs

@@ -415,6 +415,7 @@ prettyError = \case
   ReturnDataOutOfBounds -> "Return data out of bounds"
   NonceOverflow -> "Nonce overflow"
   BadCheatCode a -> "Bad cheat code: sig: " <> show a
+  NonexistentFork a -> "Fork ID does not exist: " <> show a
 
 prettyvmresult :: Expr End -> String
 prettyvmresult (Failure _ _ (Revert (ConcreteBuf ""))) = "Revert"

src/EVM/Types.hs

@@ -37,6 +37,7 @@ import Data.Map (Map)
 import Data.Map qualified as Map
 import Data.Maybe (fromMaybe)
 import Data.Set (Set)
+import Data.Sequence (Seq)
 import Data.Sequence qualified as Seq
 import Data.Serialize qualified as Cereal
 import Data.Text qualified as T
@@ -536,6 +537,7 @@ data EvmError
   | ReturnDataOutOfBounds
   | NonceOverflow
   | BadCheatCode FunctionSelector
+  | NonexistentFork Int
   deriving (Show, Eq, Ord)
 
 -- | Sometimes we can only partially execute a given program
@@ -620,9 +622,19 @@ data VM (t :: VMType) s = VM
   -- ^ how many times we've visited a loc, and what the contents of the stack were when we were there last
   , constraints    :: [Prop]
   , config         :: RuntimeConfig
+  , forks          :: Seq ForkState
+  , currentFork    :: Int
   }
   deriving (Generic)
 
+data ForkState = ForkState
+  { env :: Env
+  , block :: Block
+  , cache :: Cache
+  , urlOrAlias :: String
+  }
+  deriving (Show, Generic)
+
 deriving instance Show (VM Symbolic s)
 deriving instance Show (VM Concrete s)
 

test/contracts/pass/cheatCodesFork.sol

@@ -0,0 +1,88 @@
+pragma experimental ABIEncoderV2;
+
+import "ds-test/test.sol";
+
+interface Hevm {
+    function warp(uint256) external;
+    function roll(uint256) external;
+    function load(address,bytes32) external returns (bytes32);
+    function store(address,bytes32,bytes32) external;
+    function sign(uint256,bytes32) external returns (uint8,bytes32,bytes32);
+    function addr(uint256) external returns (address);
+    function ffi(string[] calldata) external returns (bytes memory);
+    function prank(address) external;
+    function deal(address,uint256) external;
+    function createFork(string calldata urlOrAlias) external returns (uint256);
+    function selectFork(uint256 forkId) external;
+}
+
+/// @dev This contract's state should depend on which fork we are on
+contract TestState {
+    uint256 public state;
+    function setState(uint256 _state) external {
+        state = _state;
+    }
+}
+
+/// @dev This contract's state should be persistent across forks, because it's the contract calling `selectFork`
+contract CheatCodesForkDeployee is DSTest {
+    Hevm hevm = Hevm(HEVM_ADDRESS);
+    address stateContract;
+    uint256 forkId1;
+    uint256 forkId2;
+    uint256 persistentState;
+
+    constructor() {
+        stateContract = address(new TestState());
+        forkId1 = hevm.createFork("foo"); // If the default fork's id is 0, then this would be 1
+        forkId2 = hevm.createFork("bar"); // and this would be 2
+        persistentState = 0;
+    }
+
+    function deployee_prove_ForkedState() external {
+        hevm.selectFork(0);
+        persistentState = 1;                            // Make sure this contract maintains its own state across fork
+        hevm.selectFork(forkId1);                       // Fork 1
+        assert(TestState(stateContract).state() == 0);  //      Check initial external state
+        assert(persistentState == 1);                   //      Check persistent state
+        persistentState = 2;                            //      Set persistent state
+        TestState(stateContract).setState(1);           //      Set unique external state
+        hevm.roll(12345678);                            //      Set unique block number
+        hevm.selectFork(forkId2);                       // Fork 2
+        assert(TestState(stateContract).state() == 0);  //      Check initial external state
+        assert(persistentState == 2);                   //      Check persistent state
+        persistentState = 3;                            //      Set persistent state
+        TestState(stateContract).setState(2);           //      Set unique external state
+        hevm.roll(23456789);                            //      Set unique block number
+        hevm.selectFork(forkId1);                       // Fork 1
+        assert(block.number == 12345678);               //      Check unique block number
+        assert(TestState(stateContract).state() == 1);  //      Check unique external state
+        assert(persistentState == 3);                   //      Check persistent state
+        persistentState = 4;                            //      Set persistent state
+        TestState(stateContract).setState(0);           //      Set initial external state
+        hevm.selectFork(forkId2);                       // Fork 2
+        assert(block.number == 23456789);               //      Check unique block number
+        assert(TestState(stateContract).state() == 2);  //      Check unique external state
+        assert(persistentState == 4);                   //      Check persistent state
+        persistentState = 5;                            //      Set persistent state
+        TestState(stateContract).setState(0);           //      Set initial external state
+        hevm.selectFork(forkId1);                       // Fork 1
+        hevm.deal(address(this), 10);                   //      Get some eth
+        payable(msg.sender).send(10);                   //      Send eth to msg.sender
+        uint256 senderBalance = msg.sender.balance;     //      Record msg.sender's balance
+        hevm.selectFork(forkId2);                       // Fork 2
+        assert(msg.sender.balance == senderBalance);    //      Check msg.sender's balance
+        hevm.selectFork(0);                             // Default fork
+        assert(persistentState == 5);                   //      Check persistent state
+    }
+}
+
+/// @dev This contract's state should be persistent across forks, because it's the `msg.sender` when running `deployee_prove_ForkedState`.
+///  We need this "deployer/deployee" architecture so that `msg.sender` will be concrete when running `deployee_prove_ForkedState`.
+///  If we were to only use the `CheatCodesForkDeployee` contract, the `msg.sender` would be abstract.
+contract CheatCodesFork is DSTest {
+  CheatCodesForkDeployee testContract = new CheatCodesForkDeployee();
+  function prove_ForkedState() external {
+    testContract.deployee_prove_ForkedState();
+  }
+}

test/test.hs

@@ -1255,6 +1255,9 @@ tests = testGroup "hevm"
     , test "Cheat-Codes-Pass" $ do
         let testFile = "test/contracts/pass/cheatCodes.sol"
         runSolidityTest testFile ".*" >>= assertEqualM "test result" True
+    , test "Cheat-Codes-Fork-Pass" $ do
+        let testFile = "test/contracts/pass/cheatCodesFork.sol"
+        runSolidityTest testFile ".*" >>= assertEqualM "test result" True
     , test "Unwind" $ do
         let testFile = "test/contracts/pass/unwind.sol"
         runSolidityTest testFile ".*" >>= assertEqualM "test result" True