state upgrade with nonce

Author: havan

params/state_upgrade.go

@@ -27,6 +27,7 @@ type StateUpgradeAccount struct {
 	Code          hexutil.Bytes               `json:"code,omitempty"`
 	Storage       map[common.Hash]common.Hash `json:"storage,omitempty"`
 	BalanceChange *math.HexOrDecimal256       `json:"balanceChange,omitempty"`
+	Nonce         *uint64                     `json:"nonce,omitempty"`
 }
 
 func (s *StateUpgrade) Equal(other *StateUpgrade) bool {

params/state_upgrade_test.go

@@ -160,7 +160,8 @@ func TestUnmarshalStateUpgradeJSON(t *testing.T) {
 					"blockTimestamp": 1677608400,
 					"accounts": {
 						"0x8db97C7cEcE249c2b98bDC0226Cc4C2A57BF52FC": {
-							"balanceChange": "100"
+							"balanceChange": "100",
+							"nonce": 42
 						}
 					}
 				}
@@ -175,6 +176,7 @@ func TestUnmarshalStateUpgradeJSON(t *testing.T) {
 				StateUpgradeAccounts: map[common.Address]StateUpgradeAccount{
 					common.HexToAddress("0x8db97C7cEcE249c2b98bDC0226Cc4C2A57BF52FC"): {
 						BalanceChange: (*math.HexOrDecimal256)(big.NewInt(100)),
+						Nonce:         utils.NewUint64(42),
 					},
 				},
 			},

plugin/evm/vm_upgrade_bytes_test.go

@@ -259,13 +259,23 @@ func TestVMStateUpgrade(t *testing.T) {
 		Code:          upgradedCode,
 	}
 
+	// Create a new account just for testing nonce setting
+	nonceTestAccount := common.Address{99}
+	explicitNonce := uint64(10)
+	nonceTestAccountUpgrade := &params.StateUpgradeAccount{
+		BalanceChange: (*math.HexOrDecimal256)(big.NewInt(50)),
+		Nonce:         &explicitNonce,
+		// No code for this account - should still set the nonce directly
+	}
+
 	upgradeTimestamp := upgrade.InitiallyActiveTime.Add(10 * time.Hour)
 	upgradeBytesJSON := fmt.Sprintf(
 		`{
 			"stateUpgrades": [
 				{
 					"blockTimestamp": %d,
 					"accounts": {
+						"%s": %s,
 						"%s": %s,
 						"%s": %s
 					}
@@ -277,6 +287,8 @@ func TestVMStateUpgrade(t *testing.T) {
 		mustMarshal(t, genesisAccountUpgrade),
 		newAccount.Hex(),
 		mustMarshal(t, newAccountUpgrade),
+		nonceTestAccount.Hex(),
+		mustMarshal(t, nonceTestAccountUpgrade),
 	)
 	require.Contains(t, upgradeBytesJSON, upgradedCodeStr)
 
@@ -329,6 +341,12 @@ func TestVMStateUpgrade(t *testing.T) {
 	require.Equal(t, state.GetCodeHash(newAccount), crypto.Keccak256Hash(upgradedCode))
 	require.Equal(t, state.GetNonce(newAccount), uint64(1)) // Nonce should be set to 1 when code is set if nonce was 0
 	require.Equal(t, state.GetState(newAccount, storageKey), newAccountUpgrade.Storage[storageKey])
+
+	// Test the nonce-specific account
+	expectedNonceTestAccountBalance := uint256.MustFromBig((*big.Int)(nonceTestAccountUpgrade.BalanceChange))
+	require.Equal(t, state.GetNonce(nonceTestAccount), explicitNonce)
+	require.Equal(t, state.GetBalance(nonceTestAccount), expectedNonceTestAccountBalance)
+	require.Empty(t, state.GetCode(nonceTestAccount)) // No code should be set
 }
 
 func TestVMEupgradeActivatesCancun(t *testing.T) {

stateupgrade/state_upgrade.go

@@ -33,12 +33,19 @@ func upgradeAccount(account common.Address, upgrade params.StateUpgradeAccount,
 		balanceChange, _ := uint256.FromBig((*big.Int)(upgrade.BalanceChange))
 		state.AddBalance(account, balanceChange)
 	}
-	if len(upgrade.Code) != 0 {
-		// if the nonce is 0, set the nonce to 1 as we would when deploying a contract at
-		// the address.
+
+	// Set nonce if explicitly provided
+	if upgrade.Nonce != nil {
+		state.SetNonce(account, *upgrade.Nonce)
+	} else if len(upgrade.Code) != 0 {
+		// If no explicit nonce is provided but code is being set, set the nonce to
+		// 1 as we would when deploying a contract at the address.
 		if isEIP158 && state.GetNonce(account) == 0 {
 			state.SetNonce(account, 1)
 		}
+	}
+
+	if len(upgrade.Code) != 0 {
 		state.SetCode(account, upgrade.Code)
 	}
 	for key, value := range upgrade.Storage {