core/state: storage journal entry should revert dirtyness too

core/state/journal.go

@@ -139,8 +139,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
@@ -307,7 +308,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 {
@@ -316,9 +317,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,
 	}
 }
 

core/state/state_object.go

@@ -157,13 +157,20 @@ func (s *stateObject) getTrie() Trie {
 
 // 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.
@@ -223,23 +230,37 @@ 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,
 	})
-	s.setState(key, value)
+
+	s.setState(key, &value)
 }
 
-func (s *stateObject) setState(key, value common.Hash) {
-	//s.cachedStorage[key] = value
-	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