Trie gen PR-ception

core/state/snapshot/hextrie_generator.go

@@ -31,6 +31,14 @@ type leaf struct {
 // as the rest of geth, with no enhancements or optimizations
 type trieGenerator struct{}
 
+func (gen *trieGenerator) Generate3(in chan (leaf), out chan (common.Hash)) {
+	t := trie.NewHashTrie()
+	for leaf := range in {
+		t.TryUpdate(leaf.key[:], leaf.value)
+	}
+	out <- t.Hash()
+}
+
 //BenchmarkTrieGeneration/4K-6         	      94	  12598506 ns/op	 6162370 B/op	   57921 allocs/op
 //BenchmarkTrieGeneration/10K-6        	      37	  33790908 ns/op	17278751 B/op	  151002 allocs/op
 func (gen *trieGenerator) Generate2(in chan (leaf), out chan (common.Hash)) {
@@ -41,7 +49,6 @@ func (gen *trieGenerator) Generate2(in chan (leaf), out chan (common.Hash)) {
 	out <- t.Hash()
 }
 
-
 //BenchmarkTrieGeneration/4K-6         	     115	  12755614 ns/op	 2303051 B/op	   42678 allocs/op
 //BenchmarkTrieGeneration/10K-6        	      46	  25374595 ns/op	 5754446 B/op	  106676 allocs/op
 func (gen *trieGenerator) Generate(in chan (leaf), out chan (common.Hash)) {

core/state/snapshot/trie_generator_test.go

@@ -28,16 +28,16 @@ import (
 
 func generateTrie(it AccountIterator, generator *trieGenerator) common.Hash {
 	var (
-		in  = make(chan leaf)        // chan to pass leafs
+		in  = make(chan leaf)        // chan to pass leaves
 		out = make(chan common.Hash) // chan to collect result
 		wg  sync.WaitGroup
 	)
 	wg.Add(1)
 	go func() {
-		generator.Generate2(in, out)
+		generator.Generate3(in, out)
 		wg.Done()
 	}()
-	// Feed leafs
+	// Feed leaves
 	for it.Next() {
 		in <- leaf{it.Hash(), it.Account()}
 	}

trie/appendtrie.go

@@ -1,4 +1,4 @@
-// Copyright 2014 The go-ethereum Authors
+// Copyright 2020 The go-ethereum Authors
 // This file is part of the go-ethereum library.
 //
 // The go-ethereum library is free software: you can redistribute it and/or modify
@@ -28,7 +28,7 @@ type AppendOnlyTrie struct {
 }
 
 func NewAppendOnlyTrie() *AppendOnlyTrie {
-	return &AppendOnlyTrie{root:nil}
+	return &AppendOnlyTrie{root: nil}
 }
 
 func (t *AppendOnlyTrie) TryUpdate(key, value []byte) error {
@@ -55,9 +55,9 @@ func (t *AppendOnlyTrie) insert(n node, prefix, key []byte, value node) node {
 		}
 		// Otherwise branch out at the index where they differ.
 		branch := &fullNode{flags: nodeFlag{dirty: true}}
-		branch.Children[n.Key[matchlen]]= t.insert(nil, append(prefix, n.Key[:matchlen+1]...), n.Key[matchlen+1:], n.Val)
+		branch.Children[n.Key[matchlen]] = t.insert(nil, append(prefix, n.Key[:matchlen+1]...), n.Key[matchlen+1:], n.Val)
 		// TODO: We can now shoot off n.Val for hashing
-		branch.Children[key[matchlen]]= t.insert(nil, append(prefix, key[:matchlen+1]...), key[matchlen+1:], value)
+		branch.Children[key[matchlen]] = t.insert(nil, append(prefix, key[:matchlen+1]...), key[matchlen+1:], value)
 
 		// Replace this shortNode with the branch if it occurs at index 0.
 		if matchlen == 0 {
@@ -96,4 +96,4 @@ func (t *AppendOnlyTrie) Hash() common.Hash {
 	hashed, cached := h.hash(t.root, true)
 	t.root = cached
 	return common.BytesToHash(hashed.(hashNode))
-}
+}

trie/hashtrie.go

@@ -0,0 +1,117 @@
+// Copyright 2020 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package trie
+
+import (
+	"fmt"
+	"github.com/ethereum/go-ethereum/common"
+)
+
+// HashTrie is a Merkle Patricia Trie, which can only be used for
+// constructing a trie from a sequence of sorted leafs, in descending order
+type HashTrie struct {
+	root    node
+	rootKey []byte
+	build   []node
+}
+
+func NewHashTrie() *HashTrie {
+	return &HashTrie{root: nil, rootKey: nil, build: nil}
+}
+
+func (t *HashTrie) TryUpdate(key, value []byte) error {
+	k := keybytesToHex(key)
+	if len(value) == 0 {
+		panic("deletion not supported")
+	}
+	t.root = t.insert(t.root, nil, k, valueNode(value))
+	return nil
+}
+
+func (t *HashTrie) insert(n node, prefix, key []byte, value node) node {
+	if len(key) == 0 {
+		return value
+	}
+	switch n := n.(type) {
+	case *shortNode:
+		matchlen := prefixLen(key, n.Key)
+		// If the whole key matches, it already exists
+		if matchlen == len(n.Key) {
+			n.Val = t.insert(n.Val, append(prefix, key[:matchlen]...), key[matchlen:], value)
+			n.flags = nodeFlag{dirty: true}
+			return n
+		}
+		// Otherwise branch out at the index where they differ.
+		branch := &fullNode{flags: nodeFlag{dirty: true}}
+		branch.Children[n.Key[matchlen]] = t.insert(nil, append(prefix, n.Key[:matchlen+1]...), n.Key[matchlen+1:], n.Val)
+		// Hashing the sub-node, nothing will be added to this sub-branch
+		hashed, _ := newHasher(false).hash(t.insert(nil, append(prefix, key[:matchlen+1]...), key[matchlen+1:], value), true)
+		branch.Children[key[matchlen]] = hashed.(hashNode)
+
+		// Replace this shortNode with the branch if it occurs at index 0.
+		if matchlen == 0 {
+			return branch
+		}
+		// Otherwise, replace it with a short node leading up to the branch.
+		n.Key = key[:matchlen]
+		n.Val = branch
+		n.flags = nodeFlag{dirty: true}
+		return n
+
+	case *fullNode:
+		n.flags = nodeFlag{dirty: true}
+		// If any previous child wasn't already hashed, do it now since
+		// the keys arrive in order, so if a branch is here then whatever
+		// came before can safely be hashed.
+		for i := int(key[0]) - 1; i > 0; i -= 1 {
+			switch n.Children[i].(type) {
+			case *shortNode, *fullNode, *valueNode:
+				hashed, _ := newHasher(false).hash(n.Children[i], true)
+				n.Children[i] = hashed
+			// hash encountred, the rest has already been hashed
+			case hashNode:
+				break
+			default:
+				panic("invalid node")
+			}
+		}
+		n.Children[key[0]] = t.insert(n.Children[key[0]], append(prefix, key[0]), key[1:], value)
+		return n
+
+	case nil:
+		return &shortNode{key, value, nodeFlag{dirty: true}}
+
+	case hashNode:
+		// We've hit a part of the trie that isn't loaded yet -- this means
+		// someone inserted
+		panic("hash resolution not supported")
+
+	default:
+		panic(fmt.Sprintf("%T: invalid node: %v", n, n))
+	}
+}
+
+func (t *HashTrie) Hash() common.Hash {
+	if t.root == nil {
+		return emptyRoot
+	}
+	h := newHasher(false)
+	defer returnHasherToPool(h)
+	hashed, cached := h.hash(t.root, true)
+	t.root = cached
+	return common.BytesToHash(hashed.(hashNode))
+}