From f08695c1d54415d79146999faab155c9bc684f58 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Fri, 23 Apr 2021 13:39:18 +0200
Subject: [PATCH] trie: reuse dirty data instead of hitting disk when
 generating #22667
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* core/state/snapshot: reuse memory data instead of hitting disk when generating

* trie: minor nitpicks wrt the resolver optimization

* core/state/snapshot, trie: use key/value store for resolver

* trie: fix linter

Co-authored-by: Péter Szilágyi <peterke@gmail.com>
---
 trie/iterator.go | 52 ++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 46 insertions(+), 6 deletions(-)

diff --git a/trie/iterator.go b/trie/iterator.go
index 5cf9b448f34a..13457e4d848a 100644
--- a/trie/iterator.go
+++ b/trie/iterator.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/XinFinOrg/XDPoSChain/common"
 	"github.com/XinFinOrg/XDPoSChain/core/types"
+	"github.com/XinFinOrg/XDPoSChain/ethdb"
 )
 
 // Iterator is a key-value trie iterator that traverses a Trie.
@@ -102,6 +103,19 @@ type NodeIterator interface {
 	// iterator is not positioned at a leaf. Callers must not retain references
 	// to the value after calling Next.
 	LeafProof() [][]byte
+
+	// AddResolver sets an intermediate database to use for looking up trie nodes
+	// before reaching into the real persistent layer.
+	//
+	// This is not required for normal operation, rather is an optimization for
+	// cases where trie nodes can be recovered from some external mechanism without
+	// reading from disk. In those cases, this resolver allows short circuiting
+	// accesses and returning them from memory.
+	//
+	// Before adding a similar mechanism to any other place in Geth, consider
+	// making trie.Database an interface and wrapping at that level. It's a huge
+	// refactor, but it could be worth it if another occurrence arises.
+	AddResolver(ethdb.KeyValueStore)
 }
 
 // nodeIteratorState represents the iteration state at one particular Node of the
@@ -119,6 +133,8 @@ type nodeIterator struct {
 	stack []*nodeIteratorState // Hierarchy of trie nodes persisting the iteration state
 	path  []byte               // Path to the current Node
 	err   error                // Failure set in case of an internal error in the iterator
+
+	resolver ethdb.KeyValueStore // Optional intermediate resolver above the disk layer
 }
 
 // errIteratorEnd is stored in nodeIterator.err when iteration is done.
@@ -143,6 +159,10 @@ func newNodeIterator(trie *Trie, start []byte) NodeIterator {
 	return it
 }
 
+func (it *nodeIterator) AddResolver(resolver ethdb.KeyValueStore) {
+	it.resolver = resolver
+}
+
 func (it *nodeIterator) Hash() common.Hash {
 	if len(it.stack) == 0 {
 		return common.Hash{}
@@ -261,7 +281,7 @@ func (it *nodeIterator) init() (*nodeIteratorState, error) {
 	if root != types.EmptyRootHash {
 		state.hash = root
 	}
-	return state, state.resolve(it.trie, nil)
+	return state, state.resolve(it, nil)
 }
 
 // peek creates the next state of the iterator.
@@ -285,7 +305,7 @@ func (it *nodeIterator) peek(descend bool) (*nodeIteratorState, *int, []byte, er
 		}
 		state, path, ok := it.nextChild(parent, ancestor)
 		if ok {
-			if err := state.resolve(it.trie, path); err != nil {
+			if err := state.resolve(it, path); err != nil {
 				return parent, &parent.index, path, err
 			}
 			return state, &parent.index, path, nil
@@ -318,7 +338,7 @@ func (it *nodeIterator) peekSeek(seekKey []byte) (*nodeIteratorState, *int, []by
 		}
 		state, path, ok := it.nextChildAt(parent, ancestor, seekKey)
 		if ok {
-			if err := state.resolve(it.trie, path); err != nil {
+			if err := state.resolve(it, path); err != nil {
 				return parent, &parent.index, path, err
 			}
 			return state, &parent.index, path, nil
@@ -329,9 +349,21 @@ func (it *nodeIterator) peekSeek(seekKey []byte) (*nodeIteratorState, *int, []by
 	return nil, nil, nil, errIteratorEnd
 }
 
-func (st *nodeIteratorState) resolve(tr *Trie, path []byte) error {
+func (it *nodeIterator) resolveHash(hash hashNode, path []byte) (node, error) {
+	if it.resolver != nil {
+		if blob, err := it.resolver.Get(hash); err == nil && len(blob) > 0 {
+			if resolved, err := decodeNode(hash, blob); err == nil {
+				return resolved, nil
+			}
+		}
+	}
+	resolved, err := it.trie.resolveHash(hash, path)
+	return resolved, err
+}
+
+func (st *nodeIteratorState) resolve(it *nodeIterator, path []byte) error {
 	if hash, ok := st.node.(hashNode); ok {
-		resolved, err := tr.resolveHash(hash, path)
+		resolved, err := it.resolveHash(hash, path)
 		if err != nil {
 			return err
 		}
@@ -516,6 +548,10 @@ func (it *differenceIterator) Path() []byte {
 	return it.b.Path()
 }
 
+func (it *differenceIterator) AddResolver(resolver ethdb.KeyValueStore) {
+	panic("not implemented")
+}
+
 func (it *differenceIterator) Next(bool) bool {
 	// Invariants:
 	// - We always advance at least one element in b.
@@ -623,7 +659,11 @@ func (it *unionIterator) Path() []byte {
 	return (*it.items)[0].Path()
 }
 
-// Next returns the next Node in the union of tries being iterated over.
+func (it *unionIterator) AddResolver(resolver ethdb.KeyValueStore) {
+	panic("not implemented")
+}
+
+// Next returns the next node in the union of tries being iterated over.
 //
 // It does this by maintaining a heap of iterators, sorted by the iteration
 // order of their next elements, with one entry for each source trie. Each