op-chain-ops: state forking, Go script forking, cheatcode access-control (#11919)

* op-chain-ops: state forking, Go script forking, cheatcode access-control

* add forking unit tests

* goimports

* linter

* fix state dump

* Add script test, hammer out bugs

* semgrep ignore scripts

* op-chain-ops: script state dump test

* fix merge error

* goimports

---------
Co-authored-by: Matthew Slipper <me@matthewslipper.com>

.semgrepignore

@@ -13,3 +13,6 @@ vendor/
 
 # Semgrep-action log folder
 .semgrep_logs/
+
+# Test contracts the scripts folder
+op-chain-ops/script/testdata/scripts/
\ No newline at end of file

op-chain-ops/script/addresses/addresses.go

+package addresses
+
+import "github.com/ethereum/go-ethereum/common"
+
+var (
+	// DefaultSenderAddr is known as DEFAULT_SENDER = address(uint160(uint256(keccak256("foundry default caller"))))
+	DefaultSenderAddr = common.HexToAddress("0x1804c8AB1F12E6bbf3894d4083f33e07309d1f38")
+	// DefaultScriptAddr is the address of the initial executing script, computed from:
+	// cast compute-address --nonce 1 0x1804c8AB1F12E6bbf3894d4083f33e07309d1f38
+	DefaultScriptAddr = common.HexToAddress("0x7FA9385bE102ac3EAc297483Dd6233D62b3e1496")
+	// VMAddr is known as VM_ADDRESS = address(uint160(uint256(keccak256("hevm cheat code"))));
+	VMAddr = common.HexToAddress("0x7109709ECfa91a80626fF3989D68f67F5b1DD12D")
+	// ConsoleAddr is known as CONSOLE, "console.log" in ascii.
+	// Utils like console.sol and console2.sol work by executing a staticcall to this address.
+	ConsoleAddr = common.HexToAddress("0x000000000000000000636F6e736F6c652e6c6f67")
+	// ScriptDeployer is used for temporary scripts address(uint160(uint256(keccak256("op-stack script deployer"))))
+	ScriptDeployer = common.HexToAddress("0x76Ce131128F3616871f8CDA86d18fAB44E4d0D8B")
+	// ForgeDeployer is used by some scripts as a default deployer address, e.g. makeAddr("deployer")
+	ForgeDeployer = common.HexToAddress("0xaE0bDc4eEAC5E950B67C6819B118761CaAF61946")
+)

op-chain-ops/script/cheatcodes.go

 package script
 
+import (
+	"fmt"
+
+	"github.com/ethereum/go-ethereum/core/vm"
+)
+
 // CheatCodesPrecompile implements the Forge vm cheatcodes.
 // Note that forge-std wraps these cheatcodes,
 // and provides additional convenience functions that use these cheatcodes.
 type CheatCodesPrecompile struct {
 	h *Host
 }
+
+// AccessControlledPrecompile wraps a precompile,
+// and checks that the caller has cheatcode access.
+type AccessControlledPrecompile struct {
+	h     *Host
+	inner vm.PrecompiledContract
+}
+
+var _ vm.PrecompiledContract = (*AccessControlledPrecompile)(nil)
+
+func (c *AccessControlledPrecompile) RequiredGas(input []byte) uint64 {
+	// call-frame is not open yet, and prank is ignored for cheatcode access-checking.
+	accessor := c.h.SelfAddress()
+	_, ok := c.h.allowedCheatcodes[accessor]
+	if !ok {
+		// Don't just return infinite gas, we can allow it to run,
+		// and then revert with a proper error message.
+		return 0
+	}
+	return c.inner.RequiredGas(input)
+}
+
+func (c *AccessControlledPrecompile) Run(input []byte) ([]byte, error) {
+	// call-frame is not open yet, and prank is ignored for cheatcode access-checking.
+	accessor := c.h.SelfAddress()
+	if !c.h.AllowedCheatcodes(accessor) {
+		c.h.log.Error("Cheatcode access denied!", "caller", accessor, "label", c.h.labels[accessor])
+		return encodeRevert(fmt.Errorf("call by %s to cheatcode precompile is not allowed", accessor))
+	}
+	return c.inner.Run(input)
+}

op-chain-ops/script/cheatcodes_environment.go

@@ -67,6 +67,10 @@ func (c *CheatCodesPrecompile) Load(account common.Address, slot [32]byte) [32]b
 // Etch implements https://book.getfoundry.sh/cheatcodes/etch
 func (c *CheatCodesPrecompile) Etch(who common.Address, code []byte) {
 	c.h.state.SetCode(who, bytes.Clone(code)) // important to clone; geth EVM will reuse the calldata memory.
+	if len(code) > 0 {
+		// if we're not just zeroing out the account: allow it to access cheatcodes
+		c.h.AllowCheatcodes(who)
+	}
 }
 
 // Deal implements https://book.getfoundry.sh/cheatcodes/deal

op-chain-ops/script/cheatcodes_forking.go

+package script
+
+import (
+	"errors"
+	"fmt"
+	"math/big"
+
+	"github.com/holiman/uint256"
+
+	"github.com/ethereum/go-ethereum/common"
+
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/forking"
+)
+
+func (c *CheatCodesPrecompile) CreateFork_31ba3498(urlOrAlias string) (*big.Int, error) {
+	return c.createFork(ForkWithURLOrAlias(urlOrAlias))
+}
+
+func (c *CheatCodesPrecompile) CreateFork_6ba3ba2b(urlOrAlias string, block *big.Int) (*big.Int, error) {
+	return c.createFork(ForkWithURLOrAlias(urlOrAlias), ForkWithBlockNumberU256(block))
+}
+
+func (c *CheatCodesPrecompile) CreateFork_7ca29682(urlOrAlias string, txHash common.Hash) (*big.Int, error) {
+	return c.createFork(ForkWithURLOrAlias(urlOrAlias), ForkWithTransaction(txHash))
+}
+
+// createFork implements vm.createFork:
+// https://book.getfoundry.sh/cheatcodes/create-fork
+func (c *CheatCodesPrecompile) createFork(opts ...ForkOption) (*big.Int, error) {
+	src, err := c.h.onFork(opts...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to setup fork source: %w", err)
+	}
+	id, err := c.h.state.CreateFork(src)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create fork: %w", err)
+	}
+	return id.U256().ToBig(), nil
+}
+
+func (c *CheatCodesPrecompile) CreateSelectFork_98680034(urlOrAlias string) (*big.Int, error) {
+	return c.createSelectFork(ForkWithURLOrAlias(urlOrAlias))
+}
+
+func (c *CheatCodesPrecompile) CreateSelectFork_71ee464d(urlOrAlias string, block *big.Int) (*big.Int, error) {
+	return c.createSelectFork(ForkWithURLOrAlias(urlOrAlias), ForkWithBlockNumberU256(block))
+}
+
+func (c *CheatCodesPrecompile) CreateSelectFork_84d52b7a(urlOrAlias string, txHash common.Hash) (*big.Int, error) {
+	return c.createSelectFork(ForkWithURLOrAlias(urlOrAlias), ForkWithTransaction(txHash))
+}
+
+// createSelectFork implements vm.createSelectFork:
+// https://book.getfoundry.sh/cheatcodes/create-select-fork
+func (c *CheatCodesPrecompile) createSelectFork(opts ...ForkOption) (*big.Int, error) {
+	src, err := c.h.onFork(opts...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to setup fork source: %w", err)
+	}
+	id, err := c.h.state.CreateSelectFork(src)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create-select fork: %w", err)
+	}
+	return id.U256().ToBig(), nil
+}
+
+// ActiveFork implements vm.activeFork:
+// https://book.getfoundry.sh/cheatcodes/active-fork
+func (c *CheatCodesPrecompile) ActiveFork() (*uint256.Int, error) {
+	id, active := c.h.state.ActiveFork()
+	if !active {
+		return nil, errors.New("no active fork")
+	}
+	return id.U256(), nil
+}
+
+// convenience method, to repeat the same URLOrAlias as the given fork when setting up a new fork
+func (c *CheatCodesPrecompile) forkURLOption(id forking.ForkID) ForkOption {
+	return func(cfg *ForkConfig) error {
+		urlOrAlias, err := c.h.state.ForkURLOrAlias(id)
+		if err != nil {
+			return err
+		}
+		return ForkWithURLOrAlias(urlOrAlias)(cfg)
+	}
+}
+
+func (c *CheatCodesPrecompile) RollFork_d9bbf3a1(block *big.Int) error {
+	id, ok := c.h.state.ActiveFork()
+	if !ok {
+		return errors.New("no active fork")
+	}
+	return c.rollFork(id, c.forkURLOption(id), ForkWithBlockNumberU256(block))
+}
+
+func (c *CheatCodesPrecompile) RollFork_0f29772b(txHash common.Hash) error {
+	id, ok := c.h.state.ActiveFork()
+	if !ok {
+		return errors.New("no active fork")
+	}
+	return c.rollFork(id, c.forkURLOption(id), ForkWithTransaction(txHash))
+}
+
+func (c *CheatCodesPrecompile) RollFork_d74c83a4(forkID *big.Int, block *big.Int) error {
+	id := forking.ForkIDFromBig(forkID)
+	return c.rollFork(id, c.forkURLOption(id), ForkWithBlockNumberU256(block))
+}
+
+func (c *CheatCodesPrecompile) RollFork_f2830f7b(forkID *uint256.Int, txHash common.Hash) error {
+	id := forking.ForkID(*forkID)
+	return c.rollFork(id, c.forkURLOption(id), ForkWithTransaction(txHash))
+}
+
+// rollFork implements vm.rollFork:
+// https://book.getfoundry.sh/cheatcodes/roll-fork
+func (c *CheatCodesPrecompile) rollFork(id forking.ForkID, opts ...ForkOption) error {
+	src, err := c.h.onFork(opts...)
+	if err != nil {
+		return fmt.Errorf("cannot setup fork source for roll-fork change: %w", err)
+	}
+	return c.h.state.ResetFork(id, src)
+}
+
+// MakePersistent_57e22dde implements vm.makePersistent:
+// https://book.getfoundry.sh/cheatcodes/make-persistent
+func (c *CheatCodesPrecompile) MakePersistent_57e22dde(account0 common.Address) {
+	c.h.state.MakePersistent(account0)
+}
+
+func (c *CheatCodesPrecompile) MakePersistent_4074e0a8(account0, account1 common.Address) {
+	c.h.state.MakePersistent(account0)
+	c.h.state.MakePersistent(account1)
+}
+
+func (c *CheatCodesPrecompile) MakePersistent_efb77a75(account0, account1, account2 common.Address) {
+	c.h.state.MakePersistent(account0)
+	c.h.state.MakePersistent(account1)
+	c.h.state.MakePersistent(account2)
+}
+
+func (c *CheatCodesPrecompile) MakePersistent_1d9e269e(accounts []common.Address) {
+	for _, addr := range accounts {
+		c.h.state.MakePersistent(addr)
+	}
+}
+
+// RevokePersistent_997a0222 implements vm.revokePersistent:
+// https://book.getfoundry.sh/cheatcodes/revoke-persistent
+func (c *CheatCodesPrecompile) RevokePersistent_997a0222(addr common.Address) {
+	c.h.state.RevokePersistent(addr)
+}
+
+func (c *CheatCodesPrecompile) RevokePersistent_3ce969e6(addrs []common.Address) {
+	for _, addr := range addrs {
+		c.h.state.RevokePersistent(addr)
+	}
+}
+
+// IsPersistent implements vm.isPersistent:
+// https://book.getfoundry.sh/cheatcodes/is-persistent
+func (c *CheatCodesPrecompile) IsPersistent(addr common.Address) bool {
+	return c.h.state.IsPersistent(addr)
+}
+
+// AllowCheatcodes implements vm.allowCheatcodes:
+// https://book.getfoundry.sh/cheatcodes/allow-cheatcodes
+func (c *CheatCodesPrecompile) AllowCheatcodes(addr common.Address) {
+	c.h.AllowCheatcodes(addr)
+}

op-chain-ops/script/console_test.go

@@ -7,6 +7,8 @@ import (
 	"math/rand" // nosemgrep
 	"testing"
 
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/addresses"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/ethereum/go-ethereum/log"
@@ -62,8 +64,8 @@ func TestFormatter(t *testing.T) {
 	require.Equal(t, "4.2", consoleFormat("%8e", big.NewInt(420000000)))
 	require.Equal(t, "foo true bar false", consoleFormat("foo %s bar %s", true, false))
 	require.Equal(t, "foo 1 bar 0", consoleFormat("foo %d bar %d", true, false))
-	require.Equal(t, "sender: "+DefaultSenderAddr.String(),
-		consoleFormat("sender: %s", DefaultSenderAddr))
+	require.Equal(t, "sender: "+addresses.DefaultSenderAddr.String(),
+		consoleFormat("sender: %s", addresses.DefaultSenderAddr))
 	require.Equal(t, "long 0.000000000000000042 number", consoleFormat("long %18e number", big.NewInt(42)))
 	require.Equal(t, "long 4200.000000000000000003 number", consoleFormat("long %18e number",
 		new(big.Int).Add(new(big.Int).Mul(

op-chain-ops/script/context.go

@@ -3,24 +3,9 @@ package script
 import (
 	"math/big"
 
-	"github.com/ethereum/go-ethereum/common"
-)
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/addresses"
 
-var (
-	// DefaultSenderAddr is known as DEFAULT_SENDER = address(uint160(uint256(keccak256("foundry default caller"))))
-	DefaultSenderAddr = common.HexToAddress("0x1804c8AB1F12E6bbf3894d4083f33e07309d1f38")
-	// DefaultScriptAddr is the address of the initial executing script, computed from:
-	// cast compute-address --nonce 1 0x1804c8AB1F12E6bbf3894d4083f33e07309d1f38
-	DefaultScriptAddr = common.HexToAddress("0x7FA9385bE102ac3EAc297483Dd6233D62b3e1496")
-	// VMAddr is known as VM_ADDRESS = address(uint160(uint256(keccak256("hevm cheat code"))));
-	VMAddr = common.HexToAddress("0x7109709ECfa91a80626fF3989D68f67F5b1DD12D")
-	// ConsoleAddr is known as CONSOLE, "console.log" in ascii.
-	// Utils like console.sol and console2.sol work by executing a staticcall to this address.
-	ConsoleAddr = common.HexToAddress("0x000000000000000000636F6e736F6c652e6c6f67")
-	// ScriptDeployer is used for temporary scripts address(uint160(uint256(keccak256("op-stack script deployer"))))
-	ScriptDeployer = common.HexToAddress("0x76Ce131128F3616871f8CDA86d18fAB44E4d0D8B")
-	// ForgeDeployer is used by some scripts as a default deployer address, e.g. makeAddr("deployer")
-	ForgeDeployer = common.HexToAddress("0xaE0bDc4eEAC5E950B67C6819B118761CaAF61946")
+	"github.com/ethereum/go-ethereum/common"
 )
 
 const (
@@ -42,8 +27,8 @@ type Context struct {
 
 var DefaultContext = Context{
 	ChainID:      big.NewInt(1337),
-	Sender:       DefaultSenderAddr,
-	Origin:       DefaultSenderAddr,
+	Sender:       addresses.DefaultSenderAddr,
+	Origin:       addresses.DefaultSenderAddr,
 	FeeRecipient: common.Address{},
 	GasLimit:     DefaultFoundryGasLimit,
 	BlockNum:     0,

op-chain-ops/script/fork.go

+package script
+
+import (
+	"fmt"
+	"math/big"
+
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/forking"
+	"github.com/ethereum/go-ethereum/common"
+)
+
+// ForkOption modifies a ForkConfig, and can be used by Host internals,
+// like the forking cheatcodes, to customize the forking action.
+type ForkOption func(cfg *ForkConfig) error
+
+// ForkHook is a callback to the user of the Host,
+// to translate an intent to fork into a source of data that can be forked with.
+type ForkHook func(opts *ForkConfig) (forking.ForkSource, error)
+
+// ForkConfig is a bundle of data to express a fork intent
+type ForkConfig struct {
+	URLOrAlias  string
+	BlockNumber *uint64      // latest if nil
+	Transaction *common.Hash // up to pre-state of given transaction
+}
+
+func ForkWithURLOrAlias(urlOrAlias string) ForkOption {
+	return func(cfg *ForkConfig) error {
+		cfg.URLOrAlias = urlOrAlias
+		return nil
+	}
+}
+
+func ForkWithBlockNumberU256(num *big.Int) ForkOption {
+	return func(cfg *ForkConfig) error {
+		if !num.IsUint64() {
+			return fmt.Errorf("block number %s is too large", num.String())
+		}
+		v := num.Uint64()
+		cfg.BlockNumber = &v
+		return nil
+	}
+}
+
+func ForkWithTransaction(txHash common.Hash) ForkOption {
+	return func(cfg *ForkConfig) error {
+		cfg.Transaction = &txHash
+		return nil
+	}
+}
+
+// onFork is called by script-internals to translate a fork-intent into forks data-source.
+func (h *Host) onFork(opts ...ForkOption) (forking.ForkSource, error) {
+	cfg := &ForkConfig{}
+	for _, opt := range opts {
+		if err := opt(cfg); err != nil {
+			return nil, err
+		}
+	}
+	return h.hooks.OnFork(cfg)
+}

op-chain-ops/script/forking/cache.go

+package forking
+
+import (
+	lru "github.com/hashicorp/golang-lru/v2"
+	"github.com/holiman/uint256"
+
+	"github.com/ethereum/go-ethereum/common"
+)
+
+type storageKey struct {
+	Addr common.Address
+	Slot common.Hash
+}
+
+// CachedSource wraps a ForkSource, and caches the retrieved data for faster repeat-queries.
+// The ForkSource should be immutable (as per the StateRoot value).
+// All cache data accumulates in-memory in LRU collections per data type.
+type CachedSource struct {
+	stateRoot common.Hash
+	src       ForkSource
+
+	nonces   *lru.Cache[common.Address, uint64]
+	balances *lru.Cache[common.Address, *uint256.Int]
+	storage  *lru.Cache[storageKey, common.Hash]
+	code     *lru.Cache[common.Address, []byte]
+}
+
+var _ ForkSource = (*CachedSource)(nil)
+
+func mustNewLRU[K comparable, V any](size int) *lru.Cache[K, V] {
+	out, err := lru.New[K, V](size)
+	if err != nil {
+		panic(err) // bad size parameter may produce an error
+	}
+	return out
+}
+
+func Cache(src ForkSource) *CachedSource {
+	return &CachedSource{
+		stateRoot: src.StateRoot(),
+		src:       src,
+		nonces:    mustNewLRU[common.Address, uint64](1000),
+		balances:  mustNewLRU[common.Address, *uint256.Int](1000),
+		storage:   mustNewLRU[storageKey, common.Hash](1000),
+		code:      mustNewLRU[common.Address, []byte](100),
+	}
+}
+
+func (c *CachedSource) URLOrAlias() string {
+	return c.src.URLOrAlias()
+}
+
+func (c *CachedSource) StateRoot() common.Hash {
+	return c.stateRoot
+}
+
+func (c *CachedSource) Nonce(addr common.Address) (uint64, error) {
+	if v, ok := c.nonces.Get(addr); ok {
+		return v, nil
+	}
+	v, err := c.src.Nonce(addr)
+	if err != nil {
+		return 0, err
+	}
+	c.nonces.Add(addr, v)
+	return v, nil
+}
+
+func (c *CachedSource) Balance(addr common.Address) (*uint256.Int, error) {
+	if v, ok := c.balances.Get(addr); ok {
+		return v.Clone(), nil
+	}
+	v, err := c.src.Balance(addr)
+	if err != nil {
+		return nil, err
+	}
+	c.balances.Add(addr, v)
+	return v.Clone(), nil
+}
+
+func (c *CachedSource) StorageAt(addr common.Address, key common.Hash) (common.Hash, error) {
+	if v, ok := c.storage.Get(storageKey{Addr: addr, Slot: key}); ok {
+		return v, nil
+	}
+	v, err := c.src.StorageAt(addr, key)
+	if err != nil {
+		return common.Hash{}, err
+	}
+	c.storage.Add(storageKey{Addr: addr, Slot: key}, v)
+	return v, nil
+}
+
+func (c *CachedSource) Code(addr common.Address) ([]byte, error) {
+	if v, ok := c.code.Get(addr); ok {
+		return v, nil
+	}
+	v, err := c.src.Code(addr)
+	if err != nil {
+		return nil, err
+	}
+	c.code.Add(addr, v)
+	return v, nil
+}

op-chain-ops/script/forking/cache_test.go

+package forking
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/holiman/uint256"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+// MockForkSource implements ForkSource interface for testing
+type MockForkSource struct {
+	mock.Mock
+}
+
+func (m *MockForkSource) URLOrAlias() string {
+	args := m.Called()
+	return args.String(0)
+}
+
+func (m *MockForkSource) StateRoot() common.Hash {
+	args := m.Called()
+	return args.Get(0).(common.Hash)
+}
+
+func (m *MockForkSource) Nonce(addr common.Address) (uint64, error) {
+	args := m.Called(addr)
+	return args.Get(0).(uint64), args.Error(1)
+}
+
+func (m *MockForkSource) Balance(addr common.Address) (*uint256.Int, error) {
+	args := m.Called(addr)
+	if args.Get(0) == nil {
+		return nil, args.Error(1)
+	}
+	return args.Get(0).(*uint256.Int), args.Error(1)
+}
+
+func (m *MockForkSource) StorageAt(addr common.Address, key common.Hash) (common.Hash, error) {
+	args := m.Called(addr, key)
+	return args.Get(0).(common.Hash), args.Error(1)
+}
+
+func (m *MockForkSource) Code(addr common.Address) ([]byte, error) {
+	args := m.Called(addr)
+	return args.Get(0).([]byte), args.Error(1)
+}
+
+func setupCache(t *testing.T) (*CachedSource, *MockForkSource) {
+	mockSource := new(MockForkSource)
+	stateRoot := common.HexToHash("0x1234")
+	mockSource.On("StateRoot").Return(stateRoot)
+	mockSource.On("URLOrAlias").Return("test_source")
+
+	cached := Cache(mockSource)
+	require.NotNil(t, cached)
+	require.Equal(t, stateRoot, cached.StateRoot())
+	require.Equal(t, "test_source", cached.URLOrAlias())
+
+	return cached, mockSource
+}
+
+func TestCachedSource_Nonce(t *testing.T) {
+	cached, mockSource := setupCache(t)
+	addr := common.HexToAddress("0x1234")
+	expectedNonce := uint64(42)
+
+	// First call should hit the source
+	mockSource.On("Nonce", addr).Return(expectedNonce, nil).Once()
+
+	nonce, err := cached.Nonce(addr)
+	require.NoError(t, err)
+	require.Equal(t, expectedNonce, nonce)
+
+	// Second call should use cache
+	nonce, err = cached.Nonce(addr)
+	require.NoError(t, err)
+	require.Equal(t, expectedNonce, nonce)
+
+	mockSource.AssertNumberOfCalls(t, "Nonce", 1)
+}
+
+func TestCachedSource_Balance(t *testing.T) {
+	cached, mockSource := setupCache(t)
+	addr := common.HexToAddress("0x5678")
+	expectedBalance := uint256.NewInt(1000)
+
+	// First call should hit the source
+	mockSource.On("Balance", addr).Return(expectedBalance, nil).Once()
+
+	balance, err := cached.Balance(addr)
+	require.NoError(t, err)
+	require.Equal(t, expectedBalance, balance)
+
+	// Second call should use cache
+	balance, err = cached.Balance(addr)
+	require.NoError(t, err)
+	require.Equal(t, expectedBalance, balance)
+
+	// Verify the returned balance is a clone
+	balance.Add(balance, uint256.NewInt(1))
+	cachedBalance, _ := cached.Balance(addr)
+	require.Equal(t, expectedBalance, cachedBalance)
+
+	mockSource.AssertNumberOfCalls(t, "Balance", 1)
+}
+
+func TestCachedSource_Storage(t *testing.T) {
+	cached, mockSource := setupCache(t)
+	addr := common.HexToAddress("0x9abc")
+	slot := common.HexToHash("0xdef0")
+	expectedValue := common.HexToHash("0x1234")
+
+	// First call should hit the source
+	mockSource.On("StorageAt", addr, slot).Return(expectedValue, nil).Once()
+
+	value, err := cached.StorageAt(addr, slot)
+	require.NoError(t, err)
+	require.Equal(t, expectedValue, value)
+
+	// Second call should use cache
+	value, err = cached.StorageAt(addr, slot)
+	require.NoError(t, err)
+	require.Equal(t, expectedValue, value)
+
+	mockSource.AssertNumberOfCalls(t, "StorageAt", 1)
+}
+
+func TestCachedSource_Code(t *testing.T) {
+	cached, mockSource := setupCache(t)
+	addr := common.HexToAddress("0xdef0")
+	expectedCode := []byte{1, 2, 3, 4}
+
+	// First call should hit the source
+	mockSource.On("Code", addr).Return(expectedCode, nil).Once()
+
+	code, err := cached.Code(addr)
+	require.NoError(t, err)
+	require.Equal(t, expectedCode, code)
+
+	// Second call should use cache
+	code, err = cached.Code(addr)
+	require.NoError(t, err)
+	require.Equal(t, expectedCode, code)
+
+	mockSource.AssertNumberOfCalls(t, "Code", 1)
+}
+
+func TestCachedSource_CacheEviction(t *testing.T) {
+	cached, mockSource := setupCache(t)
+
+	// Test nonce cache eviction
+	for i := 0; i < 1001; i++ { // Cache size is 1000
+		addr := common.BigToAddress(big.NewInt(int64(i)))
+		mockSource.On("Nonce", addr).Return(uint64(i), nil).Once()
+		_, _ = cached.Nonce(addr)
+	}
+
+	// This should cause first address to be evicted
+	firstAddr := common.BytesToAddress([]byte{0})
+	mockSource.On("Nonce", firstAddr).Return(uint64(0), nil).Once()
+	_, _ = cached.Nonce(firstAddr)
+
+	mockSource.AssertNumberOfCalls(t, "Nonce", 1002) // 1001 + 1 for evicted key
+}
+
+func TestCachedSource_MultipleStorageSlots(t *testing.T) {
+	cached, mockSource := setupCache(t)
+	addr := common.HexToAddress("0xabcd")
+	slot1 := common.HexToHash("0x1111")
+	slot2 := common.HexToHash("0x2222")
+	value1 := common.HexToHash("0x3333")
+	value2 := common.HexToHash("0x4444")
+
+	mockSource.On("StorageAt", addr, slot1).Return(value1, nil).Once()
+	mockSource.On("StorageAt", addr, slot2).Return(value2, nil).Once()
+
+	// Different slots should trigger separate cache entries
+	val1, err := cached.StorageAt(addr, slot1)
+	require.NoError(t, err)
+	require.Equal(t, value1, val1)
+
+	val2, err := cached.StorageAt(addr, slot2)
+	require.NoError(t, err)
+	require.Equal(t, value2, val2)
+
+	// Verify both are cached
+	val1Again, _ := cached.StorageAt(addr, slot1)
+	val2Again, _ := cached.StorageAt(addr, slot2)
+	require.Equal(t, value1, val1Again)
+	require.Equal(t, value2, val2Again)
+
+	mockSource.AssertNumberOfCalls(t, "StorageAt", 2)
+}

op-chain-ops/script/forking/db.go

+package forking
+
+import (
+	"fmt"
+
+	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/triedb/pathdb"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/state/snapshot"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/ethdb"
+	"github.com/ethereum/go-ethereum/trie/utils"
+	"github.com/ethereum/go-ethereum/triedb"
+)
+
+// ForkDB is a virtual state database: it wraps a forked accounts trie,
+// and can maintain a state diff, so we can mutate the forked state,
+// and even finalize state changes (so we can accurately measure things like cold storage gas cost).
+type ForkDB struct {
+	active *ForkedAccountsTrie
+}
+
+// Reader for read-only access to a known state. All cold reads go through this.
+// So the state-DB creates one initially, and then holds on to it.
+// The diff will be overlayed on the reader still. To get rid of the diff, it has to be explicitly cleared.
+// Warning: diffs applied to the original state that the reader wraps will be visible.
+// Geth StateDB is meant to be reinitialized after committing state.
+func (f *ForkDB) Reader(root common.Hash) (state.Reader, error) {
+	if root != f.active.stateRoot {
+		return nil, fmt.Errorf("current state is at %s, cannot open state at %s", f.active.stateRoot, root)
+	}
+	return &forkStateReader{
+		f.active,
+	}, nil
+}
+
+func (f *ForkDB) Snapshot() *snapshot.Tree {
+	return nil
+}
+
+var _ state.Database = (*ForkDB)(nil)
+
+func NewForkDB(source ForkSource) *ForkDB {
+	return &ForkDB{active: &ForkedAccountsTrie{
+		stateRoot: source.StateRoot(),
+		src:       source,
+		diff:      NewExportDiff(),
+	}}
+}
+
+// fakeRoot is just a marker; every account we load into the fork-db has this storage-root.
+// When opening a storage-trie, we sanity-check we have this root, or an empty trie.
+// And then just return the same global trie view for storage reads/writes.
+var fakeRoot = common.Hash{0: 42}
+
+func (f *ForkDB) OpenTrie(root common.Hash) (state.Trie, error) {
+	if f.active.stateRoot != root {
+		return nil, fmt.Errorf("active fork is at %s, but tried to open %s", f.active.stateRoot, root)
+	}
+	return f.active, nil
+}
+
+func (f *ForkDB) OpenStorageTrie(stateRoot common.Hash, address common.Address, root common.Hash, trie state.Trie) (state.Trie, error) {
+	if f.active.stateRoot != stateRoot {
+		return nil, fmt.Errorf("active fork is at %s, but tried to open account %s of state %s", f.active.stateRoot, address, stateRoot)
+	}
+	if _, ok := trie.(*ForkedAccountsTrie); !ok {
+		return nil, fmt.Errorf("ForkDB tried to open non-fork storage-trie %v", trie)
+	}
+	if root != fakeRoot && root != types.EmptyRootHash {
+		return nil, fmt.Errorf("ForkDB unexpectedly was queried with real looking storage root: %s", root)
+	}
+	return f.active, nil
+}
+
+func (f *ForkDB) CopyTrie(trie state.Trie) state.Trie {
+	if st, ok := trie.(*ForkedAccountsTrie); ok {
+		return st.Copy()
+	}
+	panic(fmt.Errorf("ForkDB tried to copy non-fork trie %v", trie))
+}
+
+func (f *ForkDB) ContractCode(addr common.Address, codeHash common.Hash) ([]byte, error) {
+	return f.active.ContractCode(addr, codeHash)
+}
+
+func (f *ForkDB) ContractCodeSize(addr common.Address, codeHash common.Hash) (int, error) {
+	return f.active.ContractCodeSize(addr, codeHash)
+}
+
+func (f *ForkDB) DiskDB() ethdb.KeyValueStore {
+	panic("DiskDB() during active Fork is not supported")
+}
+
+func (f *ForkDB) PointCache() *utils.PointCache {
+	panic("PointCache() is not supported")
+}
+
+func (f *ForkDB) TrieDB() *triedb.Database {
+	// The TrieDB is unused, but geth does use to check if Verkle is activated.
+	// So we have to create a read-only dummy one, to communicate that verkle really is disabled.
+	diskDB := rawdb.NewMemoryDatabase()
+	tdb := triedb.NewDatabase(diskDB, &triedb.Config{
+		Preimages: false,
+		IsVerkle:  false,
+		HashDB:    nil,
+		PathDB: &pathdb.Config{
+			StateHistory:   0,
+			CleanCacheSize: 0,
+			DirtyCacheSize: 0,
+			ReadOnly:       true,
+		},
+	})
+	return tdb
+}

op-chain-ops/script/forking/diff.go

+package forking
+
+import (
+	"bytes"
+	"maps"
+
+	"github.com/holiman/uint256"
+
+	"github.com/ethereum/go-ethereum/common"
+)
+
+// AccountDiff represents changes to an account. Unchanged values of the account are not included.
+type AccountDiff struct {
+	// Nonce change.
+	// No diff if nil.
+	Nonce *uint64 `json:"nonce"`
+
+	// Balance change.
+	// No diff if nil.
+	Balance *uint256.Int `json:"balance"`
+
+	// Storage changes.
+	// No diff if not present in map. Deletions are zero-value entries.
+	Storage map[common.Hash]common.Hash `json:"storage"`
+
+	// CodeHash, for lookup of contract bytecode in the code diff map.
+	// No code-diff if nil.
+	CodeHash *common.Hash `json:"codeHash"`
+}
+
+func (d *AccountDiff) Copy() *AccountDiff {
+	var out AccountDiff
+	if d.Nonce != nil {
+		v := *d.Nonce // copy the value
+		out.Nonce = &v
+	}
+	if d.Balance != nil {
+		out.Balance = d.Balance.Clone()
+	}
+	if d.Storage != nil {
+		out.Storage = maps.Clone(d.Storage)
+	}
+	if d.CodeHash != nil {
+		h := *d.CodeHash
+		out.CodeHash = &h
+	}
+	return &out
+}
+
+type ExportDiff struct {
+	// Accounts diff. Deleted accounts are set to nil.
+	// Warning: this only contains finalized state changes.
+	// The state itself holds on to non-flushed changes.
+	Account map[common.Address]*AccountDiff `json:"account"`
+
+	// Stores new contract codes by code-hash
+	Code map[common.Hash][]byte `json:"code"`
+}
+
+func NewExportDiff() *ExportDiff {
+	return &ExportDiff{
+		Account: make(map[common.Address]*AccountDiff),
+		Code:    make(map[common.Hash][]byte),
+	}
+}
+
+func (ed *ExportDiff) Copy() *ExportDiff {
+	out := &ExportDiff{
+		Account: make(map[common.Address]*AccountDiff),
+		Code:    make(map[common.Hash][]byte),
+	}
+	for addr, acc := range ed.Account {
+		out.Account[addr] = acc.Copy()
+	}
+	for addr, code := range ed.Code {
+		out.Code[addr] = bytes.Clone(code)
+	}
+	return out
+}
+
+func (ed *ExportDiff) Any() bool {
+	return len(ed.Code) > 0 || len(ed.Account) > 0
+}
+
+func (ed *ExportDiff) Clear() {
+	ed.Account = make(map[common.Address]*AccountDiff)
+	ed.Code = make(map[common.Hash][]byte)
+}

op-chain-ops/script/forking/forking_test.go

+package forking
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/holiman/uint256"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/tracing"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/triedb"
+	"github.com/ethereum/go-ethereum/triedb/hashdb"
+)
+
+type TestForkSource struct {
+	urlOrAlias string
+	stateRoot  common.Hash
+	nonces     map[common.Address]uint64
+	balances   map[common.Address]*uint256.Int
+	storage    map[common.Address]map[common.Hash]common.Hash
+	code       map[common.Address][]byte
+}
+
+func (t TestForkSource) URLOrAlias() string {
+	return t.urlOrAlias
+}
+
+func (t TestForkSource) StateRoot() common.Hash {
+	return t.stateRoot
+}
+
+func (t TestForkSource) Nonce(addr common.Address) (uint64, error) {
+	return t.nonces[addr], nil
+}
+
+func (t TestForkSource) Balance(addr common.Address) (*uint256.Int, error) {
+	b, ok := t.balances[addr]
+	if !ok {
+		return uint256.NewInt(0), nil
+	}
+	return b.Clone(), nil
+}
+
+func (t TestForkSource) StorageAt(addr common.Address, key common.Hash) (common.Hash, error) {
+	storage, ok := t.storage[addr]
+	if !ok {
+		return common.Hash{}, nil
+	}
+	return storage[key], nil
+}
+
+func (t TestForkSource) Code(addr common.Address) ([]byte, error) {
+	return t.code[addr], nil
+}
+
+var _ ForkSource = (*TestForkSource)(nil)
+
+func TestForking(t *testing.T) {
+	// create regular DB
+	rawDB := rawdb.NewMemoryDatabase()
+	stateDB := state.NewDatabase(triedb.NewDatabase(rawDB, &triedb.Config{
+		Preimages: true, // To be able to iterate the state we need the Preimages
+		IsVerkle:  false,
+		HashDB:    hashdb.Defaults,
+		PathDB:    nil,
+	}), nil)
+	baseState, err := state.New(types.EmptyRootHash, stateDB)
+	if err != nil {
+		panic(fmt.Errorf("failed to create memory state db: %w", err))
+	}
+	forkState := NewForkableState(baseState)
+
+	// No active fork yet
+	id, active := forkState.ActiveFork()
+	require.False(t, active)
+	require.Equal(t, ForkID{}, id)
+
+	name, err := forkState.ForkURLOrAlias(ForkID{})
+	require.ErrorContains(t, err, "default")
+	require.Equal(t, "", name)
+
+	alice := common.Address(bytes.Repeat([]byte{0xaa}, 20))
+	bob := common.Address(bytes.Repeat([]byte{0xbb}, 20))
+
+	forkState.CreateAccount(alice)
+	forkState.SetNonce(alice, 3)
+	forkState.AddBalance(alice, uint256.NewInt(123), tracing.BalanceChangeUnspecified)
+	// Check if writes worked
+	require.Equal(t, uint64(123), forkState.GetBalance(alice).Uint64())
+	require.Equal(t, uint64(3), forkState.GetNonce(alice))
+	// No active fork yet, balance change should be applied to underlying base-state
+	require.Equal(t, uint64(123), baseState.GetBalance(alice).Uint64())
+	require.Equal(t, uint64(3), baseState.GetNonce(alice))
+
+	src1 := &TestForkSource{
+		urlOrAlias: "src 1",
+		stateRoot:  crypto.Keccak256Hash([]byte("test fork state 1")),
+		nonces: map[common.Address]uint64{
+			alice: uint64(42),
+			bob:   uint64(1000),
+		},
+		balances: make(map[common.Address]*uint256.Int),
+		storage:  make(map[common.Address]map[common.Hash]common.Hash),
+		code:     make(map[common.Address][]byte),
+	}
+	forkA, err := forkState.CreateSelectFork(src1)
+	require.NoError(t, err)
+	// Check that we selected A
+	id, active = forkState.ActiveFork()
+	require.True(t, active)
+	require.Equal(t, forkA, id)
+	name, err = forkState.ForkURLOrAlias(forkA)
+	require.NoError(t, err)
+	require.Equal(t, "src 1", name)
+
+	// the fork has a different nonce for alice
+	require.Equal(t, uint64(42), forkState.GetNonce(alice))
+	// the fork has Bob, which didn't exist thus far
+	require.Equal(t, uint64(1000), forkState.GetNonce(bob))
+
+	// Apply a diff change on top of the fork
+	forkState.SetNonce(bob, 99999)
+
+	// Now unselect the fork, going back to the default again.
+	require.NoError(t, forkState.SelectFork(ForkID{}))
+	// No longer active fork
+	id, active = forkState.ActiveFork()
+	require.False(t, active)
+	require.Equal(t, ForkID{}, id)
+
+	// Check that things are back to normal
+	require.Equal(t, uint64(3), forkState.GetNonce(alice))
+	require.Equal(t, uint64(0), forkState.GetNonce(bob))
+
+	// Make a change to the base-state, to see if it survives going back to the fork.
+	forkState.SetNonce(bob, 5)
+
+	// Re-select the fork, see if the changes come back, including the diff we made
+	require.NoError(t, forkState.SelectFork(forkA))
+	require.Equal(t, uint64(42), forkState.GetNonce(alice))
+	require.Equal(t, uint64(99999), forkState.GetNonce(bob))
+
+	// This change will continue to be visible across forks,
+	// alice is going to be persistent.
+	forkState.SetNonce(alice, 777)
+
+	// Now make Alice persistent, see if we can get the original value
+	forkState.MakePersistent(alice)
+
+	// Activate a fork, to see if alice is really persistent
+	src2 := &TestForkSource{
+		urlOrAlias: "src 2",
+		stateRoot:  crypto.Keccak256Hash([]byte("test fork state 2")),
+		nonces: map[common.Address]uint64{
+			alice: uint64(2222),
+			bob:   uint64(222),
+		},
+		balances: make(map[common.Address]*uint256.Int),
+		storage:  make(map[common.Address]map[common.Hash]common.Hash),
+		code:     make(map[common.Address][]byte),
+	}
+	tmpFork, err := forkState.CreateSelectFork(src2)
+	require.NoError(t, err)
+	require.Equal(t, uint64(777), forkState.GetNonce(alice), "persistent original value")
+	// While bob is still read from the fork
+	require.Equal(t, uint64(222), forkState.GetNonce(bob), "bob is forked")
+
+	// Mutate both, and undo the fork, to test if the persistent change is still there in non-fork mode
+	forkState.SetNonce(alice, 1001) // this mutates forkA, because alice was made persistent there
+	forkState.SetNonce(bob, 1002)
+	require.NoError(t, forkState.SelectFork(ForkID{}))
+	require.Equal(t, uint64(1001), forkState.GetNonce(alice), "alice is persistent")
+	require.Equal(t, uint64(5), forkState.GetNonce(bob), "bob is not persistent")
+
+	// Stop alice persistence. Forks can now override it again.
+	forkState.RevokePersistent(alice)
+	// This foundry behavior is unspecified/undocumented.
+	// Not sure if correctly doing it by dropping the previously persisted state if it comes from another fork.
+	require.Equal(t, uint64(3), forkState.GetNonce(alice))
+	require.Equal(t, uint64(3), baseState.GetNonce(alice))
+	require.Equal(t, uint64(5), forkState.GetNonce(bob))
+
+	// Create another fork, don't select it immediately
+	src3 := &TestForkSource{
+		urlOrAlias: "src 3",
+		stateRoot:  crypto.Keccak256Hash([]byte("test fork state 3")),
+		nonces: map[common.Address]uint64{
+			alice: uint64(3333),
+		},
+		balances: make(map[common.Address]*uint256.Int),
+		storage:  make(map[common.Address]map[common.Hash]common.Hash),
+		code:     make(map[common.Address][]byte),
+	}
+	forkB, err := forkState.CreateFork(src3)
+	require.NoError(t, err)
+
+	id, active = forkState.ActiveFork()
+	require.False(t, active)
+	require.Equal(t, ForkID{}, id)
+
+	// forkA is still bound to src 1
+	name, err = forkState.ForkURLOrAlias(forkA)
+	require.NoError(t, err)
+	require.Equal(t, "src 1", name)
+	// tmpFork is still bound to src 2
+	name, err = forkState.ForkURLOrAlias(tmpFork)
+	require.NoError(t, err)
+	require.Equal(t, "src 2", name)
+	// forkB is on src 3
+	name, err = forkState.ForkURLOrAlias(forkB)
+	require.NoError(t, err)
+	require.Equal(t, "src 3", name)
+
+	require.Equal(t, uint64(3), forkState.GetNonce(alice), "not forked yet")
+	require.NoError(t, forkState.SelectFork(forkB))
+	id, active = forkState.ActiveFork()
+	require.True(t, active)
+	require.Equal(t, forkB, id)
+
+	// check if successfully forked now
+	require.Equal(t, uint64(3333), forkState.GetNonce(alice), "fork B active now")
+	// Bob is not in this fork. But that doesn't mean the base-state should be used.
+	require.Equal(t, uint64(0), forkState.GetNonce(bob))
+
+	// See if we can go from B straight to A
+	require.NoError(t, forkState.SelectFork(forkA))
+	require.Equal(t, uint64(1001), forkState.GetNonce(alice), "alice from A says hi")
+	// And back to B
+	require.NoError(t, forkState.SelectFork(forkB))
+	require.Equal(t, uint64(3333), forkState.GetNonce(alice), "alice from B says hi")
+
+	// And a fork on top of a fork; forks don't stack, they are their own individual contexts.
+	src4 := &TestForkSource{
+		urlOrAlias: "src 4",
+		stateRoot:  crypto.Keccak256Hash([]byte("test fork state 4")),
+		nonces: map[common.Address]uint64{
+			bob: uint64(9000),
+		},
+		balances: make(map[common.Address]*uint256.Int),
+		storage:  make(map[common.Address]map[common.Hash]common.Hash),
+		code:     make(map[common.Address][]byte),
+	}
+	forkC, err := forkState.CreateSelectFork(src4)
+	require.NoError(t, err)
+	// No alice in this fork.
+	require.Equal(t, uint64(0), forkState.GetNonce(alice))
+	// But bob is set
+	require.Equal(t, uint64(9000), forkState.GetNonce(bob))
+
+	// Put in some mutations, for the fork-diff testing
+	forkState.SetNonce(alice, 1234)
+	forkState.SetBalance(alice, uint256.NewInt(100_000), tracing.BalanceChangeUnspecified)
+	forkState.SetState(alice, common.Hash{4}, common.Hash{42})
+	forkState.SetState(alice, common.Hash{5}, common.Hash{100})
+	forkState.SetCode(alice, []byte("hello world"))
+
+	// Check the name
+	name, err = forkState.ForkURLOrAlias(forkC)
+	require.NoError(t, err)
+	require.Equal(t, "src 4", name)
+
+	// Now test our fork-diff exporting:
+	// it needs to reflect the changes we made to the fork, but not other fork contents.
+	forkADiff, err := forkState.ExportDiff(forkA)
+	require.NoError(t, err)
+	require.NotNil(t, forkADiff.Account[alice])
+	require.Equal(t, uint64(1001), *forkADiff.Account[alice].Nonce)
+	require.Equal(t, uint64(99999), *forkADiff.Account[bob].Nonce)
+
+	forkBDiff, err := forkState.ExportDiff(forkB)
+	require.NoError(t, err)
+	require.Len(t, forkBDiff.Account, 0, "no changes to fork B")
+
+	forkCDiff, err := forkState.ExportDiff(forkC)
+	require.NoError(t, err)
+	require.Contains(t, forkCDiff.Account, alice)
+	require.NotContains(t, forkCDiff.Account, bob)
+	require.Equal(t, uint64(1234), *forkCDiff.Account[alice].Nonce)
+	require.Equal(t, uint64(100_000), forkCDiff.Account[alice].Balance.Uint64())
+	require.Equal(t, common.Hash{42}, forkCDiff.Account[alice].Storage[common.Hash{4}])
+	require.Equal(t, common.Hash{100}, forkCDiff.Account[alice].Storage[common.Hash{5}])
+	require.Equal(t, crypto.Keccak256Hash([]byte("hello world")), *forkCDiff.Account[alice].CodeHash)
+	require.Equal(t, []byte("hello world"), forkCDiff.Code[*forkCDiff.Account[alice].CodeHash])
+}

op-chain-ops/script/forking/iface.go

+package forking
+
+import (
+	"math/big"
+
+	"github.com/ethereum/go-ethereum/core/tracing"
+	"github.com/holiman/uint256"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/vm"
+)
+
+type VMStateDB interface {
+	vm.StateDB
+	Finalise(deleteEmptyObjects bool)
+	// SetBalance sets the balance of an account. Not part of the geth VM StateDB interface (add/sub balance are).
+	SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason)
+}
+
+// ForkID is an identifier of a fork
+type ForkID uint256.Int
+
+func ForkIDFromBig(b *big.Int) ForkID {
+	return ForkID(*uint256.MustFromBig(b))
+}
+
+// U256 returns a uint256 copy of the fork ID, for usage inside the EVM.
+func (id *ForkID) U256() *uint256.Int {
+	return new(uint256.Int).Set((*uint256.Int)(id))
+}
+
+func (id ForkID) String() string {
+	return (*uint256.Int)(&id).String()
+}
+
+// ForkSource is a read-only source for ethereum state,
+// that can be used to fork a ForkableState.
+type ForkSource interface {
+	// URLOrAlias returns the URL or alias that the fork uses. This is not unique to a single fork.
+	URLOrAlias() string
+	// StateRoot returns the accounts-trie root of the committed-to state.
+	// This root must never change.
+	StateRoot() common.Hash
+	// Nonce returns 0, without error, if the account does not exist.
+	Nonce(addr common.Address) (uint64, error)
+	// Balance returns 0, without error, if the account does not exist.
+	Balance(addr common.Address) (*uint256.Int, error)
+	// StorageAt returns a zeroed hash, without error, if the storage does not exist.
+	StorageAt(addr common.Address, key common.Hash) (common.Hash, error)
+	// Code returns an empty byte slice, without error, if no code exists.
+	Code(addr common.Address) ([]byte, error)
+}

op-chain-ops/script/forking/reader.go

+package forking
+
+import (
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/types"
+)
+
+// forkStateReader implements the state.Reader abstraction,
+// for read-only access to a state-trie at a particular state-root.
+type forkStateReader struct {
+	trie *ForkedAccountsTrie
+}
+
+var _ state.Reader = (*forkStateReader)(nil)
+
+func (f *forkStateReader) Account(addr common.Address) (*types.StateAccount, error) {
+	acc, err := f.trie.GetAccount(addr)
+	if err != nil {
+		return nil, err
+	}
+	// We copy because the Reader interface defines that it should be safe to modify after returning.
+	return acc.Copy(), nil
+}
+
+func (f *forkStateReader) Storage(addr common.Address, slot common.Hash) (common.Hash, error) {
+	v, err := f.trie.GetStorage(addr, slot[:])
+	if err != nil {
+		return common.Hash{}, err
+	}
+	return common.Hash(v), nil
+}
+
+func (f *forkStateReader) Copy() state.Reader {
+	return f
+}

op-chain-ops/script/forking/rpc.go

+package forking
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/holiman/uint256"
+
+	"github.com/ethereum/go-ethereum"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/common/hexutil"
+
+	"github.com/ethereum-optimism/optimism/op-service/retry"
+)
+
+type RPCClient interface {
+	CallContext(ctx context.Context, result any, method string, args ...any) error
+}
+
+type RPCSource struct {
+	stateRoot common.Hash
+	blockHash common.Hash
+
+	maxAttempts int
+	timeout     time.Duration
+	strategy    retry.Strategy
+
+	ctx    context.Context
+	cancel context.CancelFunc
+
+	client     RPCClient
+	urlOrAlias string
+}
+
+var _ ForkSource = (*RPCSource)(nil)
+
+func RPCSourceByNumber(urlOrAlias string, cl RPCClient, num uint64) (*RPCSource, error) {
+	src := newRPCSource(urlOrAlias, cl)
+	err := src.init(hexutil.Uint64(num))
+	return src, err
+}
+
+func RPCSourceByHash(urlOrAlias string, cl RPCClient, h common.Hash) (*RPCSource, error) {
+	src := newRPCSource(urlOrAlias, cl)
+	err := src.init(h)
+	return src, err
+}
+
+func newRPCSource(urlOrAlias string, cl RPCClient) *RPCSource {
+	ctx, cancel := context.WithCancel(context.Background())
+	return &RPCSource{
+		maxAttempts: 10,
+		timeout:     time.Second * 10,
+		strategy:    retry.Exponential(),
+		ctx:         ctx,
+		cancel:      cancel,
+		client:      cl,
+		urlOrAlias:  urlOrAlias,
+	}
+}
+
+type Header struct {
+	StateRoot common.Hash `json:"stateRoot"`
+	BlockHash common.Hash `json:"hash"`
+}
+
+func (r *RPCSource) init(id any) error {
+	head, err := retry.Do[*Header](r.ctx, r.maxAttempts, r.strategy, func() (*Header, error) {
+		var result *Header
+		err := r.client.CallContext(r.ctx, &result, "eth_getBlockByNumber", id, false)
+		if err == nil && result == nil {
+			err = ethereum.NotFound
+		}
+		return result, err
+	})
+	if err != nil {
+		return fmt.Errorf("failed to initialize RPC fork source around block %v: %w", id, err)
+	}
+	r.blockHash = head.BlockHash
+	r.stateRoot = head.StateRoot
+	return nil
+}
+
+func (c *RPCSource) URLOrAlias() string {
+	return c.urlOrAlias
+}
+
+func (r *RPCSource) BlockHash() common.Hash {
+	return r.blockHash
+}
+
+func (r *RPCSource) StateRoot() common.Hash {
+	return r.stateRoot
+}
+
+func (r *RPCSource) Nonce(addr common.Address) (uint64, error) {
+	return retry.Do[uint64](r.ctx, r.maxAttempts, r.strategy, func() (uint64, error) {
+		ctx, cancel := context.WithTimeout(r.ctx, r.timeout)
+		defer cancel()
+		var result hexutil.Uint64
+		err := r.client.CallContext(ctx, &result, "eth_getTransactionCount", addr, r.blockHash)
+		return uint64(result), err
+	})
+}
+
+func (r *RPCSource) Balance(addr common.Address) (*uint256.Int, error) {
+	return retry.Do[*uint256.Int](r.ctx, r.maxAttempts, r.strategy, func() (*uint256.Int, error) {
+		ctx, cancel := context.WithTimeout(r.ctx, r.timeout)
+		defer cancel()
+		var result hexutil.U256
+		err := r.client.CallContext(ctx, &result, "eth_getBalance", addr, r.blockHash)
+		return (*uint256.Int)(&result), err
+	})
+}
+
+func (r *RPCSource) StorageAt(addr common.Address, key common.Hash) (common.Hash, error) {
+	return retry.Do[common.Hash](r.ctx, r.maxAttempts, r.strategy, func() (common.Hash, error) {
+		ctx, cancel := context.WithTimeout(r.ctx, r.timeout)
+		defer cancel()
+		var result common.Hash
+		err := r.client.CallContext(ctx, &result, "eth_getStorageAt", addr, key, r.blockHash)
+		return result, err
+	})
+}
+
+func (r *RPCSource) Code(addr common.Address) ([]byte, error) {
+	return retry.Do[[]byte](r.ctx, r.maxAttempts, r.strategy, func() ([]byte, error) {
+		ctx, cancel := context.WithTimeout(r.ctx, r.timeout)
+		defer cancel()
+		var result hexutil.Bytes
+		err := r.client.CallContext(ctx, &result, "eth_getCode", addr, r.blockHash)
+		return result, err
+	})
+}
+
+// Close stops any ongoing RPC requests by cancelling the RPC context
+func (r *RPCSource) Close() {
+	r.cancel()
+}

op-chain-ops/script/forking/rpc_test.go

+package forking
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/ethereum-optimism/optimism/op-service/retry"
+
+	"github.com/ethereum/go-ethereum"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/common/hexutil"
+	"github.com/holiman/uint256"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+// MockRPCClient implements RPCClient interface for testing
+type MockRPCClient struct {
+	mock.Mock
+}
+
+func (m *MockRPCClient) CallContext(ctx context.Context, result any, method string, args ...any) error {
+	return m.Called(ctx, result, method, args).Error(0)
+}
+
+func TestRPCSourceInitialization(t *testing.T) {
+	mockClient := new(MockRPCClient)
+	expectedStateRoot := common.HexToHash("0x1234")
+	expectedBlockHash := common.HexToHash("0x5678")
+
+	t.Run("initialization by block number", func(t *testing.T) {
+		mockClient.On("CallContext", mock.Anything, mock.AnythingOfType("**forking.Header"),
+			"eth_getBlockByNumber", []any{hexutil.Uint64(123), false}).
+			Run(func(args mock.Arguments) {
+				result := args.Get(1).(**Header)
+				*result = &Header{
+					StateRoot: expectedStateRoot,
+					BlockHash: expectedBlockHash,
+				}
+			}).
+			Return(nil).Once()
+
+		source, err := RPCSourceByNumber("test_url", mockClient, 123)
+		require.NoError(t, err)
+		require.Equal(t, expectedStateRoot, source.StateRoot())
+		require.Equal(t, expectedBlockHash, source.BlockHash())
+	})
+
+	t.Run("initialization by block hash", func(t *testing.T) {
+		blockHash := common.HexToHash("0xabcd")
+		mockClient.On("CallContext", mock.Anything, mock.AnythingOfType("**forking.Header"),
+			"eth_getBlockByNumber", []any{blockHash, false}).
+			Run(func(args mock.Arguments) {
+				result := args.Get(1).(**Header)
+				*result = &Header{
+					StateRoot: expectedStateRoot,
+					BlockHash: expectedBlockHash,
+				}
+			}).
+			Return(nil).Once()
+
+		source, err := RPCSourceByHash("test_url", mockClient, blockHash)
+		require.NoError(t, err)
+		require.Equal(t, expectedStateRoot, source.StateRoot())
+		require.Equal(t, expectedBlockHash, source.BlockHash())
+	})
+
+	t.Run("initialization failure", func(t *testing.T) {
+		mockClient.On("CallContext", mock.Anything, mock.AnythingOfType("**forking.Header"),
+			"eth_getBlockByNumber", []any{hexutil.Uint64(999), false}).
+			Return(ethereum.NotFound).Times(2)
+
+		src := newRPCSource("test_url", mockClient)
+		strategy := retry.Exponential()
+		strategy.(*retry.ExponentialStrategy).Max = 100 * time.Millisecond
+		src.strategy = strategy
+		src.maxAttempts = 2
+		require.Error(t, src.init(hexutil.Uint64(999)))
+	})
+}
+
+func TestRPCSourceDataRetrieval(t *testing.T) {
+	mockClient := new(MockRPCClient)
+	testAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
+	blockHash := common.HexToHash("0xabcd")
+
+	source := &RPCSource{
+		blockHash:   blockHash,
+		client:      mockClient,
+		ctx:         context.Background(),
+		strategy:    retry.Exponential(),
+		maxAttempts: 10,
+		timeout:     time.Second * 10,
+	}
+
+	t.Run("get nonce", func(t *testing.T) {
+		expectedNonce := uint64(5)
+		mockClient.On("CallContext", mock.Anything, mock.AnythingOfType("*hexutil.Uint64"),
+			"eth_getTransactionCount", []any{testAddr, blockHash}).
+			Run(func(args mock.Arguments) {
+				result := args.Get(1).(*hexutil.Uint64)
+				*result = hexutil.Uint64(expectedNonce)
+			}).
+			Return(nil).Once()
+
+		nonce, err := source.Nonce(testAddr)
+		require.NoError(t, err)
+		require.Equal(t, expectedNonce, nonce)
+	})
+
+	t.Run("get balance", func(t *testing.T) {
+		expectedBalance := uint256.NewInt(1000)
+		mockClient.On("CallContext", mock.Anything, mock.AnythingOfType("*hexutil.U256"),
+			"eth_getBalance", []any{testAddr, blockHash}).
+			Run(func(args mock.Arguments) {
+				result := args.Get(1).(*hexutil.U256)
+				*(*uint256.Int)(result) = *expectedBalance
+			}).
+			Return(nil).Once()
+
+		balance, err := source.Balance(testAddr)
+		require.NoError(t, err)
+		require.Equal(t, expectedBalance, balance)
+	})
+
+	t.Run("get storage", func(t *testing.T) {
+		storageKey := common.HexToHash("0x1234")
+		expectedValue := common.HexToHash("0x5678")
+		mockClient.On("CallContext", mock.Anything, mock.AnythingOfType("*common.Hash"),
+			"eth_getStorageAt", []any{testAddr, storageKey, blockHash}).
+			Run(func(args mock.Arguments) {
+				result := args.Get(1).(*common.Hash)
+				*result = expectedValue
+			}).
+			Return(nil).Once()
+
+		value, err := source.StorageAt(testAddr, storageKey)
+		require.NoError(t, err)
+		require.Equal(t, expectedValue, value)
+	})
+
+	t.Run("get code", func(t *testing.T) {
+		expectedCode := []byte{1, 2, 3, 4}
+		mockClient.On("CallContext", mock.Anything, mock.AnythingOfType("*hexutil.Bytes"),
+			"eth_getCode", []any{testAddr, blockHash}).
+			Run(func(args mock.Arguments) {
+				result := args.Get(1).(*hexutil.Bytes)
+				*result = expectedCode
+			}).
+			Return(nil).Once()
+
+		code, err := source.Code(testAddr)
+		require.NoError(t, err)
+		require.Equal(t, expectedCode, code)
+	})
+}
+
+func TestRPCSourceRetry(t *testing.T) {
+	mockClient := new(MockRPCClient)
+	testAddr := common.HexToAddress("0x1234")
+	blockHash := common.HexToHash("0xabcd")
+	strategy := retry.Exponential()
+	strategy.(*retry.ExponentialStrategy).Max = 100 * time.Millisecond
+
+	source := &RPCSource{
+		blockHash:   blockHash,
+		client:      mockClient,
+		ctx:         context.Background(),
+		strategy:    strategy,
+		maxAttempts: 3,
+		timeout:     time.Second * 10,
+	}
+
+	t.Run("retry on temporary error", func(t *testing.T) {
+		tempError := errors.New("temporary network error")
+
+		// Fail twice, succeed on third attempt
+		mockClient.On("CallContext", mock.Anything, mock.AnythingOfType("*hexutil.Uint64"),
+			"eth_getTransactionCount", []any{testAddr, blockHash}).
+			Return(tempError).Times(2)
+
+		mockClient.On("CallContext", mock.Anything, mock.AnythingOfType("*hexutil.Uint64"),
+			"eth_getTransactionCount", []any{testAddr, blockHash}).
+			Run(func(args mock.Arguments) {
+				result := args.Get(1).(*hexutil.Uint64)
+				*result = hexutil.Uint64(5)
+			}).
+			Return(nil).Once()
+
+		nonce, err := source.Nonce(testAddr)
+		require.NoError(t, err)
+		require.Equal(t, uint64(5), nonce)
+	})
+}
+
+func TestRPCSourceClose(t *testing.T) {
+	mockClient := new(MockRPCClient)
+	source := newRPCSource("test_url", mockClient)
+
+	// Verify context is active before close
+	require.NoError(t, source.ctx.Err())
+
+	source.Close()
+
+	// Verify context is cancelled after close
+	require.Error(t, source.ctx.Err())
+	require.Equal(t, context.Canceled, source.ctx.Err())
+}

op-chain-ops/script/forking/state.go

+package forking
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/addresses"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/stateless"
+	"github.com/ethereum/go-ethereum/core/tracing"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/params"
+	"github.com/ethereum/go-ethereum/trie/utils"
+	"github.com/holiman/uint256"
+)
+
+type forkStateEntry struct {
+	state *state.StateDB
+}
+
+func (fe *forkStateEntry) DB() *ForkDB {
+	return fe.state.Database().(*ForkDB)
+}
+
+// ForkableState implements the vm.StateDB interface,
+// and a few other methods as defined in the VMStateDB interface.
+// This state can be forked in-place,
+// swapping over operations to route to in-memory states that wrap fork sources.
+type ForkableState struct {
+	selected VMStateDB
+
+	activeFork ForkID
+	forks      map[ForkID]*forkStateEntry
+
+	// persistent accounts will override any interactions
+	// to be directly with the forkID that was active at the time it was made persistent,
+	// rather than whatever fork is currently active.
+	persistent map[common.Address]ForkID
+
+	fallback VMStateDB
+
+	idCounter uint64
+}
+
+var _ VMStateDB = (*ForkableState)(nil)
+
+func NewForkableState(base VMStateDB) *ForkableState {
+	return &ForkableState{
+		selected:   base,
+		activeFork: ForkID{},
+		forks:      make(map[ForkID]*forkStateEntry),
+		persistent: map[common.Address]ForkID{
+			addresses.DefaultSenderAddr: ForkID{},
+			addresses.VMAddr:            ForkID{},
+			addresses.ConsoleAddr:       ForkID{},
+			addresses.ScriptDeployer:    ForkID{},
+			addresses.ForgeDeployer:     ForkID{},
+		},
+		fallback:  base,
+		idCounter: 0,
+	}
+}
+
+// ExportDiff exports a state diff. Warning: diffs are like flushed states.
+// So we flush the state, making all the contents cold, losing transient storage, etc.
+func (fst *ForkableState) ExportDiff(id ForkID) (*ExportDiff, error) {
+	if id == (ForkID{}) {
+		return nil, errors.New("default no-fork state does not have an exportable diff")
+	}
+	f, ok := fst.forks[id]
+	if !ok {
+		return nil, fmt.Errorf("unknown fork %q", id)
+	}
+	// Finalize the state content, so we can get an accurate diff.
+	f.state.IntermediateRoot(true)
+	tr := f.state.GetTrie()
+	ft, ok := tr.(*ForkedAccountsTrie)
+	if !ok {
+		return nil, fmt.Errorf("forked state trie is unexpectedly not a ForkedAccountsTrie: %T", tr)
+	}
+	diff := ft.ExportDiff()
+	// Now re-init the state, so we can use it again (albeit it cold).
+	forkDB := &ForkDB{active: ft}
+	st, err := state.New(forkDB.active.stateRoot, forkDB)
+	if err != nil {
+		return nil, fmt.Errorf("failed to construct fork state: %w", err)
+	}
+	fst.forks[id].state = st
+	if fst.activeFork == id {
+		fst.selected = st
+	}
+	return diff, nil
+}
+
+// CreateSelectFork is like vm.createSelectFork, it creates a fork, and selects it immediately.
+func (fst *ForkableState) CreateSelectFork(source ForkSource) (ForkID, error) {
+	id, err := fst.CreateFork(source)
+	if err != nil {
+		return id, err
+	}
+	return id, fst.SelectFork(id)
+}
+
+// CreateFork is like vm.createFork, it creates a fork, but does not select it yet.
+func (fst *ForkableState) CreateFork(source ForkSource) (ForkID, error) {
+	fst.idCounter += 1 // increment first, don't use ID 0
+	id := ForkID(*uint256.NewInt(fst.idCounter))
+	_, ok := fst.forks[id]
+	if ok { // sanity check our ID counter is consistent with the tracked forks
+		return id, fmt.Errorf("cannot create fork, fork %q already exists", id)
+	}
+	forkDB := NewForkDB(source)
+	st, err := state.New(forkDB.active.stateRoot, forkDB)
+	if err != nil {
+		return id, fmt.Errorf("failed to construct fork state: %w", err)
+	}
+	fst.forks[id] = &forkStateEntry{
+		state: st,
+	}
+	return id, nil
+}
+
+// SelectFork is like vm.selectFork, it activates the usage of a previously created fork.
+func (fst *ForkableState) SelectFork(id ForkID) error {
+	if id == (ForkID{}) {
+		fst.selected = fst.fallback
+		fst.activeFork = ForkID{}
+		return nil
+	}
+	f, ok := fst.forks[id]
+	if !ok {
+		return fmt.Errorf("cannot select fork, fork %q is unknown", id)
+	}
+	fst.selected = f.state
+	fst.activeFork = id
+	return nil
+}
+
+// ResetFork resets the fork to be coupled to the given fork-source.
+// Any ephemeral state changes (transient storage, warm s-loads, etc.)
+// as well as any uncommitted state, as well as any previously flushed diffs, will be lost.
+func (fst *ForkableState) ResetFork(id ForkID, src ForkSource) error {
+	if id == (ForkID{}) {
+		return errors.New("default no-fork state cannot change its ForkSource")
+	}
+	f, ok := fst.forks[id]
+	if !ok {
+		return fmt.Errorf("unknown fork %q", id)
+	}
+	// Now create a new state
+	forkDB := NewForkDB(src)
+	st, err := state.New(src.StateRoot(), forkDB)
+	if err != nil {
+		return fmt.Errorf("failed to construct fork state: %w", err)
+	}
+	f.state = st
+	if fst.activeFork == id {
+		fst.selected = st
+	}
+	return nil
+}
+
+// ActiveFork returns the ID current active fork, or active == false if no fork is active.
+func (fst *ForkableState) ActiveFork() (id ForkID, active bool) {
+	return fst.activeFork, fst.activeFork != (ForkID{})
+}
+
+// ForkURLOrAlias returns the URL or alias that the fork was configured with as source.
+// Returns an error if no fork is active
+func (fst *ForkableState) ForkURLOrAlias(id ForkID) (string, error) {
+	if id == (ForkID{}) {
+		return "", errors.New("default no-fork state does not have an URL or Alias")
+	}
+	f, ok := fst.forks[id]
+	if !ok {
+		return "", fmt.Errorf("unknown fork %q", id)
+	}
+	return f.DB().active.src.URLOrAlias(), nil
+}
+
+// SubstituteBaseState substitutes in a fallback state.
+func (fst *ForkableState) SubstituteBaseState(base VMStateDB) {
+	fst.fallback = base
+	// If the fallback is currently selected, also updated the fallback.
+	if fst.activeFork == (ForkID{}) {
+		fst.selected = base
+	}
+}
+
+// MakePersistent is like vm.makePersistent, it maintains this account context across all forks.
+// It does not make the account of a fork persistent, it makes an account override what might be in a fork.
+func (fst *ForkableState) MakePersistent(addr common.Address) {
+	fst.persistent[addr] = fst.activeFork
+}
+
+// RevokePersistent is like vm.revokePersistent, it undoes a previous vm.makePersistent.
+func (fst *ForkableState) RevokePersistent(addr common.Address) {
+	delete(fst.persistent, addr)
+}
+
+// IsPersistent is like vm.isPersistent, it checks if an account persists across forks.
+func (fst *ForkableState) IsPersistent(addr common.Address) bool {
+	_, ok := fst.persistent[addr]
+	return ok
+}
+
+func (fst *ForkableState) stateFor(addr common.Address) VMStateDB {
+	// if forked, check if we persisted this account across forks
+	persistedForkID, ok := fst.persistent[addr]
+	if ok {
+		if persistedForkID == (ForkID{}) {
+			return fst.fallback
+		}
+		return fst.forks[persistedForkID].state
+	}
+	// This may be the fallback state, if no fork is active.
+	return fst.selected
+}
+
+// Finalise finalises the state by removing the destructed objects and clears
+// the journal as well as the refunds. Finalise, however, will not push any updates
+// into the tries just yet.
+//
+// The changes will be flushed to the underlying DB.
+// A *ForkDB if the state is currently forked.
+func (fst *ForkableState) Finalise(deleteEmptyObjects bool) {
+	fst.selected.Finalise(deleteEmptyObjects)
+}
+
+func (fst *ForkableState) CreateAccount(address common.Address) {
+	fst.stateFor(address).CreateAccount(address)
+}
+
+func (fst *ForkableState) CreateContract(address common.Address) {
+	fst.stateFor(address).CreateContract(address)
+}
+
+func (fst *ForkableState) SubBalance(address common.Address, u *uint256.Int, reason tracing.BalanceChangeReason) {
+	fst.stateFor(address).SubBalance(address, u, reason)
+}
+
+func (fst *ForkableState) AddBalance(address common.Address, u *uint256.Int, reason tracing.BalanceChangeReason) {
+	fst.stateFor(address).AddBalance(address, u, reason)
+}
+
+func (fst *ForkableState) GetBalance(address common.Address) *uint256.Int {
+	return fst.stateFor(address).GetBalance(address)
+}
+
+func (fst *ForkableState) GetNonce(address common.Address) uint64 {
+	return fst.stateFor(address).GetNonce(address)
+}
+
+func (fst *ForkableState) SetNonce(address common.Address, u uint64) {
+	fst.stateFor(address).SetNonce(address, u)
+}
+
+func (fst *ForkableState) GetCodeHash(address common.Address) common.Hash {
+	return fst.stateFor(address).GetCodeHash(address)
+}
+
+func (fst *ForkableState) GetCode(address common.Address) []byte {
+	return fst.stateFor(address).GetCode(address)
+}
+
+func (fst *ForkableState) SetCode(address common.Address, bytes []byte) {
+	fst.stateFor(address).SetCode(address, bytes)
+}
+
+func (fst *ForkableState) GetCodeSize(address common.Address) int {
+	return fst.stateFor(address).GetCodeSize(address)
+}
+
+func (fst *ForkableState) AddRefund(u uint64) {
+	fst.selected.AddRefund(u)
+}
+
+func (fst *ForkableState) SubRefund(u uint64) {
+	fst.selected.SubRefund(u)
+}
+
+func (fst *ForkableState) GetRefund() uint64 {
+	return fst.selected.GetRefund()
+}
+
+func (fst *ForkableState) GetCommittedState(address common.Address, hash common.Hash) common.Hash {
+	return fst.stateFor(address).GetCommittedState(address, hash)
+}
+
+func (fst *ForkableState) GetState(address common.Address, k common.Hash) common.Hash {
+	return fst.stateFor(address).GetState(address, k)
+}
+
+func (fst *ForkableState) SetState(address common.Address, k common.Hash, v common.Hash) {
+	fst.stateFor(address).SetState(address, k, v)
+}
+
+func (fst *ForkableState) GetStorageRoot(addr common.Address) common.Hash {
+	return fst.stateFor(addr).GetStorageRoot(addr)
+}
+
+func (fst *ForkableState) GetTransientState(addr common.Address, key common.Hash) common.Hash {
+	return fst.stateFor(addr).GetTransientState(addr, key)
+}
+
+func (fst *ForkableState) SetTransientState(addr common.Address, key, value common.Hash) {
+	fst.stateFor(addr).SetTransientState(addr, key, value)
+}
+
+func (fst *ForkableState) SelfDestruct(address common.Address) {
+	fst.stateFor(address).SelfDestruct(address)
+}
+
+func (fst *ForkableState) HasSelfDestructed(address common.Address) bool {
+	return fst.stateFor(address).HasSelfDestructed(address)
+}
+
+func (fst *ForkableState) Selfdestruct6780(address common.Address) {
+	fst.stateFor(address).Selfdestruct6780(address)
+}
+
+func (fst *ForkableState) Exist(address common.Address) bool {
+	return fst.stateFor(address).Exist(address)
+}
+
+func (fst *ForkableState) Empty(address common.Address) bool {
+	return fst.stateFor(address).Empty(address)
+}
+
+func (fst *ForkableState) AddressInAccessList(addr common.Address) bool {
+	return fst.stateFor(addr).AddressInAccessList(addr)
+}
+
+func (fst *ForkableState) SlotInAccessList(addr common.Address, slot common.Hash) (addressOk bool, slotOk bool) {
+	return fst.stateFor(addr).SlotInAccessList(addr, slot)
+}
+
+func (fst *ForkableState) AddAddressToAccessList(addr common.Address) {
+	fst.stateFor(addr).AddAddressToAccessList(addr)
+}
+
+func (fst *ForkableState) AddSlotToAccessList(addr common.Address, slot common.Hash) {
+	fst.stateFor(addr).AddSlotToAccessList(addr, slot)
+}
+
+func (fst *ForkableState) PointCache() *utils.PointCache {
+	return fst.selected.PointCache()
+}
+
+func (fst *ForkableState) Prepare(rules params.Rules, sender, coinbase common.Address, dest *common.Address, precompiles []common.Address, txAccesses types.AccessList) {
+	fst.selected.Prepare(rules, sender, coinbase, dest, precompiles, txAccesses)
+}
+
+func (fst *ForkableState) RevertToSnapshot(i int) {
+	fst.selected.RevertToSnapshot(i)
+}
+
+func (fst *ForkableState) Snapshot() int {
+	return fst.selected.Snapshot()
+}
+
+func (fst *ForkableState) AddLog(log *types.Log) {
+	fst.selected.AddLog(log)
+}
+
+func (fst *ForkableState) AddPreimage(hash common.Hash, img []byte) {
+	fst.selected.AddPreimage(hash, img)
+}
+
+func (fst *ForkableState) Witness() *stateless.Witness {
+	return fst.selected.Witness()
+}
+
+func (fst *ForkableState) SetBalance(addr common.Address, amount *uint256.Int, reason tracing.BalanceChangeReason) {
+	fst.stateFor(addr).SetBalance(addr, amount, reason)
+}

op-chain-ops/script/forking/trie.go

+package forking
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/holiman/uint256"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/ethdb"
+	"github.com/ethereum/go-ethereum/trie"
+	"github.com/ethereum/go-ethereum/trie/trienode"
+)
+
+type ForkedAccountsTrie struct {
+	// stateRoot that this diff is based on top of
+	stateRoot common.Hash
+
+	// source to retrieve data from when it's not in the diff
+	src ForkSource
+
+	diff *ExportDiff
+}
+
+var _ state.Trie = (*ForkedAccountsTrie)(nil)
+
+func (f *ForkedAccountsTrie) Copy() *ForkedAccountsTrie {
+	return &ForkedAccountsTrie{
+		stateRoot: f.stateRoot,
+		diff:      f.diff.Copy(),
+	}
+}
+
+func (f *ForkedAccountsTrie) ExportDiff() *ExportDiff {
+	return f.diff.Copy()
+}
+
+func (f *ForkedAccountsTrie) HasDiff() bool {
+	return len(f.diff.Code) > 0 || len(f.diff.Account) > 0
+}
+
+// ClearDiff clears the flushed changes. This does not clear the warm state changes.
+// To fully clear, first Finalise the forked state that uses this trie, and then clear the diff.
+func (f *ForkedAccountsTrie) ClearDiff() {
+	f.diff.Clear()
+}
+
+// ContractCode is not directly part of the vm.State interface,
+// but is used by the ForkDB to retrieve the contract code.
+func (f *ForkedAccountsTrie) ContractCode(addr common.Address, codeHash common.Hash) ([]byte, error) {
+	diffAcc, ok := f.diff.Account[addr]
+	if ok {
+		if diffAcc.CodeHash != nil && *diffAcc.CodeHash != codeHash {
+			return nil, fmt.Errorf("account code changed to %s, cannot get code %s of account %s", *diffAcc.CodeHash, codeHash, addr)
+		}
+		if code, ok := f.diff.Code[codeHash]; ok {
+			return code, nil
+		}
+		// if not in codeDiff, the actual code has not changed.
+	}
+	code, err := f.src.Code(addr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to retrieve code: %w", err)
+	}
+	// sanity-check the retrieved code matches the expected codehash
+	if h := crypto.Keccak256Hash(code); h != codeHash {
+		return nil, fmt.Errorf("retrieved code of %s hashed to %s, but expected %s", addr, h, codeHash)
+	}
+	return code, nil
+}
+
+// ContractCodeSize is not directly part of the vm.State interface,
+// but is used by the ForkDB to retrieve the contract code-size.
+func (f *ForkedAccountsTrie) ContractCodeSize(addr common.Address, codeHash common.Hash) (int, error) {
+	code, err := f.ContractCode(addr, codeHash)
+	if err != nil {
+		return 0, fmt.Errorf("cannot get contract code to determine code size: %w", err)
+	}
+	return len(code), nil
+}
+
+func (f *ForkedAccountsTrie) GetKey(bytes []byte) []byte {
+	panic("arbitrary key lookups on ForkedAccountsTrie are not supported")
+}
+
+func (f *ForkedAccountsTrie) GetAccount(address common.Address) (*types.StateAccount, error) {
+	acc := &types.StateAccount{
+		Nonce:    0,
+		Balance:  nil,
+		Root:     fakeRoot,
+		CodeHash: nil,
+	}
+	diffAcc := f.diff.Account[address]
+	if diffAcc != nil && diffAcc.Nonce != nil {
+		acc.Nonce = *diffAcc.Nonce
+	} else {
+		v, err := f.src.Nonce(address)
+		if err != nil {
+			return nil, fmt.Errorf("failed to retrieve nonce of account %s: %w", address, err)
+		}
+		acc.Nonce = v
+	}
+	if diffAcc != nil && diffAcc.Balance != nil {
+		acc.Balance = new(uint256.Int).Set(diffAcc.Balance)
+	} else {
+		v, err := f.src.Balance(address)
+		if err != nil {
+			return nil, fmt.Errorf("failed to retrieve balance of account %s: %w", address, err)
+		}
+		acc.Balance = new(uint256.Int).Set(v)
+	}
+	if diffAcc != nil && diffAcc.CodeHash != nil {
+		cpy := *diffAcc.CodeHash
+		acc.CodeHash = cpy.Bytes()
+	} else {
+		v, err := f.src.Code(address)
+		if err != nil {
+			return nil, fmt.Errorf("failed to retrieve code of account %s: %w", address, err)
+		}
+		acc.CodeHash = crypto.Keccak256Hash(v).Bytes()
+	}
+	return acc, nil
+}
+
+func (f *ForkedAccountsTrie) GetStorage(addr common.Address, key []byte) ([]byte, error) {
+	k := common.BytesToHash(key)
+	diffAcc, ok := f.diff.Account[addr]
+	if ok { // if there is a diff, try and see if it contains a storage diff
+		v, ok := diffAcc.Storage[k]
+		if ok { // if the storage has changed, return that change
+			return v.Bytes(), nil
+		}
+	}
+	v, err := f.src.StorageAt(addr, k)
+	if err != nil {
+		return nil, err
+	}
+	return v.Bytes(), nil
+}
+
+func (f *ForkedAccountsTrie) UpdateAccount(address common.Address, account *types.StateAccount, codeLen int) error {
+	// Ignored, account contains the code details we need.
+	// Also see the trie.StateTrie of geth itself, which ignores this arg too.
+	_ = codeLen
+
+	nonce := account.Nonce
+	b := account.Balance.Clone()
+	codeHash := common.BytesToHash(account.CodeHash)
+	out := &AccountDiff{
+		Nonce:    &nonce,
+		Balance:  b,
+		Storage:  nil,
+		CodeHash: &codeHash,
+	}
+	// preserve the storage diff
+	if diffAcc, ok := f.diff.Account[address]; ok {
+		out.Storage = diffAcc.Storage
+	}
+	f.diff.Account[address] = out
+	return nil
+}
+
+func (f *ForkedAccountsTrie) UpdateStorage(addr common.Address, key, value []byte) error {
+	diffAcc, ok := f.diff.Account[addr]
+	if !ok {
+		diffAcc = &AccountDiff{}
+		f.diff.Account[addr] = diffAcc
+	}
+	if diffAcc.Storage == nil {
+		diffAcc.Storage = make(map[common.Hash]common.Hash)
+	}
+	k := common.BytesToHash(key)
+	v := common.BytesToHash(value)
+	diffAcc.Storage[k] = v
+	return nil
+}
+
+func (f *ForkedAccountsTrie) DeleteAccount(address common.Address) error {
+	f.diff.Account[address] = nil
+	return nil
+}
+
+func (f *ForkedAccountsTrie) DeleteStorage(addr common.Address, key []byte) error {
+	return f.UpdateStorage(addr, key, nil)
+}
+
+func (f *ForkedAccountsTrie) UpdateContractCode(addr common.Address, codeHash common.Hash, code []byte) error {
+	diffAcc, ok := f.diff.Account[addr]
+	if !ok {
+		diffAcc = &AccountDiff{}
+		f.diff.Account[addr] = diffAcc
+	}
+	diffAcc.CodeHash = &codeHash
+	f.diff.Code[codeHash] = code
+	return nil
+}
+
+func (f *ForkedAccountsTrie) Hash() common.Hash {
+	return f.stateRoot
+}
+
+func (f *ForkedAccountsTrie) Commit(collectLeaf bool) (common.Hash, *trienode.NodeSet) {
+	panic("cannot commit state-changes of a forked trie")
+}
+
+func (f *ForkedAccountsTrie) Witness() map[string]struct{} {
+	panic("witness generation of a ForkedAccountsTrie is not supported")
+}
+
+func (f *ForkedAccountsTrie) NodeIterator(startKey []byte) (trie.NodeIterator, error) {
+	return nil, errors.New("node iteration of a ForkedAccountsTrie is not supported")
+}
+
+func (f *ForkedAccountsTrie) Prove(key []byte, proofDb ethdb.KeyValueWriter) error {
+	return errors.New("proving of a ForkedAccountsTrie is not supported")
+}
+
+func (f *ForkedAccountsTrie) IsVerkle() bool {
+	return false
+}

op-chain-ops/script/forking/trie_test.go

+package forking
+
+import (
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/holiman/uint256"
+	"github.com/stretchr/testify/require"
+)
+
+func setupTrie(t *testing.T) (*ForkedAccountsTrie, *MockForkSource) {
+	mockSource := new(MockForkSource)
+	stateRoot := common.HexToHash("0x1234")
+	mockSource.On("StateRoot").Return(stateRoot)
+
+	trie := &ForkedAccountsTrie{
+		stateRoot: stateRoot,
+		src:       mockSource,
+		diff:      NewExportDiff(),
+	}
+	return trie, mockSource
+}
+
+func TestForkedAccountsTrie_GetAccount(t *testing.T) {
+	trie, mockSource := setupTrie(t)
+	addr := common.HexToAddress("0x1234")
+
+	// Setup mock responses
+	expectedNonce := uint64(1)
+	expectedBalance := uint256.NewInt(100)
+	expectedCode := []byte{1, 2, 3, 4}
+	expectedCodeHash := crypto.Keccak256Hash(expectedCode)
+
+	mockSource.On("Nonce", addr).Return(expectedNonce, nil)
+	mockSource.On("Balance", addr).Return(expectedBalance, nil)
+	mockSource.On("Code", addr).Return(expectedCode, nil)
+
+	// Test initial account retrieval
+	account, err := trie.GetAccount(addr)
+	require.NoError(t, err)
+	require.Equal(t, expectedNonce, account.Nonce)
+	require.Equal(t, expectedBalance, uint256.NewInt(0).SetBytes(account.Balance.Bytes()))
+	require.Equal(t, expectedCodeHash.Bytes(), account.CodeHash)
+
+	// Update account and verify diff
+	newNonce := uint64(2)
+	newBalance := uint256.NewInt(200)
+	account.Nonce = newNonce
+	account.Balance = newBalance
+
+	err = trie.UpdateAccount(addr, account, 0)
+	require.NoError(t, err)
+
+	// Verify updated account
+	updatedAccount, err := trie.GetAccount(addr)
+	require.NoError(t, err)
+	require.Equal(t, newNonce, updatedAccount.Nonce)
+	require.Equal(t, newBalance, uint256.NewInt(0).SetBytes(updatedAccount.Balance.Bytes()))
+}
+
+func TestForkedAccountsTrie_Storage(t *testing.T) {
+	trie, mockSource := setupTrie(t)
+	addr := common.HexToAddress("0x1234")
+	key := common.HexToHash("0x1")
+	value := common.HexToHash("0x2")
+
+	// Setup mock for initial storage value
+	mockSource.On("StorageAt", addr, key).Return(value, nil)
+
+	// Test initial storage retrieval
+	storageValue, err := trie.GetStorage(addr, key.Bytes())
+	require.NoError(t, err)
+	require.Equal(t, value.Bytes(), storageValue)
+
+	// Update storage
+	newValue := common.HexToHash("0x3")
+	err = trie.UpdateStorage(addr, key.Bytes(), newValue.Bytes())
+	require.NoError(t, err)
+
+	// Verify updated storage
+	updatedValue, err := trie.GetStorage(addr, key.Bytes())
+	require.NoError(t, err)
+	require.Equal(t, newValue.Bytes(), updatedValue)
+}
+
+func TestForkedAccountsTrie_ContractCode(t *testing.T) {
+	trie, mockSource := setupTrie(t)
+	addr := common.HexToAddress("0x1234")
+	code := []byte{1, 2, 3, 4}
+	codeHash := crypto.Keccak256Hash(code)
+
+	// Setup mock for code retrieval
+	mockSource.On("Code", addr).Return(code, nil)
+
+	// Test initial code retrieval
+	retrievedCode, err := trie.ContractCode(addr, codeHash)
+	require.NoError(t, err)
+	require.Equal(t, code, retrievedCode)
+
+	// Update code
+	newCode := []byte{5, 6, 7, 8}
+	newCodeHash := crypto.Keccak256Hash(newCode)
+
+	err = trie.UpdateContractCode(addr, newCodeHash, newCode)
+	require.NoError(t, err)
+
+	// Verify updated code
+	updatedCode, err := trie.ContractCode(addr, newCodeHash)
+	require.NoError(t, err)
+	require.Equal(t, newCode, updatedCode)
+}
+
+func TestForkedAccountsTrie_DeleteAccount(t *testing.T) {
+	trie, _ := setupTrie(t)
+	addr := common.HexToAddress("0x1234")
+
+	// Setup initial account
+	account := &types.StateAccount{
+		Nonce:    1,
+		Balance:  uint256.NewInt(100),
+		CodeHash: crypto.Keccak256([]byte{1, 2, 3, 4}),
+	}
+
+	err := trie.UpdateAccount(addr, account, 0)
+	require.NoError(t, err)
+
+	// Delete account
+	err = trie.DeleteAccount(addr)
+	require.NoError(t, err)
+
+	// Verify account is marked as deleted in diff
+	require.Nil(t, trie.diff.Account[addr])
+}
+
+func TestForkedAccountsTrie_Copy(t *testing.T) {
+	trie, _ := setupTrie(t)
+	addr := common.HexToAddress("0x1234")
+
+	// Setup some initial state
+	account := &types.StateAccount{
+		Nonce:    1,
+		Balance:  uint256.NewInt(100),
+		CodeHash: crypto.Keccak256([]byte{1, 2, 3, 4}),
+	}
+	err := trie.UpdateAccount(addr, account, 0)
+	require.NoError(t, err)
+
+	// Make a copy
+	cpy := trie.Copy()
+
+	// Verify copy has same state
+	require.Equal(t, trie.stateRoot, cpy.stateRoot)
+	require.Equal(t, trie.diff.Account[addr].Nonce, cpy.diff.Account[addr].Nonce)
+	require.True(t, trie.diff.Account[addr].Balance.Eq(cpy.diff.Account[addr].Balance))
+
+	// Modify copy and verify original is unchanged
+	newAccount := &types.StateAccount{
+		Nonce:    2,
+		Balance:  uint256.NewInt(200),
+		CodeHash: crypto.Keccak256([]byte{5, 6, 7, 8}),
+	}
+	err = cpy.UpdateAccount(addr, newAccount, 0)
+	require.NoError(t, err)
+
+	originalAccount, err := trie.GetAccount(addr)
+	require.NoError(t, err)
+	require.Equal(t, uint64(1), originalAccount.Nonce)
+	require.True(t, uint256.NewInt(100).Eq(uint256.NewInt(0).SetBytes(originalAccount.Balance.Bytes())))
+}
+
+func TestForkedAccountsTrie_HasDiff(t *testing.T) {
+	trie, _ := setupTrie(t)
+
+	// Initially no diff
+	require.False(t, trie.HasDiff())
+
+	// Add account change
+	addr := common.HexToAddress("0x1234")
+	account := &types.StateAccount{
+		Nonce:    1,
+		Balance:  uint256.NewInt(100),
+		CodeHash: crypto.Keccak256([]byte{1, 2, 3, 4}),
+	}
+	err := trie.UpdateAccount(addr, account, 0)
+	require.NoError(t, err)
+
+	// Verify diff exists
+	require.True(t, trie.HasDiff())
+
+	// Clear diff
+	trie.ClearDiff()
+	require.False(t, trie.HasDiff())
+}
+
+func TestForkedAccountsTrie_UnsupportedOperations(t *testing.T) {
+	trie, _ := setupTrie(t)
+
+	require.Panics(t, func() { trie.GetKey([]byte{1, 2, 3}) })
+	require.Panics(t, func() { trie.Commit(false) })
+	require.Panics(t, func() { trie.Witness() })
+
+	_, err := trie.NodeIterator(nil)
+	require.Error(t, err)
+
+	err = trie.Prove(nil, nil)
+	require.Error(t, err)
+}

op-chain-ops/script/prank.go

@@ -8,6 +8,8 @@ import (
 	"fmt"
 	"math/big"
 
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/addresses"
+
 	"github.com/holiman/uint256"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -53,7 +55,7 @@ func (h *Host) handleCaller(caller vm.ContractRef) vm.ContractRef {
 	// apply prank, if top call-frame had set up a prank
 	if len(h.callStack) > 0 {
 		parentCallFrame := h.callStack[len(h.callStack)-1]
-		if parentCallFrame.Prank != nil && caller.Address() != VMAddr { // pranks do not apply to the cheatcode precompile
+		if parentCallFrame.Prank != nil && caller.Address() != addresses.VMAddr { // pranks do not apply to the cheatcode precompile
 			if parentCallFrame.Prank.Broadcast && parentCallFrame.LastOp == vm.CREATE2 && h.useCreate2Deployer {
 				return &prankRef{
 					prank: DeterministicDeployerAddress,

op-chain-ops/script/precompile.go

@@ -9,6 +9,8 @@ import (
 	"reflect"
 	"strings"
 
+	"github.com/holiman/uint256"
+
 	"github.com/ethereum/go-ethereum/accounts/abi"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
@@ -352,6 +354,8 @@ type ABIInt256 big.Int
 
 var abiInt256Type = typeFor[ABIInt256]()
 
+var abiUint256Type = typeFor[uint256.Int]()
+
 // goTypeToSolidityType converts a Go type to the solidity ABI type definition.
 // The "internalType" is a quirk of the Geth ABI utils, for nested structures.
 // Unfortunately we have to convert to string, not directly to ABI type structure,
@@ -364,6 +368,9 @@ func goTypeToSolidityType(typ reflect.Type) (typeDef, internalType string, err e
 		reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
 		return strings.ToLower(typ.Kind().String()), "", nil
 	case reflect.Array:
+		if typ.AssignableTo(abiUint256Type) { // uint256.Int underlying Go type is [4]uint64
+			return "uint256", "", nil
+		}
 		if typ.Elem().Kind() == reflect.Uint8 {
 			if typ.Len() == 20 && typ.Name() == "Address" {
 				return "address", "", nil

op-chain-ops/script/script.go

@@ -4,9 +4,12 @@ import (
 	"bytes"
 	"encoding/binary"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/big"
 
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/addresses"
+
 	"github.com/holiman/uint256"
 
 	"github.com/ethereum/go-ethereum/accounts/abi"
@@ -19,13 +22,13 @@ import (
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/crypto"
-	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/triedb"
 	"github.com/ethereum/go-ethereum/triedb/hashdb"
 
 	"github.com/ethereum-optimism/optimism/op-chain-ops/foundry"
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/forking"
 	"github.com/ethereum-optimism/optimism/op-chain-ops/srcmap"
 )
 
@@ -72,9 +75,12 @@ type Host struct {
 	af       *foundry.ArtifactsFS
 	chainCfg *params.ChainConfig
 	env      *vm.EVM
-	state    *state.StateDB
-	stateDB  state.Database
-	rawDB    ethdb.Database
+
+	state     *forking.ForkableState
+	baseState *state.StateDB
+
+	// only known contracts may utilize cheatcodes and logging
+	allowedCheatcodes map[common.Address]struct{}
 
 	cheatcodes *Precompile[*CheatCodesPrecompile]
 	console    *Precompile[*ConsolePrecompile]
@@ -117,6 +123,7 @@ type BroadcastHook func(broadcast Broadcast)
 
 type Hooks struct {
 	OnBroadcast BroadcastHook
+	OnFork      ForkHook
 }
 
 func WithBroadcastHook(hook BroadcastHook) HostOption {
@@ -125,6 +132,12 @@ func WithBroadcastHook(hook BroadcastHook) HostOption {
 	}
 }
 
+func WithForkHook(hook ForkHook) HostOption {
+	return func(h *Host) {
+		h.hooks.OnFork = hook
+	}
+}
+
 // WithIsolatedBroadcasts makes each broadcast clean the context,
 // by flushing the dirty storage changes, and preparing the ephemeral state again.
 // This then produces more accurate gas estimation for broadcast calls.
@@ -167,7 +180,11 @@ func NewHost(
 		srcMaps:          make(map[common.Address]*srcmap.SourceMap),
 		hooks: &Hooks{
 			OnBroadcast: func(broadcast Broadcast) {},
+			OnFork: func(opts *ForkConfig) (forking.ForkSource, error) {
+				return nil, errors.New("no forking configured")
+			},
 		},
+		allowedCheatcodes: make(map[common.Address]struct{}),
 	}
 
 	for _, opt := range options {
@@ -212,18 +229,19 @@ func NewHost(
 	}
 
 	// Create an in-memory database, to host our temporary script state changes
-	h.rawDB = rawdb.NewMemoryDatabase()
-	h.stateDB = state.NewDatabase(triedb.NewDatabase(h.rawDB, &triedb.Config{
+	rawDB := rawdb.NewMemoryDatabase()
+	stateDB := state.NewDatabase(triedb.NewDatabase(rawDB, &triedb.Config{
 		Preimages: true, // To be able to iterate the state we need the Preimages
 		IsVerkle:  false,
 		HashDB:    hashdb.Defaults,
 		PathDB:    nil,
 	}), nil)
 	var err error
-	h.state, err = state.New(types.EmptyRootHash, h.stateDB)
+	h.baseState, err = state.New(types.EmptyRootHash, stateDB)
 	if err != nil {
 		panic(fmt.Errorf("failed to create memory state db: %w", err))
 	}
+	h.state = forking.NewForkableState(h.baseState)
 
 	// Initialize a block-context for the EVM to access environment variables.
 	// The block context (after embedding inside of the EVM environment) may be mutated later.
@@ -252,7 +270,7 @@ func NewHost(
 		GasPrice:     big.NewInt(0),
 		BlobHashes:   executionContext.BlobHashes,
 		BlobFeeCap:   big.NewInt(0),
-		AccessEvents: state.NewAccessEvents(h.stateDB.PointCache()),
+		AccessEvents: state.NewAccessEvents(h.baseState.PointCache()),
 	}
 
 	// Hook up the Host to capture the EVM environment changes
@@ -278,6 +296,18 @@ func NewHost(
 	return h
 }
 
+// AllowCheatcodes allows the given address to utilize the cheatcodes and logging precompiles
+func (h *Host) AllowCheatcodes(addr common.Address) {
+	h.log.Debug("Allowing cheatcodes", "address", addr, "label", h.labels[addr])
+	h.allowedCheatcodes[addr] = struct{}{}
+}
+
+// AllowedCheatcodes returns whether the given address is allowed to use cheatcodes
+func (h *Host) AllowedCheatcodes(addr common.Address) bool {
+	_, ok := h.allowedCheatcodes[addr]
+	return ok
+}
+
 // EnableCheats enables the Forge/HVM cheat-codes precompile and the Hardhat-style console2 precompile.
 func (h *Host) EnableCheats() error {
 	vmPrecompile, err := NewPrecompile[*CheatCodesPrecompile](&CheatCodesPrecompile{h: h})
@@ -288,8 +318,8 @@ func (h *Host) EnableCheats() error {
 	// Solidity does EXTCODESIZE checks on functions without return-data.
 	// We need to insert some placeholder code to prevent it from aborting calls.
 	// Emulates Forge script: https://github.com/foundry-rs/foundry/blob/224fe9cbf76084c176dabf7d3b2edab5df1ab818/crates/evm/evm/src/executors/mod.rs#L108
-	h.state.SetCode(VMAddr, []byte{0x00})
-	h.precompiles[VMAddr] = h.cheatcodes
+	h.state.SetCode(addresses.VMAddr, []byte{0x00})
+	h.precompiles[addresses.VMAddr] = h.cheatcodes
 
 	consolePrecompile, err := NewPrecompile[*ConsolePrecompile](&ConsolePrecompile{
 		logger: h.log,
@@ -299,7 +329,7 @@ func (h *Host) EnableCheats() error {
 		return fmt.Errorf("failed to init console precompile: %w", err)
 	}
 	h.console = consolePrecompile
-	h.precompiles[ConsoleAddr] = h.console
+	h.precompiles[addresses.ConsoleAddr] = h.console
 	// The Console precompile does not need bytecode,
 	// calls all go through a console lib, which avoids the EXTCODESIZE.
 	return nil
@@ -414,7 +444,7 @@ func (h *Host) ImportAccount(addr common.Address, account types.Account) {
 // getPrecompile overrides any accounts during runtime, to insert special precompiles, if activated.
 func (h *Host) getPrecompile(rules params.Rules, original vm.PrecompiledContract, addr common.Address) vm.PrecompiledContract {
 	if p, ok := h.precompiles[addr]; ok {
-		return p
+		return &AccessControlledPrecompile{h: h, inner: p}
 	}
 	return original
 }
@@ -462,7 +492,7 @@ func (h *Host) onEnter(depth int, typ byte, from common.Address, to common.Addre
 	if !parentCallFrame.Prank.Broadcast {
 		return
 	}
-	if to == VMAddr || to == ConsoleAddr { // no broadcasts to the cheatcode or console address
+	if to == addresses.VMAddr || to == addresses.ConsoleAddr { // no broadcasts to the cheatcode or console address
 		return
 	}
 
@@ -561,6 +591,13 @@ func (h *Host) unwindCallstack(depth int) {
 func (h *Host) onOpcode(pc uint64, op byte, gas, cost uint64, scope tracing.OpContext, rData []byte, depth int, err error) {
 	h.unwindCallstack(depth)
 	scopeCtx := scope.(*vm.ScopeContext)
+	if scopeCtx.Contract.IsDeployment {
+		// If we are not yet allowed access to cheatcodes, but if the caller is,
+		// and if this is a contract-creation, then we are automatically granted cheatcode access.
+		if !h.AllowedCheatcodes(scopeCtx.Address()) && h.AllowedCheatcodes(scopeCtx.Caller()) {
+			h.AllowCheatcodes(scopeCtx.Address())
+		}
+	}
 	// Check if we are entering a new depth, add it to the call-stack if so.
 	// We do this here, instead of onEnter, to capture an initialized scope.
 	if len(h.callStack) == 0 || h.callStack[len(h.callStack)-1].Depth < depth {
@@ -609,11 +646,11 @@ func (h *Host) onLog(ev *types.Log) {
 
 // CurrentCall returns the top of the callstack. Or zeroed if there was no call frame yet.
 // If zeroed, the call-frame has a nil scope context.
-func (h *Host) CurrentCall() CallFrame {
+func (h *Host) CurrentCall() *CallFrame {
 	if len(h.callStack) == 0 {
-		return CallFrame{}
+		return &CallFrame{}
 	}
-	return *h.callStack[len(h.callStack)-1]
+	return h.callStack[len(h.callStack)-1]
 }
 
 // MsgSender returns the msg.sender of the current active EVM call-frame,
@@ -652,27 +689,35 @@ func (h *Host) SetEnvVar(key string, value string) {
 // After flushing the EVM state also cannot revert to a previous snapshot state:
 // the state should not be dumped within contract-execution that needs to revert.
 func (h *Host) StateDump() (*foundry.ForgeAllocs, error) {
+	if id, ok := h.state.ActiveFork(); ok {
+		return nil, fmt.Errorf("cannot state-dump while fork %s is active", id)
+	}
+	baseState := h.baseState
 	// We have to commit the existing state to the trie,
 	// for all the state-changes to be captured by the trie iterator.
-	root, err := h.state.Commit(h.env.Context.BlockNumber.Uint64(), true)
+	root, err := baseState.Commit(h.env.Context.BlockNumber.Uint64(), true)
 	if err != nil {
 		return nil, fmt.Errorf("failed to commit state: %w", err)
 	}
 	// We need a state object around the state DB
-	st, err := state.New(root, h.stateDB)
+	st, err := state.New(root, baseState.Database())
 	if err != nil {
 		return nil, fmt.Errorf("failed to create state object for state-dumping: %w", err)
 	}
 	// After Commit we cannot reuse the old State, so we update the host to use the new one
-	h.state = st
-	h.env.StateDB = st
+	h.baseState = st
+	h.state.SubstituteBaseState(st)
 
+	// We use the new state object for state-dumping & future state-access, wrapped around
+	// the just committed trie that has all changes in it.
+	// I.e. the trie is committed and ready to provide all data,
+	// and the state is new and iterable, prepared specifically for FromState(state).
 	var allocs foundry.ForgeAllocs
 	allocs.FromState(st)
 
 	// Sanity check we have no lingering scripts.
-	for i := uint64(0); i <= allocs.Accounts[ScriptDeployer].Nonce; i++ {
-		scriptAddr := crypto.CreateAddress(ScriptDeployer, i)
+	for i := uint64(0); i <= allocs.Accounts[addresses.ScriptDeployer].Nonce; i++ {
+		scriptAddr := crypto.CreateAddress(addresses.ScriptDeployer, i)
 		h.log.Info("removing script from state-dump", "addr", scriptAddr, "label", h.labels[scriptAddr])
 		delete(allocs.Accounts, scriptAddr)
 	}
@@ -694,12 +739,12 @@ func (h *Host) StateDump() (*foundry.ForgeAllocs, error) {
 	}
 
 	// Remove the script deployer from the output
-	delete(allocs.Accounts, ScriptDeployer)
-	delete(allocs.Accounts, ForgeDeployer)
+	delete(allocs.Accounts, addresses.ScriptDeployer)
+	delete(allocs.Accounts, addresses.ForgeDeployer)
 
 	// The cheatcodes VM has a placeholder bytecode,
 	// because solidity checks if the code exists prior to regular EVM-calls to it.
-	delete(allocs.Accounts, VMAddr)
+	delete(allocs.Accounts, addresses.VMAddr)
 
 	// Precompile overrides come with temporary state account placeholders. Ignore those.
 	for addr := range h.precompiles {
@@ -776,7 +821,7 @@ func (h *Host) Label(addr common.Address, label string) {
 
 // NewScriptAddress creates a new address for the ScriptDeployer account, and bumps the nonce.
 func (h *Host) NewScriptAddress() common.Address {
-	deployer := ScriptDeployer
+	deployer := addresses.ScriptDeployer
 	deployNonce := h.state.GetNonce(deployer)
 	// compute address of script contract to be deployed
 	addr := crypto.CreateAddress(deployer, deployNonce)

op-chain-ops/script/script_test.go

@@ -2,12 +2,18 @@ package script
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"math/big"
 	"strings"
 	"testing"
 
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/addresses"
+
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/forking"
+	"github.com/stretchr/testify/mock"
+
 	"github.com/holiman/uint256"
 	"github.com/stretchr/testify/require"
 
@@ -23,16 +29,27 @@ import (
 
 //go:generate ./testdata/generate.sh
 
+// MockRPCClient implements RPCClient interface for testing
+type MockRPCClient struct {
+	mock.Mock
+}
+
+func (m *MockRPCClient) CallContext(ctx context.Context, result any, method string, args ...any) error {
+	return m.Called(ctx, result, method, args).Error(0)
+}
+
 func TestScript(t *testing.T) {
 	logger, captLog := testlog.CaptureLogger(t, log.LevelInfo)
 	af := foundry.OpenArtifactsDir("./testdata/test-artifacts")
 
 	scriptContext := DefaultContext
 	h := NewHost(logger, af, nil, scriptContext)
+	require.NoError(t, h.EnableCheats())
+
 	addr, err := h.LoadContract("ScriptExample.s.sol", "ScriptExample")
 	require.NoError(t, err)
-
-	require.NoError(t, h.EnableCheats())
+	h.AllowCheatcodes(addr)
+	t.Logf("allowing %s to access cheatcodes", addr)
 
 	h.SetEnvVar("EXAMPLE_BOOL", "true")
 	input := bytes4("run()")
@@ -45,18 +62,18 @@ func TestScript(t *testing.T) {
 	require.NoError(t, h.cheatcodes.Precompile.DumpState("noop"))
 }
 
-func TestScriptBroadcast(t *testing.T) {
-	logger := testlog.Logger(t, log.LevelDebug)
-	af := foundry.OpenArtifactsDir("./testdata/test-artifacts")
-
-	mustEncodeCalldata := func(method, input string) []byte {
+func mustEncodeStringCalldata(t *testing.T, method, input string) []byte {
 	packer, err := abi.JSON(strings.NewReader(fmt.Sprintf(`[{"type":"function","name":"%s","inputs":[{"type":"string","name":"input"}]}]`, method)))
 	require.NoError(t, err)
 
 	data, err := packer.Pack(method, input)
 	require.NoError(t, err)
 	return data
-	}
+}
+
+func TestScriptBroadcast(t *testing.T) {
+	logger := testlog.Logger(t, log.LevelDebug)
+	af := foundry.OpenArtifactsDir("./testdata/test-artifacts")
 
 	fooBar, err := af.ReadArtifact("ScriptExample.s.sol", "FooBar")
 	require.NoError(t, err)
@@ -74,7 +91,7 @@ func TestScriptBroadcast(t *testing.T) {
 		{
 			From:    scriptAddr,
 			To:      scriptAddr,
-			Input:   mustEncodeCalldata("call1", "single_call1"),
+			Input:   mustEncodeStringCalldata(t, "call1", "single_call1"),
 			Value:   (*hexutil.U256)(uint256.NewInt(0)),
 			GasUsed: 23421,
 			Type:    BroadcastCall,
@@ -83,7 +100,7 @@ func TestScriptBroadcast(t *testing.T) {
 		{
 			From:    coffeeAddr,
 			To:      scriptAddr,
-			Input:   mustEncodeCalldata("call1", "startstop_call1"),
+			Input:   mustEncodeStringCalldata(t, "call1", "startstop_call1"),
 			Value:   (*hexutil.U256)(uint256.NewInt(0)),
 			GasUsed: 1521,
 			Type:    BroadcastCall,
@@ -92,7 +109,7 @@ func TestScriptBroadcast(t *testing.T) {
 		{
 			From:    coffeeAddr,
 			To:      scriptAddr,
-			Input:   mustEncodeCalldata("call2", "startstop_call2"),
+			Input:   mustEncodeStringCalldata(t, "call2", "startstop_call2"),
 			Value:   (*hexutil.U256)(uint256.NewInt(0)),
 			GasUsed: 1565,
 			Type:    BroadcastCall,
@@ -101,7 +118,7 @@ func TestScriptBroadcast(t *testing.T) {
 		{
 			From:    common.HexToAddress("0x1234"),
 			To:      scriptAddr,
-			Input:   mustEncodeCalldata("nested1", "nested"),
+			Input:   mustEncodeStringCalldata(t, "nested1", "nested"),
 			Value:   (*hexutil.U256)(uint256.NewInt(0)),
 			GasUsed: 2763,
 			Type:    BroadcastCall,
@@ -142,10 +159,11 @@ func TestScriptBroadcast(t *testing.T) {
 		broadcasts = append(broadcasts, broadcast)
 	}
 	h := NewHost(logger, af, nil, DefaultContext, WithBroadcastHook(hook), WithCreate2Deployer())
+	require.NoError(t, h.EnableCheats())
+
 	addr, err := h.LoadContract("ScriptExample.s.sol", "ScriptExample")
 	require.NoError(t, err)
-
-	require.NoError(t, h.EnableCheats())
+	h.AllowCheatcodes(addr)
 
 	input := bytes4("runBroadcast()")
 	returnData, _, err := h.Call(senderAddr, addr, input[:], DefaultFoundryGasLimit, uint256.NewInt(0))
@@ -168,3 +186,163 @@ func TestScriptBroadcast(t *testing.T) {
 	// address that will perform the send to the Create2Deployer.
 	require.EqualValues(t, 1, h.GetNonce(cafeAddr))
 }
+
+func TestScriptStateDump(t *testing.T) {
+	logger := testlog.Logger(t, log.LevelDebug)
+	af := foundry.OpenArtifactsDir("./testdata/test-artifacts")
+
+	h := NewHost(logger, af, nil, DefaultContext)
+	require.NoError(t, h.EnableCheats())
+
+	addr, err := h.LoadContract("ScriptExample.s.sol", "ScriptExample")
+	require.NoError(t, err)
+	h.AllowCheatcodes(addr)
+
+	counterStorageSlot := common.Hash{}
+
+	dump, err := h.StateDump()
+	require.NoError(t, err, "dump 1")
+	require.Contains(t, dump.Accounts, addr, "has contract")
+	require.NotContains(t, dump.Accounts[addr].Storage, counterStorageSlot, "not counted yet")
+
+	dat := mustEncodeStringCalldata(t, "call1", "call A")
+	returnData, _, err := h.Call(addresses.DefaultSenderAddr, addr, dat, DefaultFoundryGasLimit, uint256.NewInt(0))
+	require.NoError(t, err, "call A failed: %x", string(returnData))
+
+	dump, err = h.StateDump()
+	require.NoError(t, err, "dump 2")
+	require.Contains(t, dump.Accounts, addr, "has contract")
+	require.Equal(t, dump.Accounts[addr].Storage[counterStorageSlot], common.Hash{31: 1}, "counted to 1")
+
+	dat = mustEncodeStringCalldata(t, "call1", "call B")
+	returnData, _, err = h.Call(addresses.DefaultSenderAddr, addr, dat, DefaultFoundryGasLimit, uint256.NewInt(0))
+	require.NoError(t, err, "call B failed: %x", string(returnData))
+
+	dump, err = h.StateDump()
+	require.NoError(t, err, "dump 3")
+	require.Contains(t, dump.Accounts, addr, "has contract")
+	require.Equal(t, dump.Accounts[addr].Storage[counterStorageSlot], common.Hash{31: 2}, "counted to 2")
+}
+
+type forkConfig struct {
+	blockNum     uint64
+	stateRoot    common.Hash
+	blockHash    common.Hash
+	nonce        uint64
+	storageValue *big.Int
+	code         []byte
+	balance      uint64
+}
+
+func TestForkingScript(t *testing.T) {
+	logger := testlog.Logger(t, log.LevelInfo)
+	af := foundry.OpenArtifactsDir("./testdata/test-artifacts")
+
+	forkedContract, err := af.ReadArtifact("ScriptExample.s.sol", "ForkedContract")
+	require.NoError(t, err)
+	code := forkedContract.DeployedBytecode.Object
+
+	fork1Config := forkConfig{
+		blockNum:     12345,
+		stateRoot:    common.HexToHash("0x1111"),
+		blockHash:    common.HexToHash("0x2222"),
+		nonce:        12345,
+		storageValue: big.NewInt(1),
+		code:         code,
+		balance:      1,
+	}
+
+	fork2Config := forkConfig{
+		blockNum:     23456,
+		stateRoot:    common.HexToHash("0x3333"),
+		blockHash:    common.HexToHash("0x4444"),
+		nonce:        23456,
+		storageValue: big.NewInt(2),
+		code:         code,
+		balance:      2,
+	}
+
+	// Map of URL/alias to RPC client
+	rpcClients := map[string]*MockRPCClient{
+		"fork1": setupMockRPC(fork1Config),
+		"fork2": setupMockRPC(fork2Config),
+	}
+	forkHook := func(opts *ForkConfig) (forking.ForkSource, error) {
+		client, ok := rpcClients[opts.URLOrAlias]
+		if !ok {
+			return nil, fmt.Errorf("unknown fork URL/alias: %s", opts.URLOrAlias)
+		}
+		return forking.RPCSourceByNumber(opts.URLOrAlias, client, *opts.BlockNumber)
+	}
+
+	scriptContext := DefaultContext
+	h := NewHost(logger, af, nil, scriptContext, WithForkHook(forkHook))
+	require.NoError(t, h.EnableCheats())
+
+	addr, err := h.LoadContract("ScriptExample.s.sol", "ForkTester")
+	require.NoError(t, err)
+	h.AllowCheatcodes(addr)
+	// Make this script persistent so it doesn't call the fork RPC.
+	h.state.MakePersistent(addr)
+	t.Logf("allowing %s to access cheatcodes", addr)
+
+	input := bytes4("run()")
+	returnData, _, err := h.Call(scriptContext.Sender, addr, input[:], DefaultFoundryGasLimit, uint256.NewInt(0))
+	require.NoError(t, err, "call failed: %x", string(returnData))
+
+	for _, client := range rpcClients {
+		client.AssertExpectations(t)
+	}
+}
+
+// setupMockRPC creates a mock RPC client with the specified fork configuration
+func setupMockRPC(config forkConfig) *MockRPCClient {
+	mockRPC := new(MockRPCClient)
+	testAddr := common.HexToAddress("0x1234")
+
+	forkArgs := []any{testAddr, config.blockHash}
+
+	// Mock block header
+	mockRPC.On("CallContext", mock.Anything, mock.AnythingOfType("**forking.Header"),
+		"eth_getBlockByNumber", []any{hexutil.Uint64(config.blockNum), false}).
+		Run(func(args mock.Arguments) {
+			result := args.Get(1).(**forking.Header)
+			*result = &forking.Header{
+				StateRoot: config.stateRoot,
+				BlockHash: config.blockHash,
+			}
+		}).Return(nil).Once()
+
+	mockRPC.On("CallContext", mock.Anything, mock.AnythingOfType("*hexutil.Uint64"),
+		"eth_getTransactionCount", forkArgs).
+		Run(func(args mock.Arguments) {
+			result := args.Get(1).(*hexutil.Uint64)
+			*result = hexutil.Uint64(config.nonce)
+		}).Return(nil)
+
+	// Mock balance
+	mockRPC.On("CallContext", mock.Anything, mock.AnythingOfType("*hexutil.U256"),
+		"eth_getBalance", forkArgs).
+		Run(func(args mock.Arguments) {
+			result := args.Get(1).(*hexutil.U256)
+			*result = hexutil.U256(*uint256.NewInt(config.balance))
+		}).Return(nil)
+
+	// Mock contract code
+	mockRPC.On("CallContext", mock.Anything, mock.AnythingOfType("*hexutil.Bytes"),
+		"eth_getCode", forkArgs).
+		Run(func(args mock.Arguments) {
+			result := args.Get(1).(*hexutil.Bytes)
+			*result = config.code
+		}).Return(nil)
+
+	// Mock storage value
+	mockRPC.On("CallContext", mock.Anything, mock.AnythingOfType("*common.Hash"),
+		"eth_getStorageAt", []any{testAddr, common.Hash{}, config.blockHash}).
+		Run(func(args mock.Arguments) {
+			result := args.Get(1).(*common.Hash)
+			*result = common.BigToHash(config.storageValue)
+		}).Return(nil)
+
+	return mockRPC
+}

op-chain-ops/script/testdata/scripts/ScriptExample.s.sol

@@ -13,6 +13,10 @@ interface Vm {
     function startBroadcast(address msgSender) external;
     function startBroadcast() external;
     function stopBroadcast() external;
+    function getDeployedCode(string calldata artifactPath) external view returns (bytes memory runtimeBytecode);
+    function etch(address target, bytes calldata newRuntimeBytecode) external;
+    function allowCheatcodes(address account) external;
+    function createSelectFork(string calldata forkName, uint256 blockNumber) external returns (uint256);
 }
 
 // console is a minimal version of the console2 lib.
@@ -96,6 +100,13 @@ contract ScriptExample {
         vm.stopPrank();
         this.hello("from original again");
 
+        // vm.etch should not give cheatcode access, unless allowed to afterwards
+        address tmpNonceGetter = address(uint160(uint256(keccak256("temp nonce test getter"))));
+        vm.etch(tmpNonceGetter, vm.getDeployedCode("ScriptExample.s.sol:NonceGetter"));
+        vm.allowCheatcodes(tmpNonceGetter);
+        uint256 v = NonceGetter(tmpNonceGetter).getNonce(address(this));
+        console.log("nonce from nonce getter, no explicit access required with vm.etch:", v);
+
         console.log("done!");
     }
 
@@ -177,3 +188,44 @@ contract FooBar {
         foo = v;
     }
 }
+
+contract NonceGetter {
+    address internal constant VM_ADDRESS = address(uint160(uint256(keccak256("hevm cheat code"))));
+    Vm internal constant vm = Vm(VM_ADDRESS);
+
+    function getNonce(address _addr) public view returns (uint256) {
+        return vm.getNonce(_addr);
+    }
+}
+
+contract ForkedContract {
+    uint256 internal v;
+
+    constructor() {
+        v = 1;
+    }
+
+    function getValue() public view returns (uint256) {
+        return v;
+    }
+}
+
+contract ForkTester {
+    address internal constant VM_ADDRESS = address(uint160(uint256(keccak256("hevm cheat code"))));
+    Vm internal constant vm = Vm(VM_ADDRESS);
+
+    function run() external {
+        address testAddr = address(uint160(0x1234));
+        ForkedContract fc = ForkedContract(testAddr);
+
+        vm.createSelectFork("fork1", 12345);
+        require(vm.getNonce(testAddr) == 12345, "nonce should be 12345");
+        require(fc.getValue() == 1, "value should be 1");
+        require(testAddr.balance == uint256(1), "balance should be 1");
+
+        vm.createSelectFork("fork2", 23456);
+        require(vm.getNonce(testAddr) == 23456, "nonce should be 12345");
+        require(fc.getValue() == 2, "value should be 2");
+        require(testAddr.balance == uint256(2), "balance should be 2");
+    }
+}
\ No newline at end of file

op-chain-ops/script/with.go

@@ -3,6 +3,8 @@ package script
 import (
 	"fmt"
 
+	"github.com/ethereum-optimism/optimism/op-chain-ops/script/addresses"
+
 	"github.com/ethereum/go-ethereum/accounts/abi"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
@@ -26,10 +28,13 @@ func WithScript[B any](h *Host, name string, contract string) (b *B, cleanup fun
 		return nil, nil, fmt.Errorf("could not load script artifact: %w", err)
 	}
 
-	deployer := ScriptDeployer
+	deployer := addresses.ScriptDeployer
 	deployNonce := h.state.GetNonce(deployer)
 	// compute address of script contract to be deployed
 	addr := crypto.CreateAddress(deployer, deployNonce)
+	h.Label(addr, contract)
+	h.AllowCheatcodes(addr)      // before constructor execution, give our script cheatcode access
+	h.state.MakePersistent(addr) // scripts are persistent across forks
 
 	// init bindings (with ABI check)
 	bindings, err := MakeBindings[B](h.ScriptBackendFn(addr), func(abiDef string) bool {
@@ -51,7 +56,6 @@ func WithScript[B any](h *Host, name string, contract string) (b *B, cleanup fun
 		return nil, nil, fmt.Errorf("deployed to unexpected address %s, expected %s", deployedAddr, addr)
 	}
 	h.RememberArtifact(addr, artifact, contract)
-	h.Label(addr, contract)
 	return bindings, func() {
 		h.Wipe(addr)
 	}, nil