From ad5ff70a9b2592a8031979607ef78559924ccce3 Mon Sep 17 00:00:00 2001
From: Jared Wasinger <j-wasinger@hotmail.com>
Date: Thu, 7 Dec 2023 23:44:05 +0800
Subject: [PATCH] prefetch all accessed accounts and storage slots regardless
 of whether snapshot is running. in trie prefetcher, a subfetcher is not
 halted until a call to triePrefetcher.trie to receive the trie that it
 corresponds to.

---
 core/state/state_object.go    |  5 ++++-
 core/state/statedb.go         |  5 +++++
 core/state/trie_prefetcher.go | 12 ++++++------
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/core/state/state_object.go b/core/state/state_object.go
index 6444690727..a5463cfc2b 100644
--- a/core/state/state_object.go
+++ b/core/state/state_object.go
@@ -204,6 +204,10 @@ func (s *stateObject) GetCommittedState(key common.Hash) common.Hash {
 			value.SetBytes(content)
 		}
 	}
+	if s.db.prefetcher != nil {
+		// always prefetch to ensure that read accounts will end up in the witness
+		s.db.prefetcher.prefetch(s.addrHash, s.data.Root, s.address, [][]byte{key[:]})
+	}
 	// If the snapshot is unavailable or reading from it fails, load from the database.
 	if s.db.snap == nil || err != nil {
 		start := time.Now()
@@ -395,7 +399,6 @@ func (s *stateObject) commit() (*trienode.NodeSet, map[string][]byte, error) {
 		return nil, nil, err
 	}
 	s.data.Root = root
-	fmt.Printf("node count %d\n", len(nodes.Leaves))
 
 	// Update original account data after commit
 	s.origin = s.data.Copy()
diff --git a/core/state/statedb.go b/core/state/statedb.go
index 3d6bcd1369..46c6056472 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -562,6 +562,11 @@ func (s *StateDB) getStateObject(addr common.Address) *stateObject {
 // flag set. This is needed by the state journal to revert to the correct s-
 // destructed object instead of wiping all knowledge about the state object.
 func (s *StateDB) getDeletedStateObject(addr common.Address) *stateObject {
+	if s.prefetcher != nil {
+		// always prefetch to ensure written/read accounts appear in the witness
+		// regardless of whether snapshot is enabled
+		s.prefetcher.prefetch(common.Hash{}, s.originalRoot, common.Address{}, [][]byte{addr[:]})
+	}
 	// Prefer live objects if any is available
 	if obj := s.stateObjects[addr]; obj != nil {
 		return obj
diff --git a/core/state/trie_prefetcher.go b/core/state/trie_prefetcher.go
index 716969e2c1..d24f18781c 100644
--- a/core/state/trie_prefetcher.go
+++ b/core/state/trie_prefetcher.go
@@ -208,7 +208,7 @@ type subfetcher struct {
 	lock    sync.Mutex // Lock protecting the task queue
 	closing bool       // set to true if the subfetcher is closing
 
-	wake chan struct{} // Wake channel if a new task is scheduled
+	wake chan bool     // Wake channel if a new task is scheduled, true if the subfetcher should continue running when there are no pending tasks
 	term chan struct{} // Channel to signal interruption
 
 	seen map[string]struct{} // Tracks the entries already loaded
@@ -225,7 +225,7 @@ func newSubfetcher(db Database, state common.Hash, owner common.Hash, root commo
 		owner: owner,
 		root:  root,
 		addr:  addr,
-		wake:  make(chan struct{}, 1),
+		wake:  make(chan bool, 1),
 		term:  make(chan struct{}),
 		seen:  make(map[string]struct{}),
 	}
@@ -240,7 +240,7 @@ func (sf *subfetcher) schedule(keys [][]byte) {
 	sf.tasks = append(sf.tasks, keys...)
 	sf.lock.Unlock()
 	// Notify the prefetcher. The wake-chan is buffered, so this is async.
-	sf.wake <- struct{}{}
+	sf.wake <- true
 }
 
 // wait waits for the subfetcher to finish it's task. It is safe to call wait multiple
@@ -255,7 +255,7 @@ func (sf *subfetcher) wait() {
 	sf.tasks = nil
 	sf.lock.Unlock()
 	// Notify the prefetcher. The wake-chan is buffered, so this is async.
-	sf.wake <- struct{}{}
+	sf.wake <- false
 	// Wait for it to terminate
 	<-sf.term
 }
@@ -286,13 +286,13 @@ func (sf *subfetcher) loop() {
 	}
 	// Trie opened successfully, keep prefetching items
 	for {
-		<-sf.wake
+		keepRunning := <-sf.wake
 		// Subfetcher was woken up, retrieve any tasks to avoid spinning the lock
 		sf.lock.Lock()
 		tasks := sf.tasks
 		sf.tasks = nil
 		sf.lock.Unlock()
-		if tasks == nil {
+		if tasks == nil && !keepRunning {
 			// No more tasks
 			return
 		}