[codex] fix giga snapshot storage rollback

giga/deps/xevm/state/balance.go

@@ -55,7 +55,6 @@ func (s *DBImpl) SubBalance(evmAddr common.Address, amtUint256 *uint256.Int, rea
 	surplus := sdk.NewIntFromBigInt(amt)
 	s.tempState.surplus = s.tempState.surplus.Add(surplus)
 	s.journal = append(s.journal, &surplusChange{delta: surplus})
-	s.journal = append(s.journal, &balanceChange{evmAddr: evmAddr, seiAddr: addr, usei: usei, wei: wei, isAdd: false})
 	return *ZeroInt
 }
 
@@ -98,7 +97,6 @@ func (s *DBImpl) AddBalance(evmAddr common.Address, amtUint256 *uint256.Int, rea
 	surplus := sdk.NewIntFromBigInt(amt).Neg()
 	s.tempState.surplus = s.tempState.surplus.Add(surplus)
 	s.journal = append(s.journal, &surplusChange{delta: surplus})
-	s.journal = append(s.journal, &balanceChange{evmAddr: evmAddr, seiAddr: addr, usei: usei, wei: wei, isAdd: true})
 	return *ZeroInt
 }
 

giga/deps/xevm/state/code.go

@@ -1,11 +1,8 @@
 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 {
@@ -18,10 +15,6 @@ 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)
@@ -30,12 +23,6 @@ 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:       slices.Clone(prevCode),
-		prevCodeExists: prevCodeExists,
-		prevMapping:    prevMapping,
-	})
 	return oldCode
 }
 

giga/deps/xevm/state/journal.go

@@ -2,7 +2,6 @@ package state
 
 import (
 	"encoding/binary"
-	"fmt"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
@@ -19,6 +18,7 @@ type journalEntry interface {
 type revision struct {
 	id           int
 	journalIndex int
+	ctxIndex     int
 }
 
 type (
@@ -52,6 +52,10 @@ type (
 		delta sdk.Int
 	}
 
+	watermark struct {
+		revision int
+	}
+
 	// storageChange records a KV storage mutation so it can be reverted.
 	storageChange struct {
 		addr common.Address
@@ -180,54 +184,19 @@ func (e *accountStatusChange) revert(s *DBImpl) {
 	}
 }
 
-func (e *storageChange) revert(s *DBImpl) {
-	s.k.SetState(s.ctx, e.addr, e.key, e.prev)
-}
+func (e *watermark) revert(s *DBImpl) {}
 
-func (e *codeChange) revert(s *DBImpl) {
-	restoreCode(s, e.addr, e.prevCode, e.prevCodeExists)
-	e.prevMapping.restore(s, e.addr)
-}
+func (e *storageChange) revert(s *DBImpl) {}
 
-func (e *nonceChange) revert(s *DBImpl) {
-	restoreNonce(s, e.addr, e.prev, e.prevExists)
-}
+func (e *codeChange) revert(s *DBImpl) {}
 
-func (e *balanceChange) revert(s *DBImpl) {
-	// Suppress events on revert
-	ctx := s.ctx.WithEventManager(sdk.NewEventManager())
-	denom := s.k.GetBaseDenom(s.ctx)
-	if e.isAdd {
-		// Was AddBalance: reverse by subtracting
-		if err := s.k.BankKeeper().SubUnlockedCoins(ctx, e.seiAddr, sdk.NewCoins(sdk.NewCoin(denom, e.usei)), true); err != nil {
-			panic(fmt.Sprintf("balanceChange revert SubUnlockedCoins: %v", err))
-		}
-		if err := s.k.BankKeeper().SubWei(ctx, e.seiAddr, e.wei); err != nil {
-			panic(fmt.Sprintf("balanceChange revert SubWei: %v", err))
-		}
-	} else {
-		// Was SubBalance: reverse by adding
-		if err := s.k.BankKeeper().AddCoins(ctx, e.seiAddr, sdk.NewCoins(sdk.NewCoin(denom, e.usei)), true); err != nil {
-			panic(fmt.Sprintf("balanceChange revert AddCoins: %v", err))
-		}
-		if err := s.k.BankKeeper().AddWei(ctx, e.seiAddr, e.wei); err != nil {
-			panic(fmt.Sprintf("balanceChange revert AddWei: %v", err))
-		}
-	}
-}
+func (e *nonceChange) revert(s *DBImpl) {}
 
-func (e *createAccountChange) revert(s *DBImpl) {
-	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)
-	}
-}
+func (e *balanceChange) revert(s *DBImpl) {}
 
-func (e *deleteMappingChange) revert(s *DBImpl) {
-	ctx := s.ctx.WithEventManager(sdk.NewEventManager())
-	s.k.SetAddressMapping(ctx, e.seiAddr, e.evmAddr)
-}
+func (e *createAccountChange) revert(s *DBImpl) {}
+
+func (e *deleteMappingChange) revert(s *DBImpl) {}
 
 func captureAddressMapping(s *DBImpl, addr common.Address) addressMappingState {
 	seiAddr, ok := s.k.GetSeiAddress(s.ctx, addr)

giga/deps/xevm/state/nonce.go

@@ -3,7 +3,6 @@ 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 {
@@ -12,12 +11,10 @@ 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, prevExists: prevExists})
 }

giga/deps/xevm/state/state.go

@@ -2,7 +2,6 @@ package state
 
 import (
 	"bytes"
-	"slices"
 	"sort"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -15,11 +14,11 @@ import (
 )
 
 func (s *DBImpl) CreateAccount(acc common.Address) {
-	// clear any existing state but keep balance untouched, journaled for revert
+	// clear any existing state but keep balance untouched
 	if !s.ctx.IsTracing() {
 		// too slow on historical DB so not doing it for tracing for now.
 		// could cause tracing to be incorrect in theory.
-		s.clearAccountStateJournaled(acc)
+		s.clearAccountState(acc)
 	}
 	s.MarkAccount(acc, AccountCreated)
 }
@@ -43,7 +42,6 @@ func (s *DBImpl) SetState(addr common.Address, key common.Hash, val common.Hash)
 	}
 
 	s.k.SetState(s.ctx, addr, key, val)
-	s.journal = append(s.journal, &storageChange{addr: addr, key: key, prev: old})
 	return old
 }
 
@@ -77,7 +75,6 @@ func (s *DBImpl) SelfDestruct(acc common.Address) uint256.Int {
 	if seiAddr, ok := s.k.GetSeiAddress(s.ctx, acc); ok {
 		// remove the association
 		s.k.DeleteAddressMapping(s.ctx, seiAddr, acc)
-		s.journal = append(s.journal, &deleteMappingChange{evmAddr: acc, seiAddr: seiAddr})
 	}
 	b := s.GetBalance(acc)
 	s.SubBalance(acc, b, tracing.BalanceDecreaseSelfdestruct)
@@ -105,25 +102,22 @@ func (s *DBImpl) HasSelfDestructed(acc common.Address) bool {
 	return bytes.Equal(val, AccountDeleted)
 }
 
-// Snapshot records the current journal length as a revision and pushes the current
-// EventManager onto the stack, creating a fresh one for subsequent events.
 func (s *DBImpl) Snapshot() int {
 	id := s.nextRevisionId
 	s.nextRevisionId++
+	newCtx := s.ctx.WithMultiStore(s.ctx.MultiStore().CacheMultiStore()).WithEventManager(sdk.NewEventManager())
+	s.snapshottedCtxs = append(s.snapshottedCtxs, s.ctx)
 	s.validRevisions = append(s.validRevisions, revision{
 		id:           id,
 		journalIndex: len(s.journal),
+		ctxIndex:     len(s.snapshottedCtxs) - 1,
 	})
-	// Push current EM and create a fresh one so reverted events are discarded.
-	s.snapshottedEventManagers = append(s.snapshottedEventManagers, s.ctx.EventManager())
-	s.ctx = s.ctx.WithEventManager(sdk.NewEventManager())
+	s.ctx = newCtx
+	s.journal = append(s.journal, &watermark{revision: id})
 	return id
 }
 
-// RevertToSnapshot reverts all journal entries back to the snapshot identified by rev,
-// restores the EventManager, and truncates the revision list.
 func (s *DBImpl) RevertToSnapshot(rev int) {
-	// Binary-search for the revision with the given id (like go-ethereum).
 	idx := sort.Search(len(s.validRevisions), func(i int) bool {
 		return s.validRevisions[i].id >= rev
 	})
@@ -132,18 +126,14 @@ func (s *DBImpl) RevertToSnapshot(rev int) {
 	}
 	snapshot := s.validRevisions[idx]
 
-	// Revert journal entries in reverse order down to the snapshot point.
+	s.ctx = s.snapshottedCtxs[snapshot.ctxIndex]
+	s.snapshottedCtxs = s.snapshottedCtxs[:snapshot.ctxIndex]
+
 	for i := len(s.journal) - 1; i >= snapshot.journalIndex; i-- {
 		s.journal[i].revert(s)
 	}
 	s.journal = s.journal[:snapshot.journalIndex]
 
-	// Restore the EventManager that was active when the snapshot was taken.
-	// snapshottedEventManagers has one entry per snapshot; idx corresponds to this snapshot.
-	s.ctx = s.ctx.WithEventManager(s.snapshottedEventManagers[idx])
-	s.snapshottedEventManagers = s.snapshottedEventManagers[:idx]
-
-	// Truncate the revision list (removing this snapshot and any taken after it).
 	s.validRevisions = s.validRevisions[:idx]
 }
 
@@ -184,56 +174,6 @@ func (s *DBImpl) clearAccountState(acc common.Address) {
 	}
 }
 
-// clearAccountStateJournaled wipes code, nonce, and storage for acc, recording
-// the previous values in the journal so a RevertToSnapshot can restore them.
-// Called from CreateAccount (when not tracing).
-func (s *DBImpl) clearAccountStateJournaled(acc common.Address) {
-	// Only clear if a code hash exists (mirrors clearAccountState logic).
-	codeHashStore := s.k.PrefixStore(s.ctx, types.CodeHashKeyPrefix)
-	if !codeHashStore.Has(acc[:]) {
-		return
-	}
-
-	// Save previous state for potential revert.
-	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,
-	// so we strip the per-address state prefix to obtain each slot hash.
-	prevSlots := make(map[common.Hash]common.Hash)
-	statePrefix := types.StateKey(acc)
-	stateStore := s.k.PrefixStore(s.ctx, statePrefix)
-	rawKeys := stateStore.GetAllKeyStrsInRange(nil, nil)
-	prefixLen := len(statePrefix)
-	for _, raw := range rawKeys {
-		if len(raw) <= prefixLen {
-			continue
-		}
-		slotKey := common.BytesToHash([]byte(raw)[prefixLen:])
-		slotVal := s.k.GetState(s.ctx, acc, slotKey)
-		if slotVal != (common.Hash{}) {
-			prevSlots[slotKey] = slotVal
-		}
-	}
-
-	// Append journal entry before making changes.
-	s.journal = append(s.journal, &createAccountChange{
-		addr:            acc,
-		prevCode:        slices.Clone(prevCode),
-		prevCodeExists:  prevCodeExists,
-		prevNonce:       prevNonce,
-		prevNonceExists: prevNonceExists,
-		prevSlots:       prevSlots,
-	})
-
-	// Clear the account state.
-	s.clearAccountState(acc)
-}
-
 func (s *DBImpl) MarkAccount(acc common.Address, status []byte) {
 	prev, ok := s.tempState.transientAccounts[acc.Hex()]
 	if !ok {

giga/deps/xevm/state/state_test.go

@@ -17,14 +17,20 @@ import (
 
 type countingCacheKVStore struct {
 	sdk.CacheKVStore
-	sets int
+	sets    int
+	deletes int
 }
 
 func (s *countingCacheKVStore) Set(key, value []byte) {
 	s.sets++
 	s.CacheKVStore.Set(key, value)
 }
 
+func (s *countingCacheKVStore) Delete(key []byte) {
+	s.deletes++
+	s.CacheKVStore.Delete(key)
+}
+
 func TestState(t *testing.T) {
 	k, ctx := testkeeper.MockEVMKeeper(t)
 	ctx = ctx.WithBlockTime(time.Now())
@@ -336,3 +342,66 @@ func TestSetState_NoopStillWritesThrough(t *testing.T) {
 	sdb.SetState(evmAddr, key, val)
 	require.Equal(t, 1, counter.sets)
 }
+
+func TestSnapshotRevertDoesNotFlushStorageWrites(t *testing.T) {
+	tests := []struct {
+		name     string
+		writeVal common.Hash
+	}{
+		{
+			name:     "same value",
+			writeVal: common.BytesToHash([]byte("v")),
+		},
+		{
+			name:     "changed value",
+			writeVal: common.BytesToHash([]byte("changed")),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			k, ctx := testkeeper.MockEVMKeeper(t)
+			ctx = ctx.WithBlockTime(time.Now())
+			_, evmAddr := testkeeper.MockAddressPair()
+
+			key := common.BytesToHash([]byte("k"))
+			val := common.BytesToHash([]byte("v"))
+			seedDB := state.NewDBImpl(ctx, k, false)
+			seedDB.SetState(evmAddr, key, val)
+			_, err := seedDB.Finalize()
+			require.NoError(t, err)
+
+			counter := &countingCacheKVStore{}
+			ms := ctx.MultiStore().(interface {
+				SetKVStores(func(sdk.StoreKey, sdk.KVStore) sdk.CacheWrap) sdk.MultiStore
+			}).SetKVStores(func(sk sdk.StoreKey, store sdk.KVStore) sdk.CacheWrap {
+				if sk == k.GetStoreKey() {
+					cacheStore, ok := store.(sdk.CacheKVStore)
+					require.True(t, ok)
+					counter.CacheKVStore = cacheStore
+					return counter
+				}
+				cacheWrap, ok := store.(sdk.CacheWrap)
+				require.True(t, ok)
+				return cacheWrap
+			})
+			ctx = ctx.WithMultiStore(ms)
+			require.NotNil(t, counter.CacheKVStore)
+
+			sdb := state.NewDBImpl(ctx, k, false)
+			rev := sdb.Snapshot()
+			sdb.SetState(evmAddr, key, tt.writeVal)
+			require.Equal(t, tt.writeVal, sdb.GetState(evmAddr, key))
+
+			sdb.RevertToSnapshot(rev)
+			require.Equal(t, val, sdb.GetState(evmAddr, key))
+
+			_, err = sdb.Finalize()
+			require.NoError(t, err)
+			require.Zero(t, counter.sets+counter.deletes)
+
+			checkDB := state.NewDBImpl(ctx, k, false)
+			require.Equal(t, val, checkDB.GetState(evmAddr, key))
+		})
+	}
+}