fix(forge): ensure selected fork contains init state for persisted accounts (#10301)

Author: grandizzy

Diff for: crates/evm/core/src/backend/mod.rs

@@ -1102,6 +1102,14 @@ impl DatabaseExt for Backend {
             // update the shared state and track
             let mut fork = self.inner.take_fork(idx);
 
+            // Make sure all persistent accounts on the newly selected fork starts from the init
+            // state (from setup).
+            for addr in &self.inner.persistent_accounts {
+                if let Some(account) = self.fork_init_journaled_state.state.get(addr) {
+                    fork.journaled_state.state.insert(*addr, account.clone());
+                }
+            }
+
             // since all forks handle their state separately, the depth can drift
             // this is a handover where the target fork starts at the same depth where it was
             // selected. This ensures that there are no gaps in depth which would

Diff for: crates/forge/tests/cli/test_cmd.rs

@@ -3521,3 +3521,82 @@ contract InterceptInitcodeTest is DSTest {
     .unwrap();
     cmd.args(["test", "-vvvvv"]).assert_success();
 });
+
+// <https://github.com/foundry-rs/foundry/issues/10296>
+forgetest_init!(should_preserve_fork_state_setup, |prj, cmd| {
+    prj.wipe_contracts();
+    prj.add_test(
+        "Counter.t.sol",
+        r#"
+import "forge-std/Test.sol";
+import {StdChains} from "forge-std/StdChains.sol";
+
+contract CounterTest is Test {
+    struct Domain {
+        StdChains.Chain chain;
+        uint256 forkId;
+    }
+
+    struct Bridge {
+        Domain source;
+        Domain destination;
+        uint256 someVal;
+    }
+
+    struct SomeStruct {
+        Domain domain;
+        Bridge[] bridges;
+    }
+
+    mapping(uint256 => SomeStruct) internal data;
+
+    function setUp() public {
+        StdChains.Chain memory chain1 = getChain("mainnet");
+        StdChains.Chain memory chain2 = getChain("base");
+        Domain memory domain1 = Domain(chain1, vm.createFork(chain1.rpcUrl, 22253716));
+        Domain memory domain2 = Domain(chain2, vm.createFork(chain2.rpcUrl, 28839981));
+        data[1].domain = domain1;
+        data[2].domain = domain2;
+
+        vm.selectFork(domain1.forkId);
+
+        data[2].bridges.push(Bridge(domain1, domain2, 123));
+        vm.selectFork(data[2].domain.forkId);
+        vm.selectFork(data[1].domain.forkId);
+        data[2].bridges.push(Bridge(domain1, domain2, 456));
+
+        assertEq(data[2].bridges.length, 2);
+    }
+
+    function test_assert_storage() public {
+        vm.selectFork(data[2].domain.forkId);
+        assertEq(data[2].bridges.length, 2);
+    }
+
+    function test_modify_and_storage() public {
+        data[3].domain = Domain(getChain("base"), vm.createFork(getChain("base").rpcUrl, 28839981));
+        data[3].bridges.push(Bridge(data[1].domain, data[2].domain, 123));
+        data[3].bridges.push(Bridge(data[1].domain, data[2].domain, 456));
+
+        vm.selectFork(data[2].domain.forkId);
+        assertEq(data[3].bridges.length, 2);
+    }
+}
+    "#,
+    )
+    .unwrap();
+
+    cmd.args(["test", "--mc", "CounterTest"]).assert_success().stdout_eq(str![[r#"
+[COMPILING_FILES] with [SOLC_VERSION]
+[SOLC_VERSION] [ELAPSED]
+Compiler run successful!
+
+Ran 2 tests for test/Counter.t.sol:CounterTest
+[PASS] test_assert_storage() ([GAS])
+[PASS] test_modify_and_storage() ([GAS])
+Suite result: ok. 2 passed; 0 failed; 0 skipped; [ELAPSED]
+
+Ran 1 test suite [ELAPSED]: 2 tests passed, 0 failed, 0 skipped (2 total tests)
+
+"#]]);
+});