Improvements/bugfixes to Go forge scripts (#11838)

* Improvements/bugfixes to Go forge scripts

Adds some improvements to the Go forge scripts:

- Adds a `GasUsed` field to the `Broadcast` struct so that transaction broadcast utilities can use it for gas estimation. Gas estimation using the RPC will fail when sending transactions in parallel since the state can change significantly between calls.
- Fixes a bug in the `vm.broadcast` cheatcode where sender nonce were not increased for `vm.CALL`s. This led to a mismatch between the contract addresses generated by the Forge tooling, and what was actually being generated onchain.

* op-chain-ops: isolate broadcast functionality

* review updates

* wrap in broadcast check

* Add nonce tests

* Update op-chain-ops/script/script.go
Co-authored-by: protolambda <proto@protolambda.com>

* Fix test

* op-chain-ops: track broadcast nonce, add sanity checks

---------
Co-authored-by: protolambda <proto@protolambda.com>

op-chain-ops/script/prank.go

@@ -195,12 +195,14 @@ func (bt *BroadcastType) UnmarshalText(data []byte) error {
 // via vm.broadcast(). Actually submitting the transaction is left up
 // to other tools.
 type Broadcast struct {
-	From  common.Address `json:"from"`
-	To    common.Address `json:"to"`    // set to expected contract address, if this is a deployment
-	Input hexutil.Bytes  `json:"input"` // set to contract-creation code, if this is a deployment
-	Value *hexutil.U256  `json:"value"`
-	Salt  common.Hash    `json:"salt"` // set if this is a Create2 broadcast
-	Type  BroadcastType  `json:"type"`
+	From    common.Address `json:"from"`
+	To      common.Address `json:"to"`    // set to expected contract address, if this is a deployment
+	Input   hexutil.Bytes  `json:"input"` // set to contract-creation code, if this is a deployment
+	Value   *hexutil.U256  `json:"value"`
+	Salt    common.Hash    `json:"salt"` // set if this is a Create2 broadcast
+	GasUsed uint64         `json:"gasUsed"`
+	Type    BroadcastType  `json:"type"`
+	Nonce   uint64         `json:"nonce"` // pre-state nonce of From, before any increment (always 0 if create2)
 }
 
 // NewBroadcast creates a Broadcast from a parent callframe, and the completed child callframe.
@@ -225,13 +227,22 @@ func NewBroadcast(parent, current *CallFrame) Broadcast {
 		From: ctx.Caller(),
 		To:   ctx.Address(),
 		// Need to clone the input below since memory is reused in the VM
-		Input: bytes.Clone(input),
-		Value: (*hexutil.U256)(value.Clone()),
+		Input:   bytes.Clone(input),
+		Value:   (*hexutil.U256)(value.Clone()),
+		GasUsed: current.GasUsed,
 	}
 
 	switch parent.LastOp {
 	case vm.CREATE:
 		bcast.Type = BroadcastCreate
+		// Nonce bump was already applied, but we need the pre-state
+		bcast.Nonce = current.CallerNonce - 1
+		expectedAddr := crypto.CreateAddress(bcast.From, bcast.Nonce)
+		if expectedAddr != bcast.To {
+			panic(fmt.Errorf("script bug: create broadcast has "+
+				"unexpected address: %s, expected %s. Sender: %s, Nonce: %d",
+				bcast.To, expectedAddr, bcast.From, bcast.Nonce))
+		}
 	case vm.CREATE2:
 		bcast.Salt = parent.LastCreate2Salt
 		initHash := crypto.Keccak256Hash(bcast.Input)
@@ -243,8 +254,11 @@ func NewBroadcast(parent, current *CallFrame) Broadcast {
 				bcast.To, expectedAddr, bcast.From, bcast.Salt, initHash))
 		}
 		bcast.Type = BroadcastCreate2
+		bcast.Nonce = 0 // always 0. The nonce should not matter for create2.
 	case vm.CALL:
 		bcast.Type = BroadcastCall
+		// Nonce bump was already applied, but we need the pre-state
+		bcast.Nonce = current.CallerNonce - 1
 	default:
 		panic(fmt.Errorf("unexpected broadcast operation %s", parent.LastOp))
 	}

op-chain-ops/script/script.go

@@ -55,6 +55,15 @@ type CallFrame struct {
 	// Forge script does not support nested pranks on the same call-depth.
 	// Pranks can also be broadcasting.
 	Prank *Prank
+
+	// GasUsed keeps track of the amount of gas used by this call frame.
+	// This is useful for broadcasts, which sometimes cannot correctly
+	// estimate gas when sending transactions in parallel.
+	GasUsed uint64
+
+	// CallerNonce keeps track of the nonce of the caller who entered the callframe
+	// (nonce of pranked caller, if pranked).
+	CallerNonce uint64
 }
 
 // Host is an EVM executor that runs Forge scripts.
@@ -91,6 +100,11 @@ type Host struct {
 	onLabel []func(name string, addr common.Address)
 
 	hooks *Hooks
+
+	// isolateBroadcasts will flush the journal changes,
+	// and prepare the ephemeral tx context again,
+	// to make gas accounting of a broadcast sub-call more accurate.
+	isolateBroadcasts bool
 }
 
 type HostOption func(h *Host)
@@ -107,6 +121,17 @@ func WithBroadcastHook(hook BroadcastHook) HostOption {
 	}
 }
 
+// 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.
+// This is not compatible with state-snapshots: upon cleaning,
+// it is assumed that the state has to never revert back, similar to the state-dump guarantees.
+func WithIsolatedBroadcasts() HostOption {
+	return func(h *Host) {
+		h.isolateBroadcasts = true
+	}
+}
+
 // NewHost creates a Host that can load contracts from the given Artifacts FS,
 // and with an EVM initialized to the given executionContext.
 // Optionally src-map loading may be enabled, by providing a non-nil srcFS to read sources from.
@@ -217,6 +242,7 @@ func NewHost(
 
 	// Hook up the Host to capture the EVM environment changes
 	trHooks := &tracing.Hooks{
+		OnEnter:         h.onEnter,
 		OnExit:          h.onExit,
 		OnOpcode:        h.onOpcode,
 		OnFault:         h.onFault,
@@ -336,6 +362,16 @@ func (h *Host) Wipe(addr common.Address) {
 	h.state.SetBalance(addr, uint256.NewInt(0), tracing.BalanceChangeUnspecified)
 }
 
+// SetNonce sets an account's nonce in state.
+func (h *Host) SetNonce(addr common.Address, nonce uint64) {
+	h.state.SetNonce(addr, nonce)
+}
+
+// GetNonce returs an account's nonce from state.
+func (h *Host) GetNonce(addr common.Address) uint64 {
+	return h.state.GetNonce(addr)
+}
+
 // 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 {
@@ -365,6 +401,52 @@ func (h *Host) HasPrecompileOverride(addr common.Address) bool {
 	return ok
 }
 
+// onEnter is a trace-hook, which we use to apply changes to the state-DB, to simulate isolated broadcast calls,
+// for better gas estimation of the exact broadcast call execution.
+func (h *Host) onEnter(depth int, typ byte, from common.Address, to common.Address, input []byte, gas uint64, value *big.Int) {
+	if len(h.callStack) == 0 {
+		return
+	}
+	parentCallFrame := h.callStack[len(h.callStack)-1]
+	if parentCallFrame.Prank == nil {
+		return
+	}
+	// sanity check our callframe is set up correctly
+	if parentCallFrame.LastOp != vm.OpCode(typ) {
+		panic(fmt.Errorf("parent call-frame has invalid last Op: %d", typ))
+	}
+	if !parentCallFrame.Prank.Broadcast {
+		return
+	}
+	if to == VMAddr || to == ConsoleAddr { // no broadcasts to the cheatcode or console address
+		return
+	}
+
+	// Bump nonce value, such that a broadcast Call appears like a tx
+	if parentCallFrame.LastOp == vm.CALL {
+		sender := parentCallFrame.Ctx.Address()
+		if parentCallFrame.Prank.Sender != nil {
+			sender = *parentCallFrame.Prank.Sender
+		}
+		h.state.SetNonce(sender, h.state.GetNonce(sender)+1)
+	}
+
+	if h.isolateBroadcasts {
+		var dest *common.Address
+		switch parentCallFrame.LastOp {
+		case vm.CREATE, vm.CREATE2:
+			dest = nil // no destination address to warm up
+		case vm.CALL:
+			dest = &to
+		default:
+			return
+		}
+		h.state.Finalise(true)
+		// the prank msg.sender, if any, has already been applied to 'from' before onEnter
+		h.prelude(from, dest)
+	}
+}
+
 // onExit is a trace-hook, which we use to maintain an accurate view of functions, and log any revert warnings.
 func (h *Host) onExit(depth int, output []byte, gasUsed uint64, err error, reverted bool) {
 	// Note: onExit runs also when going deeper, exiting the context into a nested context.
@@ -377,6 +459,8 @@ func (h *Host) onExit(depth int, output []byte, gasUsed uint64, err error, rever
 			h.log.Warn("Revert", "addr", addr, "err", err, "revertData", hexutil.Bytes(output), "depth", depth)
 		}
 	}
+
+	h.callStack[len(h.callStack)-1].GasUsed += gasUsed
 	h.unwindCallstack(depth)
 }
 
@@ -437,10 +521,11 @@ func (h *Host) onOpcode(pc uint64, op byte, gas, cost uint64, scope tracing.OpCo
 	// 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 {
 		h.callStack = append(h.callStack, &CallFrame{
-			Depth:  depth,
-			LastOp: vm.OpCode(op),
-			LastPC: pc,
-			Ctx:    scopeCtx,
+			Depth:       depth,
+			LastOp:      vm.OpCode(op),
+			LastPC:      pc,
+			Ctx:         scopeCtx,
+			CallerNonce: h.GetNonce(scopeCtx.Caller()),
 		})
 	}
 	// Sanity check that top of the call-stack matches the scope context now

op-chain-ops/script/script_test.go

@@ -66,66 +66,89 @@ func TestScriptBroadcast(t *testing.T) {
 	expectedInitCode = append(expectedInitCode, leftPad32(big.NewInt(1234).Bytes())...)
 	salt := uint256.NewInt(42).Bytes32()
 
-	senderAddr := common.HexToAddress("0x5b73C5498c1E3b4dbA84de0F1833c4a029d90519")
+	senderAddr := common.HexToAddress("0x0000000000000000000000000000000000Badc0d")
+	scriptAddr := common.HexToAddress("0x5b73c5498c1e3b4dba84de0f1833c4a029d90519")
+	coffeeAddr := common.HexToAddress("0x0000000000000000000000000000000000C0FFEE")
+	cafeAddr := common.HexToAddress("0xcafe")
 	expBroadcasts := []Broadcast{
 		{
-			From:  senderAddr,
-			To:    senderAddr,
-			Input: mustEncodeCalldata("call1", "single_call1"),
-			Value: (*hexutil.U256)(uint256.NewInt(0)),
-			Type:  BroadcastCall,
+			From:    scriptAddr,
+			To:      scriptAddr,
+			Input:   mustEncodeCalldata("call1", "single_call1"),
+			Value:   (*hexutil.U256)(uint256.NewInt(0)),
+			GasUsed: 23421,
+			Type:    BroadcastCall,
+			Nonce:   1, // first action by script (script already has a nonce of 1)
 		},
 		{
-			From:  common.HexToAddress("0x0000000000000000000000000000000000C0FFEE"),
-			To:    senderAddr,
-			Input: mustEncodeCalldata("call1", "startstop_call1"),
-			Value: (*hexutil.U256)(uint256.NewInt(0)),
-			Type:  BroadcastCall,
+			From:    coffeeAddr,
+			To:      scriptAddr,
+			Input:   mustEncodeCalldata("call1", "startstop_call1"),
+			Value:   (*hexutil.U256)(uint256.NewInt(0)),
+			GasUsed: 1521,
+			Type:    BroadcastCall,
+			Nonce:   0, // first action by 0xc0ffee
 		},
 		{
-			From:  common.HexToAddress("0x0000000000000000000000000000000000C0FFEE"),
-			To:    senderAddr,
-			Input: mustEncodeCalldata("call2", "startstop_call2"),
-			Value: (*hexutil.U256)(uint256.NewInt(0)),
-			Type:  BroadcastCall,
+			From:    coffeeAddr,
+			To:      scriptAddr,
+			Input:   mustEncodeCalldata("call2", "startstop_call2"),
+			Value:   (*hexutil.U256)(uint256.NewInt(0)),
+			GasUsed: 1565,
+			Type:    BroadcastCall,
+			Nonce:   1, // second action of 0xc0ffee
 		},
 		{
-			From:  common.HexToAddress("0x1234"),
-			To:    senderAddr,
-			Input: mustEncodeCalldata("nested1", "nested"),
-			Value: (*hexutil.U256)(uint256.NewInt(0)),
-			Type:  BroadcastCall,
+			From:    common.HexToAddress("0x1234"),
+			To:      scriptAddr,
+			Input:   mustEncodeCalldata("nested1", "nested"),
+			Value:   (*hexutil.U256)(uint256.NewInt(0)),
+			GasUsed: 2763,
+			Type:    BroadcastCall,
+			Nonce:   0, // first action of 0x1234
 		},
 		{
-			From:  common.HexToAddress("0x123456"),
-			To:    crypto.CreateAddress(common.HexToAddress("0x123456"), 0),
-			Input: expectedInitCode,
-			Value: (*hexutil.U256)(uint256.NewInt(0)),
-			Type:  BroadcastCreate,
+			From:    common.HexToAddress("0x123456"),
+			To:      crypto.CreateAddress(common.HexToAddress("0x123456"), 0),
+			Input:   expectedInitCode,
+			Value:   (*hexutil.U256)(uint256.NewInt(0)),
+			GasUsed: 39112,
+			Type:    BroadcastCreate,
+			Nonce:   0, // first action of 0x123456
 		},
 		{
-			From:  common.HexToAddress("0xcafe"),
-			To:    crypto.CreateAddress2(common.HexToAddress("0xcafe"), salt, crypto.Keccak256(expectedInitCode)),
-			Input: expectedInitCode,
-			Value: (*hexutil.U256)(uint256.NewInt(0)),
-			Type:  BroadcastCreate2,
-			Salt:  salt,
+			From:    cafeAddr,
+			To:      crypto.CreateAddress2(cafeAddr, salt, crypto.Keccak256(expectedInitCode)),
+			Input:   expectedInitCode,
+			Value:   (*hexutil.U256)(uint256.NewInt(0)),
+			Type:    BroadcastCreate2,
+			GasUsed: 39112,
+			Salt:    salt,
+			Nonce:   0, // first action of 0xcafe
+		},
+		{
+			From:    scriptAddr,
+			To:      crypto.CreateAddress(scriptAddr, 2),
+			Input:   expectedInitCode,
+			Value:   (*hexutil.U256)(uint256.NewInt(0)),
+			GasUsed: 39112,
+			Type:    BroadcastCreate,
+			Nonce:   2, // second action, on top of starting at 1.
 		},
 	}
 
-	scriptContext := DefaultContext
 	var broadcasts []Broadcast
 	hook := func(broadcast Broadcast) {
 		broadcasts = append(broadcasts, broadcast)
 	}
-	h := NewHost(logger, af, nil, scriptContext, WithBroadcastHook(hook))
+	h := NewHost(logger, af, nil, DefaultContext, WithBroadcastHook(hook))
 	addr, err := h.LoadContract("ScriptExample.s.sol", "ScriptExample")
 	require.NoError(t, err)
 
 	require.NoError(t, h.EnableCheats())
 
 	input := bytes4("runBroadcast()")
-	returnData, _, err := h.Call(scriptContext.Sender, addr, input[:], DefaultFoundryGasLimit, uint256.NewInt(0))
+	returnData, _, err := h.Call(senderAddr, addr, input[:], DefaultFoundryGasLimit, uint256.NewInt(0))
 	require.NoError(t, err, "call failed: %x", string(returnData))
 
 	expected, err := json.MarshalIndent(expBroadcasts, "  ", "  ")
@@ -133,4 +156,13 @@ func TestScriptBroadcast(t *testing.T) {
 	got, err := json.MarshalIndent(broadcasts, "  ", "  ")
 	require.NoError(t, err)
 	require.Equal(t, string(expected), string(got))
+
+	// Assert that the nonces for accounts participating in the
+	// broadcast increase. The scriptAddr check is set to 3 to
+	// account for the initial deployment of the contract and
+	// two additional calls.
+	require.EqualValues(t, 0, h.GetNonce(senderAddr))
+	require.EqualValues(t, 3, h.GetNonce(scriptAddr))
+	require.EqualValues(t, 2, h.GetNonce(coffeeAddr))
+	require.EqualValues(t, 1, h.GetNonce(cafeAddr))
 }

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

@@ -100,6 +100,8 @@ contract ScriptExample {
 
     /// @notice example function, to test vm.broadcast with.
     function runBroadcast() public {
+        console.log("nonce start", uint256(vm.getNonce(address(this))));
+
         console.log("testing single");
         vm.broadcast();
         this.call1("single_call1");
@@ -128,6 +130,12 @@ contract ScriptExample {
         FooBar y = new FooBar{salt: bytes32(uint256(42))}(1234);
         require(y.foo() == 1234);
         console.log("done!");
+
+        // Deploy a script without a pranked sender and check the nonce.
+        vm.broadcast();
+        new FooBar(1234);
+
+        console.log("nonce end", uint256(vm.getNonce(address(this))));
     }
 
     /// @notice example external function, to force a CALL, and test vm.startPrank with.