feat: loadAllocs cheatcode (Rebased) (#6207)

* chore: add test & fixture

* feat: add cheatcode impl

* chore: remove useless import

* chore: fmt

* chore: switch to foundry fs, make try_for_each loop

* chore: support genesis format

* chore: clippy

crates/cheatcodes/defs/src/vm.rs

@@ -148,6 +148,10 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Safe)]
     function load(address target, bytes32 slot) external view returns (bytes32 data);
 
+    /// Load a genesis JSON file's `allocs` into the in-memory revm state.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function loadAllocs(string calldata pathToAllocsJson) external;
+
     /// Signs data.
     #[cheatcode(group = Evm, safety = Safe)]
     function sign(uint256 privateKey, bytes32 digest) external pure returns (uint8 v, bytes32 r, bytes32 s);

crates/cheatcodes/src/evm.rs

@@ -3,14 +3,16 @@
 use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Result, Vm::*};
 use alloy_primitives::{Address, Bytes, U256};
 use alloy_sol_types::SolValue;
+use ethers_core::utils::{Genesis, GenesisAccount};
 use ethers_signers::Signer;
+use foundry_common::fs::read_json_file;
 use foundry_evm_core::backend::DatabaseExt;
 use foundry_utils::types::ToAlloy;
 use revm::{
     primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY},
     EVMData,
 };
-use std::collections::HashMap;
+use std::{collections::HashMap, path::Path};
 
 mod fork;
 pub(crate) mod mapping;
@@ -62,6 +64,29 @@ impl Cheatcode for loadCall {
     }
 }
 
+impl Cheatcode for loadAllocsCall {
+    fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { pathToAllocsJson } = self;
+
+        let path = Path::new(pathToAllocsJson);
+        ensure!(path.exists(), "allocs file does not exist: {pathToAllocsJson}");
+
+        // Let's first assume we're reading a genesis.json file.
+        let allocs: HashMap<Address, GenesisAccount> = match read_json_file::<Genesis>(path) {
+            Ok(genesis) => genesis.alloc.into_iter().map(|(k, g)| (k.to_alloy(), g)).collect(),
+            // If that fails, let's try reading a file with just the genesis accounts.
+            Err(_) => read_json_file(path)?,
+        };
+
+        // Then, load the allocs into the database.
+        ccx.data
+            .db
+            .load_allocs(&allocs, &mut ccx.data.journaled_state)
+            .map(|_| Vec::default())
+            .map_err(|e| fmt_err!("failed to load allocs: {e}"))
+    }
+}
+
 impl Cheatcode for sign_0Call {
     fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { privateKey, digest } = self;

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

@@ -3,7 +3,7 @@
 pub use foundry_abi::{
     console::{self, ConsoleEvents, CONSOLE_ABI},
     hardhat_console::{self, HardhatConsoleCalls, HARDHATCONSOLE_ABI as HARDHAT_CONSOLE_ABI},
-    hevm::{self, HEVMCalls, HEVM_ABI},
+    hevm::HEVM_ABI,
 };
 use once_cell::sync::Lazy;
 use std::collections::HashMap;

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

@@ -7,12 +7,13 @@ use crate::{
     fork::{CreateFork, ForkId},
 };
 use alloy_primitives::{Address, B256, U256};
+use ethers::utils::GenesisAccount;
 use revm::{
     db::DatabaseRef,
     primitives::{AccountInfo, Bytecode, Env, ResultAndState},
     Database, Inspector, JournaledState,
 };
-use std::borrow::Cow;
+use std::{borrow::Cow, collections::HashMap};
 
 /// A wrapper around `Backend` that ensures only `revm::DatabaseRef` functions are called.
 ///
@@ -198,6 +199,14 @@ impl<'a> DatabaseExt for FuzzBackendWrapper<'a> {
         self.backend.diagnose_revert(callee, journaled_state)
     }
 
+    fn load_allocs(
+        &mut self,
+        allocs: &HashMap<Address, GenesisAccount>,
+        journaled_state: &mut JournaledState,
+    ) -> Result<(), DatabaseError> {
+        self.backend_mut(&Env::default()).load_allocs(allocs, journaled_state)
+    }
+
     fn is_persistent(&self, acc: &Address) -> bool {
         self.backend.is_persistent(acc)
     }

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

@@ -10,6 +10,7 @@ use alloy_primitives::{b256, keccak256, Address, B256, U256, U64};
 use ethers::{
     prelude::Block,
     types::{BlockNumber, Transaction},
+    utils::GenesisAccount,
 };
 use foundry_common::{is_known_system_sender, SYSTEM_TRANSACTION_TYPE};
 use foundry_utils::types::{ToAlloy, ToEthers};
@@ -19,7 +20,7 @@ use revm::{
     precompile::{Precompiles, SpecId},
     primitives::{
         Account, AccountInfo, Bytecode, CreateScheme, Env, HashMap as Map, Log, ResultAndState,
-        TransactTo, KECCAK_EMPTY,
+        StorageSlot, TransactTo, KECCAK_EMPTY,
     },
     Database, DatabaseCommit, Inspector, JournaledState, EVM,
 };
@@ -248,6 +249,15 @@ pub trait DatabaseExt: Database<Error = DatabaseError> {
         journaled_state: &JournaledState,
     ) -> Option<RevertDiagnostic>;
 
+    /// Loads the account allocs from the given `allocs` map into the passed [JournaledState].
+    ///
+    /// Returns [Ok] if all accounts were successfully inserted into the journal, [Err] otherwise.
+    fn load_allocs(
+        &mut self,
+        allocs: &HashMap<Address, GenesisAccount>,
+        journaled_state: &mut JournaledState,
+    ) -> Result<(), DatabaseError>;
+
     /// Returns true if the given account is currently marked as persistent.
     fn is_persistent(&self, acc: &Address) -> bool;
 
@@ -1263,6 +1273,55 @@ impl DatabaseExt for Backend {
         None
     }
 
+    /// Loads the account allocs from the given `allocs` map into the passed [JournaledState].
+    ///
+    /// Returns [Ok] if all accounts were successfully inserted into the journal, [Err] otherwise.
+    fn load_allocs(
+        &mut self,
+        allocs: &HashMap<Address, GenesisAccount>,
+        journaled_state: &mut JournaledState,
+    ) -> Result<(), DatabaseError> {
+        // Loop through all of the allocs defined in the map and commit them to the journal.
+        for (addr, acc) in allocs.iter() {
+            // Fetch the account from the journaled state. Will create a new account if it does
+            // not already exist.
+            let (state_acc, _) = journaled_state.load_account(*addr, self)?;
+
+            // Set the account's bytecode and code hash, if the `bytecode` field is present.
+            if let Some(bytecode) = acc.code.as_ref() {
+                state_acc.info.code_hash = keccak256(bytecode);
+                let bytecode = Bytecode::new_raw(bytecode.0.clone().into());
+                state_acc.info.code = Some(bytecode);
+            }
+
+            // Set the account's storage, if the `storage` field is present.
+            if let Some(storage) = acc.storage.as_ref() {
+                state_acc.storage = storage
+                    .iter()
+                    .map(|(slot, value)| {
+                        let slot = U256::from_be_bytes(slot.0);
+                        (
+                            slot,
+                            StorageSlot::new_changed(
+                                state_acc
+                                    .storage
+                                    .get(&slot)
+                                    .map(|s| s.present_value)
+                                    .unwrap_or_default(),
+                                U256::from_be_bytes(value.0),
+                            ),
+                        )
+                    })
+                    .collect();
+            }
+            // Set the account's nonce and balance.
+            state_acc.info.nonce = acc.nonce.unwrap_or_default();
+            state_acc.info.balance = acc.balance.to_alloy();
+        }
+
+        Ok(())
+    }
+
     fn is_persistent(&self, acc: &Address) -> bool {
         self.inner.persistent_accounts.contains(acc)
     }

testdata/cheats/loadAllocs.t.sol

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: Unlicense
+pragma solidity 0.8.18;
+
+import "ds-test/test.sol";
+import "./Vm.sol";
+
+contract LoadAllocsTest is DSTest {
+    Vm constant vm = Vm(HEVM_ADDRESS);
+    string allocsPath;
+    address constant ALLOCD = address(0x420);
+    address constant ALLOCD_B = address(0x421);
+
+    function setUp() public {
+        allocsPath = string.concat(vm.projectRoot(), "/fixtures/Json/test_allocs.json");
+    }
+
+    function testLoadAllocsStatic() public {
+        vm.loadAllocs(allocsPath);
+
+        // Balance should be `0xabcd`
+        assertEq(ALLOCD.balance, 0xabcd);
+
+        // Code should be a simple store / return, returning `0x42`
+        (bool success, bytes memory rd) = ALLOCD.staticcall("");
+        assertTrue(success);
+        uint256 ret = abi.decode(rd, (uint256));
+        assertEq(ret, 0x42);
+
+        // Storage should have been set in slot 0x1, equal to `0xbeef`
+        assertEq(uint256(vm.load(ALLOCD, bytes32(uint256(0x10 << 248)))), 0xbeef);
+    }
+
+    function testLoadAllocsOverride() public {
+        // Populate the alloc'd account's code.
+        vm.etch(ALLOCD, hex"FF");
+        assertEq(ALLOCD.code, hex"FF");
+
+        // Store something in the alloc'd storage slot.
+        bytes32 slot = bytes32(uint256(0x10 << 248));
+        vm.store(ALLOCD, slot, bytes32(uint256(0xBADC0DE)));
+        assertEq(uint256(vm.load(ALLOCD, slot)), 0xBADC0DE);
+
+        // Populate balance.
+        vm.deal(ALLOCD, 0x1234);
+        assertEq(ALLOCD.balance, 0x1234);
+
+        vm.loadAllocs(allocsPath);
+
+        // Info should have changed.
+        assertTrue(keccak256(ALLOCD.code) != keccak256(hex"FF"));
+        assertEq(uint256(vm.load(ALLOCD, slot)), 0xbeef);
+        assertEq(ALLOCD.balance, 0xabcd);
+    }
+
+    function testLoadAllocsPartialOverride() public {
+        // Populate the alloc'd account's code.
+        vm.etch(ALLOCD_B, hex"FF");
+        assertEq(ALLOCD_B.code, hex"FF");
+
+        // Populate balance.
+        vm.deal(ALLOCD_B, 0x1234);
+        assertEq(ALLOCD_B.balance, 0x1234);
+
+        vm.loadAllocs(allocsPath);
+
+        assertEq(ALLOCD_B.code, hex"FF");
+        assertEq(ALLOCD_B.balance, 0);
+    }
+}

testdata/fixtures/Json/test_allocs.json

@@ -0,0 +1,15 @@
+{
+  "0x0000000000000000000000000000000000000420": {
+    "balance": "0xabcd",
+    "code": "0x604260005260206000F3",
+    "storage": {
+      "0x1000000000000000000000000000000000000000000000000000000000000000": "0x000000000000000000000000000000000000000000000000000000000000beef"
+    }
+  },
+  "0x0000000000000000000000000000000000000421": {
+    "balance": "0x0",
+    "storage": {
+      "0x1000000000000000000000000000000000000000000000000000000000000000": "0x000000000000000000000000000000000000000000000000000000000000beef"
+    }
+  }
+}