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/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/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_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