Fix xevm snapshot journal reverts

giga/deps/xevm/state/code.go

@@ -1,8 +1,11 @@
 package state
 
 import (
+	"slices"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/sei-protocol/sei-chain/giga/deps/xevm/types"
 )
 
 func (s *DBImpl) GetCodeHash(addr common.Address) common.Hash {
@@ -15,6 +18,10 @@ func (s *DBImpl) GetCode(addr common.Address) []byte {
 
 func (s *DBImpl) SetCode(addr common.Address, code []byte) []byte {
 	oldCode := s.GetCode(addr)
+	codeStore := s.k.PrefixStore(s.ctx, types.CodeKeyPrefix)
+	prevCode := codeStore.Get(addr[:])
+	prevCodeExists := prevCode != nil
+	prevMapping := captureAddressMapping(s, addr)
 	if s.logger != nil && s.logger.OnCodeChange != nil {
 		// The SetCode method could be modified to return the old code/hash directly.
 		oldHash := s.GetCodeHash(addr)
@@ -23,7 +30,12 @@ func (s *DBImpl) SetCode(addr common.Address, code []byte) []byte {
 	}
 
 	s.k.SetCode(s.ctx, addr, code)
-	s.journal = append(s.journal, &codeChange{addr: addr, prevCode: oldCode})
+	s.journal = append(s.journal, &codeChange{
+		addr:           addr,
+		prevCode:       slices.Clone(prevCode),
+		prevCodeExists: prevCodeExists,
+		prevMapping:    prevMapping,
+	})
 	return oldCode
 }
 

giga/deps/xevm/state/journal.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/sei-protocol/sei-chain/giga/deps/xevm/types"
 	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
 )
@@ -60,14 +61,17 @@ type (
 
 	// codeChange records a code mutation so it can be reverted.
 	codeChange struct {
-		addr     common.Address
-		prevCode []byte
+		addr           common.Address
+		prevCode       []byte
+		prevCodeExists bool
+		prevMapping    addressMappingState
 	}
 
 	// nonceChange records a nonce mutation so it can be reverted.
 	nonceChange struct {
-		addr common.Address
-		prev uint64
+		addr       common.Address
+		prev       uint64
+		prevExists bool
 	}
 
 	// balanceChange records an Add or Sub balance so it can be reverted.
@@ -81,17 +85,26 @@ type (
 
 	// createAccountChange records the previous state cleared by clearAccountStateJournaled.
 	createAccountChange struct {
-		addr      common.Address
-		prevCode  []byte
-		prevNonce uint64
-		prevSlots map[common.Hash]common.Hash
+		addr            common.Address
+		prevCode        []byte
+		prevCodeExists  bool
+		prevNonce       uint64
+		prevNonceExists bool
+		prevSlots       map[common.Hash]common.Hash
 	}
 
 	// deleteMappingChange records a DeleteAddressMapping so it can be reverted.
 	deleteMappingChange struct {
 		evmAddr common.Address
 		seiAddr sdk.AccAddress
 	}
+
+	addressMappingState struct {
+		exists              bool
+		seiAddr             sdk.AccAddress
+		accountCreated      bool
+		globalAccountNumber []byte
+	}
 )
 
 func (e *accessListAddAccountChange) revert(s *DBImpl) {
@@ -172,18 +185,12 @@ func (e *storageChange) revert(s *DBImpl) {
 }
 
 func (e *codeChange) revert(s *DBImpl) {
-	if len(e.prevCode) == 0 {
-		// Directly delete code entries since SetCode(nil) sets to empty but doesn't delete
-		deleteIfExists(s.k.PrefixStore(s.ctx, types.CodeKeyPrefix), e.addr[:])
-		deleteIfExists(s.k.PrefixStore(s.ctx, types.CodeHashKeyPrefix), e.addr[:])
-		deleteIfExists(s.k.PrefixStore(s.ctx, types.CodeSizeKeyPrefix), e.addr[:])
-	} else {
-		s.k.SetCode(s.ctx, e.addr, e.prevCode)
-	}
+	restoreCode(s, e.addr, e.prevCode, e.prevCodeExists)
+	e.prevMapping.restore(s, e.addr)
 }
 
 func (e *nonceChange) revert(s *DBImpl) {
-	s.k.SetNonce(s.ctx, e.addr, e.prev)
+	restoreNonce(s, e.addr, e.prev, e.prevExists)
 }
 
 func (e *balanceChange) revert(s *DBImpl) {
@@ -210,8 +217,8 @@ func (e *balanceChange) revert(s *DBImpl) {
 }
 
 func (e *createAccountChange) revert(s *DBImpl) {
-	s.k.SetCode(s.ctx, e.addr, e.prevCode)
-	s.k.SetNonce(s.ctx, e.addr, e.prevNonce)
+	restoreCode(s, e.addr, e.prevCode, e.prevCodeExists)
+	restoreNonce(s, e.addr, e.prevNonce, e.prevNonceExists)
 	for k, v := range e.prevSlots {
 		s.k.SetState(s.ctx, e.addr, k, v)
 	}
@@ -221,3 +228,63 @@ func (e *deleteMappingChange) revert(s *DBImpl) {
 	ctx := s.ctx.WithEventManager(sdk.NewEventManager())
 	s.k.SetAddressMapping(ctx, e.seiAddr, e.evmAddr)
 }
+
+func captureAddressMapping(s *DBImpl, addr common.Address) addressMappingState {
+	seiAddr, ok := s.k.GetSeiAddress(s.ctx, addr)
+	if !ok {
+		seiAddr = s.k.GetSeiAddressOrDefault(s.ctx, addr)
+	}
+	accountExists := s.k.AccountKeeper().HasAccount(s.ctx, seiAddr)
+	return addressMappingState{
+		exists:              ok,
+		seiAddr:             append(sdk.AccAddress(nil), seiAddr...),
+		accountCreated:      !ok && !accountExists,
+		globalAccountNumber: s.k.AccountKeeper().GetGlobalAccountNumberBytes(s.ctx),
+	}
+}
+
+func (m addressMappingState) restore(s *DBImpl, addr common.Address) {
+	currentSeiAddr, ok := s.k.GetSeiAddress(s.ctx, addr)
+	if ok && (!m.exists || !currentSeiAddr.Equals(m.seiAddr)) {
+		s.k.DeleteAddressMapping(s.ctx, currentSeiAddr, addr)
+	}
+	if m.exists && (!ok || !currentSeiAddr.Equals(m.seiAddr)) {
+		ctx := s.ctx.WithEventManager(sdk.NewEventManager())
+		s.k.SetAddressMapping(ctx, m.seiAddr, addr)
+	}
+	if m.accountCreated {
+		if acc := s.k.AccountKeeper().GetAccount(s.ctx, m.seiAddr); acc != nil {
+			s.k.AccountKeeper().RemoveAccount(s.ctx, acc)
+		}
+		s.k.AccountKeeper().SetGlobalAccountNumberBytes(s.ctx, m.globalAccountNumber)
+	}
+}
+
+func restoreCode(s *DBImpl, addr common.Address, code []byte, exists bool) {
+	if !exists {
+		deleteIfExists(s.k.PrefixStore(s.ctx, types.CodeKeyPrefix), addr[:])
+		deleteIfExists(s.k.PrefixStore(s.ctx, types.CodeHashKeyPrefix), addr[:])
+		deleteIfExists(s.k.PrefixStore(s.ctx, types.CodeSizeKeyPrefix), addr[:])
+		return
+	}
+
+	if code == nil {
+		code = []byte{}
+	}
+	s.k.PrefixStore(s.ctx, types.CodeKeyPrefix).Set(addr[:], code)
+
+	length := make([]byte, 8)
+	binary.BigEndian.PutUint64(length, uint64(len(code)))
+	s.k.PrefixStore(s.ctx, types.CodeSizeKeyPrefix).Set(addr[:], length)
+
+	hash := crypto.Keccak256Hash(code)
+	s.k.PrefixStore(s.ctx, types.CodeHashKeyPrefix).Set(addr[:], hash[:])
+}
+
+func restoreNonce(s *DBImpl, addr common.Address, nonce uint64, exists bool) {
+	if !exists {
+		deleteIfExists(s.k.PrefixStore(s.ctx, types.NonceKeyPrefix), addr[:])
+		return
+	}
+	s.k.SetNonce(s.ctx, addr, nonce)
+}

giga/deps/xevm/state/nonce.go

@@ -3,6 +3,7 @@ package state
 import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/tracing"
+	"github.com/sei-protocol/sei-chain/giga/deps/xevm/types"
 )
 
 func (s *DBImpl) GetNonce(addr common.Address) uint64 {
@@ -11,11 +12,12 @@ func (s *DBImpl) GetNonce(addr common.Address) uint64 {
 
 func (s *DBImpl) SetNonce(addr common.Address, nonce uint64, reason tracing.NonceChangeReason) {
 	prevNonce := s.GetNonce(addr)
+	prevExists := s.k.PrefixStore(s.ctx, types.NonceKeyPrefix).Has(addr[:])
 	if s.logger != nil && s.logger.OnNonceChangeV2 != nil {
 		// The SetCode method could be modified to return the old code/hash directly.
 		s.logger.OnNonceChangeV2(addr, prevNonce, nonce, reason)
 	}
 
 	s.k.SetNonce(s.ctx, addr, nonce)
-	s.journal = append(s.journal, &nonceChange{addr: addr, prev: prevNonce})
+	s.journal = append(s.journal, &nonceChange{addr: addr, prev: prevNonce, prevExists: prevExists})
 }

giga/deps/xevm/state/snapshot_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/tracing"
 	ethtypes "github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/holiman/uint256"
 	testkeeper "github.com/sei-protocol/sei-chain/giga/deps/testutil/keeper"
 	"github.com/sei-protocol/sei-chain/giga/deps/xevm/state"
@@ -124,6 +125,72 @@ func TestSnapshotReverts_CodeToEmpty(t *testing.T) {
 	require.Zero(t, sdb.GetCodeSize(evmAddr))
 }
 
+func TestSnapshotReverts_CodeRemovesCreatedMapping(t *testing.T) {
+	k, ctx := testkeeper.MockEVMKeeper(t)
+	ctx = ctx.WithBlockTime(time.Now())
+	sdb := state.NewDBImpl(ctx, k, false)
+
+	_, evmAddr := testkeeper.MockAddressPair()
+	seiAddr := sdk.AccAddress(evmAddr.Bytes())
+	require.False(t, k.AccountKeeper().HasAccount(sdb.Ctx(), seiAddr))
+	prevGlobalAccountNumber := k.AccountKeeper().GetGlobalAccountNumberBytes(sdb.Ctx())
+
+	rev := sdb.Snapshot()
+	sdb.SetCode(evmAddr, []byte("deployed"))
+
+	gotSeiAddr, ok := k.GetSeiAddress(sdb.Ctx(), evmAddr)
+	require.True(t, ok)
+	require.Equal(t, seiAddr, gotSeiAddr)
+	gotEVMAddr, ok := k.GetEVMAddress(sdb.Ctx(), seiAddr)
+	require.True(t, ok)
+	require.Equal(t, evmAddr, gotEVMAddr)
+	require.True(t, k.AccountKeeper().HasAccount(sdb.Ctx(), seiAddr))
+	require.NotEqual(t, prevGlobalAccountNumber, k.AccountKeeper().GetGlobalAccountNumberBytes(sdb.Ctx()))
+
+	sdb.RevertToSnapshot(rev)
+
+	_, ok = k.GetSeiAddress(sdb.Ctx(), evmAddr)
+	require.False(t, ok)
+	_, ok = k.GetEVMAddress(sdb.Ctx(), seiAddr)
+	require.False(t, ok)
+	require.False(t, k.AccountKeeper().HasAccount(sdb.Ctx(), seiAddr))
+	require.Equal(t, prevGlobalAccountNumber, k.AccountKeeper().GetGlobalAccountNumberBytes(sdb.Ctx()))
+}
+
+func TestSnapshotReverts_ExplicitEmptyCodePreserved(t *testing.T) {
+	k, ctx := testkeeper.MockEVMKeeper(t)
+	ctx = ctx.WithBlockTime(time.Now())
+	sdb := state.NewDBImpl(ctx, k, false)
+
+	_, evmAddr := testkeeper.MockAddressPair()
+	sdb.SetCode(evmAddr, []byte{})
+
+	codeStore := k.PrefixStore(sdb.Ctx(), types.CodeKeyPrefix)
+	codeHashStore := k.PrefixStore(sdb.Ctx(), types.CodeHashKeyPrefix)
+	codeSizeStore := k.PrefixStore(sdb.Ctx(), types.CodeSizeKeyPrefix)
+	require.True(t, codeStore.Has(evmAddr[:]))
+	require.True(t, codeHashStore.Has(evmAddr[:]))
+	require.True(t, codeSizeStore.Has(evmAddr[:]))
+	require.Nil(t, sdb.GetCode(evmAddr))
+	require.Equal(t, crypto.Keccak256Hash(nil), sdb.GetCodeHash(evmAddr))
+
+	rev := sdb.Snapshot()
+	sdb.SetCode(evmAddr, []byte("deployed"))
+
+	sdb.RevertToSnapshot(rev)
+
+	codeStore = k.PrefixStore(sdb.Ctx(), types.CodeKeyPrefix)
+	codeHashStore = k.PrefixStore(sdb.Ctx(), types.CodeHashKeyPrefix)
+	codeSizeStore = k.PrefixStore(sdb.Ctx(), types.CodeSizeKeyPrefix)
+	require.True(t, codeStore.Has(evmAddr[:]))
+	require.True(t, codeHashStore.Has(evmAddr[:]))
+	require.True(t, codeSizeStore.Has(evmAddr[:]))
+	require.Equal(t, []byte{}, codeStore.Get(evmAddr[:]))
+	require.Nil(t, sdb.GetCode(evmAddr))
+	require.Zero(t, sdb.GetCodeSize(evmAddr))
+	require.Equal(t, crypto.Keccak256Hash(nil), sdb.GetCodeHash(evmAddr))
+}
+
 // ----------------------------------------------------------------------------
 // KV-state revert: SetNonce / nonceChange
 // ----------------------------------------------------------------------------
@@ -145,6 +212,23 @@ func TestSnapshotReverts_Nonce(t *testing.T) {
 	require.EqualValues(t, 5, sdb.GetNonce(evmAddr))
 }
 
+func TestSnapshotReverts_NonceToAbsent(t *testing.T) {
+	k, ctx := testkeeper.MockEVMKeeper(t)
+	ctx = ctx.WithBlockTime(time.Now())
+	sdb := state.NewDBImpl(ctx, k, false)
+
+	_, evmAddr := testkeeper.MockAddressPair()
+	require.False(t, k.PrefixStore(sdb.Ctx(), types.NonceKeyPrefix).Has(evmAddr[:]))
+
+	rev := sdb.Snapshot()
+	sdb.SetNonce(evmAddr, 7, tracing.NonceChangeUnspecified)
+	require.True(t, k.PrefixStore(sdb.Ctx(), types.NonceKeyPrefix).Has(evmAddr[:]))
+
+	sdb.RevertToSnapshot(rev)
+	require.EqualValues(t, 0, sdb.GetNonce(evmAddr))
+	require.False(t, k.PrefixStore(sdb.Ctx(), types.NonceKeyPrefix).Has(evmAddr[:]))
+}
+
 // ----------------------------------------------------------------------------
 // Balance revert: AddBalance / SubBalance / balanceChange
 // ----------------------------------------------------------------------------
@@ -248,6 +332,24 @@ func TestSnapshotReverts_CreateAccountOnFreshAddress_NoOp(t *testing.T) {
 	require.Nil(t, sdb.GetCode(evmAddr))
 }
 
+func TestSnapshotReverts_CreateAccountPreservesAbsentNonce(t *testing.T) {
+	k, ctx := testkeeper.MockEVMKeeper(t)
+	ctx = ctx.WithBlockTime(time.Now())
+	sdb := state.NewDBImpl(ctx, k, false)
+
+	_, evmAddr := testkeeper.MockAddressPair()
+	sdb.SetCode(evmAddr, []byte("existing_code"))
+	require.False(t, k.PrefixStore(sdb.Ctx(), types.NonceKeyPrefix).Has(evmAddr[:]))
+
+	rev := sdb.Snapshot()
+	sdb.CreateAccount(evmAddr)
+
+	sdb.RevertToSnapshot(rev)
+	require.Equal(t, []byte("existing_code"), sdb.GetCode(evmAddr))
+	require.EqualValues(t, 0, sdb.GetNonce(evmAddr))
+	require.False(t, k.PrefixStore(sdb.Ctx(), types.NonceKeyPrefix).Has(evmAddr[:]))
+}
+
 // ----------------------------------------------------------------------------
 // SelfDestruct: deleteMappingChange
 // ----------------------------------------------------------------------------

giga/deps/xevm/state/state.go

@@ -2,6 +2,7 @@ package state
 
 import (
 	"bytes"
+	"slices"
 	"sort"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -34,6 +35,7 @@ func (s *DBImpl) GetState(addr common.Address, hash common.Hash) common.Hash {
 func (s *DBImpl) SetState(addr common.Address, key common.Hash, val common.Hash) common.Hash {
 	old := s.GetState(addr, key)
 	if old == val {
+		s.k.SetState(s.ctx, addr, key, val)
 		return old
 	}
 	if s.logger != nil && s.logger.OnStorageChange != nil {
@@ -193,8 +195,11 @@ func (s *DBImpl) clearAccountStateJournaled(acc common.Address) {
 	}
 
 	// Save previous state for potential revert.
-	prevCode := s.k.GetCode(s.ctx, acc)
+	codeStore := s.k.PrefixStore(s.ctx, types.CodeKeyPrefix)
+	prevCode := codeStore.Get(acc[:])
+	prevCodeExists := prevCode != nil
 	prevNonce := s.k.GetNonce(s.ctx, acc)
+	prevNonceExists := s.k.PrefixStore(s.ctx, types.NonceKeyPrefix).Has(acc[:])
 
 	// Collect all storage slots for this account using GetAllKeyStrsInRange.
 	// The prefix store's GetAllKeyStrsInRange returns raw parent-store keys,
@@ -217,10 +222,12 @@ func (s *DBImpl) clearAccountStateJournaled(acc common.Address) {
 
 	// Append journal entry before making changes.
 	s.journal = append(s.journal, &createAccountChange{
-		addr:      acc,
-		prevCode:  prevCode,
-		prevNonce: prevNonce,
-		prevSlots: prevSlots,
+		addr:            acc,
+		prevCode:        slices.Clone(prevCode),
+		prevCodeExists:  prevCodeExists,
+		prevNonce:       prevNonce,
+		prevNonceExists: prevNonceExists,
+		prevSlots:       prevSlots,
 	})
 
 	// Clear the account state.