From 8065f98b3686e0ca0050b25f3641fe9c09b101ea Mon Sep 17 00:00:00 2001
From: Jared Wasinger <j-wasinger@hotmail.com>
Date: Mon, 11 Mar 2024 13:18:47 -0700
Subject: [PATCH] simplify updateTrie

Co-authored-by: Martin HS <martin@swende.se>
---
 core/state/state_object.go | 88 ++++++++++++++------------------------
 1 file changed, 32 insertions(+), 56 deletions(-)

diff --git a/core/state/state_object.go b/core/state/state_object.go
index be2b49d269..d3d7765256 100644
--- a/core/state/state_object.go
+++ b/core/state/state_object.go
@@ -294,16 +294,39 @@ func (s *stateObject) updateTrie() (Trie, error) {
 	}
 	// Insert all the pending storage updates into the trie
 	usedStorage := make([][]byte, 0, len(s.pendingStorage))
-	// insertStorageChange updates a storage slot
-	insertStorageChange := func(key, value common.Hash, updateTr func() ([]byte, error)) error {
+
+	// Perform trie updates before deletions.  This prevents resolution of unnecessary trie nodes
+	//  in circumstances similar to the following:
+	//
+	// Consider nodes `A` and `B` who share the same full node parent `P` and have no other siblings.
+	// During the execution of a block:
+	// - `A` is deleted,
+	// - `C` is created, and also shares the parent `P`.
+	// If the deletion is handled first, then `P` would be left with only one child, thus collapsed
+	// into a shortnode. This requires `B` to be resolved from disk.
+	// Whereas if the created node is handled first, then the collapse is avoided, and `B` is not resolved.
+	var deletions []common.Hash
+	for key, value := range s.pendingStorage {
+		// Skip noop changes, persist actual changes
+		if value == s.originStorage[key] {
+			continue
+		}
 		prev := s.originStorage[key]
 		s.originStorage[key] = value
 
 		var encoded []byte // rlp-encoded value to be used by the snapshot
-		if encoded, err = updateTr(); err != nil {
-			return err
+		if (value != common.Hash{}) {
+			// Encoding []byte cannot fail, ok to ignore the error.
+			trimmed := common.TrimLeftZeroes(value[:])
+			encoded, _ = rlp.EncodeToBytes(trimmed)
+			if err := tr.UpdateStorage(s.address, key[:], trimmed); err != nil {
+				s.db.setError(err)
+				return nil, err
+			}
+			s.db.StorageUpdated += 1
+		} else {
+			deletions = append(deletions, key)
 		}
-
 		// Cache the mutated storage slots until commit
 		if storage == nil {
 			if storage = s.db.storages[s.addrHash]; storage == nil {
@@ -333,60 +356,13 @@ func (s *stateObject) updateTrie() (Trie, error) {
 		}
 		// Cache the items for preloading
 		usedStorage = append(usedStorage, common.CopyBytes(key[:])) // Copy needed for closure
-		return nil
 	}
-	applyDeletion := func(key, value common.Hash) error {
-		return insertStorageChange(key, value, func() (encoded []byte, err error) {
-			if err := tr.DeleteStorage(s.address, key[:]); err != nil {
-				s.db.setError(err)
-				return nil, err
-			}
-			s.db.StorageDeleted += 1
-			return nil, nil
-		})
-	}
-	applyUpdate := func(key, value common.Hash) error {
-		return insertStorageChange(key, value, func() (encoded []byte, err error) {
-			// Encoding []byte cannot fail, ok to ignore the error.
-			trimmed := common.TrimLeftZeroes(value[:])
-			encoded, _ = rlp.EncodeToBytes(trimmed)
-			if err := tr.UpdateStorage(s.address, key[:], trimmed); err != nil {
-				s.db.setError(err)
-				return nil, err
-			}
-			s.db.StorageUpdated += 1
-			return encoded, nil
-		})
-	}
-
-	// Perform trie updates before deletions.  This prevents resolution of unnecessary trie nodes
-	//  in circumstances similar to the following:
-	//
-	// Consider nodes `A` and `B` who share the same full node parent `P` and have no other siblings.
-	// During the execution of a block:
-	// - `A` is deleted,
-	// - `C` is created, and also shares the parent `P`.
-	// If the deletion is handled first, then `P` would be left with only one child, thus collapsed
-	// into a shortnode. This requires `B` to be resolved from disk.
-	// Whereas if the created node is handled first, then the collapse is avoided, and `B` is not resolved.
-	deletedStorages := make(map[common.Hash]common.Hash)
-	for key, value := range s.pendingStorage {
-		// Skip noop changes, persist actual changes
-		if value == s.originStorage[key] {
-			continue
-		}
-		if (value != common.Hash{}) {
-			if err := applyUpdate(key, value); err != nil {
-				return nil, err
-			}
-		} else {
-			deletedStorages[key] = value
-		}
-	}
-	for key, value := range deletedStorages {
-		if err := applyDeletion(key, value); err != nil {
+	for _, key := range deletions {
+		if err := tr.DeleteStorage(s.address, key[:]); err != nil {
+			s.db.setError(err)
 			return nil, err
 		}
+		s.db.StorageDeleted += 1
 	}
 	if s.db.prefetcher != nil {
 		s.db.prefetcher.used(s.addrHash, s.data.Root, usedStorage)