all: reworkNodeResolver for working with multiple state schemes with …

…calling ReadTrieNode underlying (#603)

cmd/ronin/snapshot.go

@@ -185,6 +185,8 @@ block is used.
 	}
 )
 
+// Deprecation: this command should be deprecated once the hash-based
+// scheme is deprecated.
 func pruneState(ctx *cli.Context) error {
 	stack, _ := makeConfigNode(ctx)
 	defer stack.Close()
@@ -392,7 +394,7 @@ func traverseRawState(ctx *cli.Context) error {
 		node := accIter.Hash()
 
 		if node != (common.Hash{}) {
-			// Check the present for non-empty hash node(embedded node doesn't
+			// Check the presence for non-empty hash node(embedded node doesn't
 			// have their own hash).
 			if !rawdb.HasLegacyTrieNode(chaindb, node) {
 				log.Error("Missing trie node(account)", "hash", node)

core/state/snapshot/generate.go

@@ -429,20 +429,25 @@ func (dl *diskLayer) generateRange(trieID *trie.ID, prefix []byte, kind string,
 
 	// We use the snap data to build up a cache which can be used by the
 	// main account trie as a primary lookup when resolving hashes
-	var snapNodeCache ethdb.Database
+	var resolver trie.NodeResolver
 	if len(result.keys) > 0 {
-		snapNodeCache = rawdb.NewMemoryDatabase()
-		snapTrieDb := trie.NewDatabase(snapNodeCache)
-		snapTrie := trie.NewEmpty(snapTrieDb)
+		mdb := rawdb.NewMemoryDatabase()
+		tdb := trie.NewDatabase(mdb)
+		snapTrie := trie.NewEmpty(tdb)
 		for i, key := range result.keys {
 			snapTrie.Update(key, result.vals[i])
 		}
-		root, nodes, _ := snapTrie.Commit(false)
-		if nodes != nil {
-			snapTrieDb.Update(root, types.EmptyRootHash, 0, trienode.NewWithNodeSet(nodes), nil)
+		root, nodes, err := snapTrie.Commit(false)
+
+		if err != nil && nodes != nil {
+			tdb.Update(root, types.EmptyRootHash, 0, trienode.NewWithNodeSet(nodes), nil)
+			tdb.Commit(root, false)
+		}
+		resolver = func(owner common.Hash, path []byte, hash common.Hash) []byte {
+			return rawdb.ReadTrieNode(mdb, owner, path, hash, tdb.Scheme()) // Read the TrieNode based on scheme
 		}
-		snapTrieDb.Commit(root, false)
 	}
+
 	tr := result.tr
 	if tr == nil {
 		tr, err = trie.New(trieID, dl.triedb)
@@ -469,7 +474,7 @@ func (dl *diskLayer) generateRange(trieID *trie.ID, prefix []byte, kind string,
 		start    = time.Now()
 		internal time.Duration
 	)
-	nodeIt.AddResolver(snapNodeCache)
+	nodeIt.AddResolver(resolver)
 	for iter.Next() {
 		if last != nil && bytes.Compare(iter.Key, last) > 0 {
 			trieMore = true

trie/encoding_test.go

@@ -78,17 +78,17 @@ func TestHexKeybytes(t *testing.T) {
 }
 
 func TestHexToCompactInPlace(t *testing.T) {
-	for i, keyS := range []string{
+	for i, key := range []string{
 		"00",
 		"060a040c0f000a090b040803010801010900080d090a0a0d0903000b10",
 		"10",
 	} {
-		hexBytes, _ := hex.DecodeString(keyS)
+		hexBytes, _ := hex.DecodeString(key)
 		exp := hexToCompact(hexBytes)
 		sz := hexToCompactInPlace(hexBytes)
 		got := hexBytes[:sz]
 		if !bytes.Equal(exp, got) {
-			t.Fatalf("test %d: encoding err\ninp %v\ngot %x\nexp %x\n", i, keyS, got, exp)
+			t.Fatalf("test %d: encoding err\ninp %v\ngot %x\nexp %x\n", i, key, got, exp)
 		}
 	}
 }

trie/iterator.go

@@ -22,10 +22,16 @@ import (
 	"errors"
 
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
+// NodeResolver is used for looking up trie nodes before reaching into the real
+// persistent layer. This is not mandatory, 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. (It should be supported for multiple schemes in iterator)
+type NodeResolver func(owner common.Hash, path []byte, hash common.Hash) []byte
+
 // Iterator is a key-value trie iterator that traverses a Trie.
 type Iterator struct {
 	nodeIt NodeIterator
@@ -108,8 +114,8 @@ type NodeIterator interface {
 	// 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.
+	// AddResolver sets a node resolver 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
@@ -119,7 +125,7 @@ type NodeIterator interface {
 	// 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)
+	AddResolver(NodeResolver)
 }
 
 // nodeIteratorState represents the iteration state at one particular node of the
@@ -138,7 +144,7 @@ type nodeIterator struct {
 	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
+	resolver NodeResolver // optional node resolver for avoiding disk hits
 }
 
 // errIteratorEnd is stored in nodeIterator.err when iteration is done.
@@ -166,7 +172,7 @@ func newNodeIterator(trie *Trie, start []byte) NodeIterator {
 	return it
 }
 
-func (it *nodeIterator) AddResolver(resolver ethdb.KeyValueStore) {
+func (it *nodeIterator) AddResolver(resolver NodeResolver) {
 	it.resolver = resolver
 }
 
@@ -371,12 +377,15 @@ func (it *nodeIterator) peekSeek(seekKey []byte) (*nodeIteratorState, *int, []by
 
 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 {
+		// Support hash/path from memory.
+		if blob := it.resolver(it.trie.owner, path, common.BytesToHash(hash)); len(blob) > 0 {
 			if resolved, err := decodeNode(hash, blob); err == nil {
 				return resolved, nil
 			}
 		}
 	}
+
+	// Retrieve the specified node from the underlying node reader.
 	blob, err := it.trie.reader.node(path, common.BytesToHash(hash))
 	if err != nil {
 		return nil, err
@@ -389,7 +398,7 @@ func (it *nodeIterator) resolveHash(hash hashNode, path []byte) (node, error) {
 
 func (it *nodeIterator) resolveBlob(hash hashNode, path []byte) ([]byte, error) {
 	if it.resolver != nil {
-		if blob, err := it.resolver.Get(hash); err == nil && len(blob) > 0 {
+		if blob := it.resolver(it.trie.owner, path, common.BytesToHash(hash)); len(blob) > 0 {
 			return blob, nil
 		}
 	}
@@ -593,7 +602,7 @@ func (it *differenceIterator) NodeBlob() []byte {
 	return it.b.NodeBlob()
 }
 
-func (it *differenceIterator) AddResolver(resolver ethdb.KeyValueStore) {
+func (it *differenceIterator) AddResolver(resolver NodeResolver) {
 	panic("not implemented")
 }
 
@@ -708,7 +717,7 @@ func (it *unionIterator) NodeBlob() []byte {
 	return (*it.items)[0].NodeBlob()
 }
 
-func (it *unionIterator) AddResolver(resolver ethdb.KeyValueStore) {
+func (it *unionIterator) AddResolver(resolver NodeResolver) {
 	panic("not implemented")
 }