Efficient reverts clean dirty state cache

core/state/multi_tx_snapshot.go

@@ -29,6 +29,10 @@ type MultiTxSnapshot struct {
 	accountNotPending map[common.Address]struct{}
 	accountNotDirty   map[common.Address]struct{}
 
+	// touched accounts are accounts that can be affected when snapshot is reverted
+	// we clear dirty storage for touched accounts when snapshot is reverted
+	touchedAccounts map[common.Address]struct{}
+
 	// TODO: snapdestructs, snapaccount storage
 }
 
@@ -51,6 +55,7 @@ func newMultiTxSnapshot() MultiTxSnapshot {
 		accountDeleted:    make(map[common.Address]bool),
 		accountNotPending: make(map[common.Address]struct{}),
 		accountNotDirty:   make(map[common.Address]struct{}),
+		touchedAccounts:   make(map[common.Address]struct{}),
 	}
 }
 
@@ -142,6 +147,7 @@ func (s *MultiTxSnapshot) objectChanged(address common.Address) bool {
 
 // updateBalanceChange updates the snapshot with the balance change.
 func (s *MultiTxSnapshot) updateBalanceChange(change balanceChange) {
+	s.touchedAccounts[*change.account] = struct{}{}
 	if s.objectChanged(*change.account) {
 		return
 	}
@@ -152,6 +158,7 @@ func (s *MultiTxSnapshot) updateBalanceChange(change balanceChange) {
 
 // updateNonceChange updates the snapshot with the nonce change.
 func (s *MultiTxSnapshot) updateNonceChange(change nonceChange) {
+	s.touchedAccounts[*change.account] = struct{}{}
 	if s.objectChanged(*change.account) {
 		return
 	}
@@ -162,6 +169,7 @@ func (s *MultiTxSnapshot) updateNonceChange(change nonceChange) {
 
 // updateCodeChange updates the snapshot with the code change.
 func (s *MultiTxSnapshot) updateCodeChange(change codeChange) {
+	s.touchedAccounts[*change.account] = struct{}{}
 	if s.objectChanged(*change.account) {
 		return
 	}
@@ -173,6 +181,7 @@ func (s *MultiTxSnapshot) updateCodeChange(change codeChange) {
 
 // updateResetObjectChange updates the snapshot with the reset object change.
 func (s *MultiTxSnapshot) updateResetObjectChange(change resetObjectChange) {
+	s.touchedAccounts[change.prev.address] = struct{}{}
 	address := change.prev.address
 	if _, ok := s.prevObjects[address]; !ok {
 		s.prevObjects[address] = change.prev
@@ -181,13 +190,15 @@ func (s *MultiTxSnapshot) updateResetObjectChange(change resetObjectChange) {
 
 // updateCreateObjectChange updates the snapshot with the createObjectChange.
 func (s *MultiTxSnapshot) updateCreateObjectChange(change createObjectChange) {
+	s.touchedAccounts[*change.account] = struct{}{}
 	if _, ok := s.prevObjects[*change.account]; !ok {
 		s.prevObjects[*change.account] = nil
 	}
 }
 
 // updateSuicideChange updates the snapshot with the suicide change.
 func (s *MultiTxSnapshot) updateSuicideChange(change suicideChange) {
+	s.touchedAccounts[*change.account] = struct{}{}
 	if s.objectChanged(*change.account) {
 		return
 	}
@@ -201,6 +212,7 @@ func (s *MultiTxSnapshot) updateSuicideChange(change suicideChange) {
 
 // updatePendingStorage updates the snapshot with the pending storage change.
 func (s *MultiTxSnapshot) updatePendingStorage(address common.Address, key, value common.Hash, ok bool) {
+	s.touchedAccounts[address] = struct{}{}
 	if s.objectChanged(address) {
 		return
 	}
@@ -219,6 +231,7 @@ func (s *MultiTxSnapshot) updatePendingStorage(address common.Address, key, valu
 
 // updatePendingStatus updates the snapshot with previous pending status.
 func (s *MultiTxSnapshot) updatePendingStatus(address common.Address, pending, dirty bool) {
+	s.touchedAccounts[address] = struct{}{}
 	if !pending {
 		s.accountNotPending[address] = struct{}{}
 	}
@@ -229,6 +242,7 @@ func (s *MultiTxSnapshot) updatePendingStatus(address common.Address, pending, d
 
 // updateObjectDeleted updates the snapshot with the object deletion.
 func (s *MultiTxSnapshot) updateObjectDeleted(address common.Address, deleted bool) {
+	s.touchedAccounts[address] = struct{}{}
 	if s.objectChanged(address) {
 		return
 	}
@@ -358,6 +372,10 @@ func (s *MultiTxSnapshot) Merge(other *MultiTxSnapshot) error {
 		}
 	}
 
+	for address := range other.touchedAccounts {
+		s.touchedAccounts[address] = struct{}{}
+	}
+
 	return nil
 }
 
@@ -428,6 +446,13 @@ func (s *MultiTxSnapshot) revertState(st *StateDB) {
 	for address := range s.accountNotDirty {
 		delete(st.stateObjectsDirty, address)
 	}
+
+	// clean dirty state of touched accounts
+	for address := range s.touchedAccounts {
+		if obj, ok := st.stateObjects[address]; ok {
+			obj.dirtyStorage = make(Storage)
+		}
+	}
 }
 
 // MultiTxSnapshotStack contains a list of snapshots for multiple transactions associated with a StateDB.

core/state/multi_tx_snapshot_test.go

@@ -346,6 +346,25 @@ func TestMultiTxSnapshotAccountChangesSimple(t *testing.T) {
 	})
 }
 
+// This test verifies that dirty account storage is properly cleaned for accounts after revert
+func TestMultiTxSnapshotAccountChangesRevertedByJournal(t *testing.T) {
+	testMultiTxSnapshot(t, func(s *StateDB) {
+		for _, addr := range addrs {
+			s.SetState(addr, common.HexToHash("0x01"), common.HexToHash("0x03"))
+		}
+		s.Finalise(true)
+		for _, addr := range addrs {
+			// we use normal snapshot here because it
+			// 1. does not mark an account dirty (even though we applied changes)
+			// 2. changes dirty, uncommitted state of the account
+			snap := s.Snapshot()
+			s.SetState(addr, common.HexToHash("0x01"), common.HexToHash("0x02"))
+			s.RevertToSnapshot(snap)
+		}
+		s.Finalise(true)
+	})
+}
+
 func TestMultiTxSnapshotRefund(t *testing.T) {
 	testMultiTxSnapshot(t, func(s *StateDB) {
 		for _, addr := range addrs {