From 4f4f9d88d3a59e972c9fe3a5e6e6a93dad2e7db4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A9ter=20Szil=C3=A1gyi?= <peterke@gmail.com>
Date: Wed, 24 Apr 2024 18:45:24 +0300
Subject: [PATCH] core/state: storage journal entry should revert dirtyness too
 (#29641)

Currently our state journal tracks each storage update to a contract, having the ability to revert those changes to the previously set value.

For the very first modification however, it behaves a bit wonky. Reverting the update doesn't actually remove the dirty-ness of the slot, rather leaves it as "change this slot to it's original value". This can cause issues down the line with for example write witnesses needing to gather an unneeded proof.

This PR modifies the storageChange journal entry to not only track the previous value of a slot, but also whether there was any previous value at all set in the current execution context. In essence, the PR changes the semantic of storageChange so it does not simply track storage changes, rather it tracks dirty storage changes, an important distinction for being able to cleanly revert the journal item.
---
 core/state/journal.go      | 13 +++++++------
 core/state/state_object.go | 40 ++++++++++++++++++++++++++++----------
 2 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/core/state/journal.go b/core/state/journal.go
index cfd3782eb0..c0f5615c98 100644
--- a/core/state/journal.go
+++ b/core/state/journal.go
@@ -129,8 +129,9 @@ type (
 		prev    uint64
 	}
 	storageChange struct {
-		account       *common.Address
-		key, prevalue common.Hash
+		account   *common.Address
+		key       common.Hash
+		prevvalue *common.Hash
 	}
 	codeChange struct {
 		account            *common.Address
@@ -277,7 +278,7 @@ func (ch codeChange) copy() journalEntry {
 }
 
 func (ch storageChange) revert(s *StateDB) {
-	s.getStateObject(*ch.account).setState(ch.key, ch.prevalue)
+	s.getStateObject(*ch.account).setState(ch.key, ch.prevvalue)
 }
 
 func (ch storageChange) dirtied() *common.Address {
@@ -286,9 +287,9 @@ func (ch storageChange) dirtied() *common.Address {
 
 func (ch storageChange) copy() journalEntry {
 	return storageChange{
-		account:  ch.account,
-		key:      ch.key,
-		prevalue: ch.prevalue,
+		account:   ch.account,
+		key:       ch.key,
+		prevvalue: ch.prevvalue,
 	}
 }
 
diff --git a/core/state/state_object.go b/core/state/state_object.go
index 8978b30542..aa748f08ac 100644
--- a/core/state/state_object.go
+++ b/core/state/state_object.go
@@ -140,13 +140,20 @@ func (s *stateObject) getTrie() (Trie, error) {
 
 // GetState retrieves a value from the account storage trie.
 func (s *stateObject) GetState(key common.Hash) common.Hash {
+	value, _ := s.getState(key)
+	return value
+}
+
+// getState retrieves a value from the account storage trie and also returns if
+// the slot is already dirty or not.
+func (s *stateObject) getState(key common.Hash) (common.Hash, bool) {
 	// If we have a dirty value for this state entry, return it
 	value, dirty := s.dirtyStorage[key]
 	if dirty {
-		return value
+		return value, true
 	}
 	// Otherwise return the entry's original value
-	return s.GetCommittedState(key)
+	return s.GetCommittedState(key), false
 }
 
 // GetCommittedState retrieves a value from the committed account storage trie.
@@ -209,25 +216,38 @@ func (s *stateObject) GetCommittedState(key common.Hash) common.Hash {
 
 // SetState updates a value in account storage.
 func (s *stateObject) SetState(key, value common.Hash) {
-	// If the new value is the same as old, don't set
-	prev := s.GetState(key)
+	// If the new value is the same as old, don't set. Otherwise, track only the
+	// dirty changes, supporting reverting all of it back to no change.
+	prev, dirty := s.getState(key)
 	if prev == value {
 		return
 	}
+	var prevvalue *common.Hash
+	if dirty {
+		prevvalue = &prev
+	}
 	// New value is different, update and journal the change
 	s.db.journal.append(storageChange{
-		account:  &s.address,
-		key:      key,
-		prevalue: prev,
+		account:   &s.address,
+		key:       key,
+		prevvalue: prevvalue,
 	})
 	if s.db.logger != nil && s.db.logger.OnStorageChange != nil {
 		s.db.logger.OnStorageChange(s.address, key, prev, value)
 	}
-	s.setState(key, value)
+	s.setState(key, &value)
 }
 
-func (s *stateObject) setState(key, value common.Hash) {
-	s.dirtyStorage[key] = value
+// setState updates a value in account dirty storage. If the value being set is
+// nil (assuming journal revert), the dirtyness is removed.
+func (s *stateObject) setState(key common.Hash, value *common.Hash) {
+	// If the first set is being reverted, undo the dirty marker
+	if value == nil {
+		delete(s.dirtyStorage, key)
+		return
+	}
+	// Otherwise restore the previous value
+	s.dirtyStorage[key] = *value
 }
 
 // finalise moves all dirty storage slots into the pending area to be hashed or