trie: bump cachegen for read access, avoiding copies

This change ensures that the cache generation of trie branches is bumped
for reads. To avoid allocating new nodes for each read, the accessed
node is modified directly if it was created in the current cache
generation.

trie/hasher.go

@@ -85,11 +85,13 @@ func (h *hasher) hash(n node, db DatabaseWriter, force bool) (node, node, error)
 		cn.flags.hash = cachedHash
 		if db != nil {
 			cn.flags.dirty = false
+			cn.flags.new = false
 		}
 	case *fullNode:
 		cn.flags.hash = cachedHash
 		if db != nil {
 			cn.flags.dirty = false
+			cn.flags.new = false
 		}
 	}
 	return hashed, cached, nil

trie/node.go

@@ -60,6 +60,7 @@ type nodeFlag struct {
 	hash  hashNode // cached hash of the node (may be nil)
 	gen   uint16   // cache generation counter
 	dirty bool     // whether the node has changes that must be written to the database
+	new   bool     // whether the node was allocated in the current generation
 }
 
 // canUnload tells whether a node can be unloaded.
@@ -138,7 +139,7 @@ func decodeShort(hash, buf, elems []byte, cachegen uint16) (node, error) {
 	if err != nil {
 		return nil, err
 	}
-	flag := nodeFlag{hash: hash, gen: cachegen}
+	flag := nodeFlag{hash: hash, gen: cachegen, new: true}
 	key := compactDecode(kbuf)
 	if key[len(key)-1] == 16 {
 		// value node
@@ -156,7 +157,7 @@ func decodeShort(hash, buf, elems []byte, cachegen uint16) (node, error) {
 }
 
 func decodeFull(hash, buf, elems []byte, cachegen uint16) (*fullNode, error) {
-	n := &fullNode{flags: nodeFlag{hash: hash, gen: cachegen}}
+	n := &fullNode{flags: nodeFlag{hash: hash, gen: cachegen, new: true}}
 	for i := 0; i < 16; i++ {
 		cld, rest, err := decodeRef(elems, cachegen)
 		if err != nil {

trie/trie.go

@@ -98,7 +98,7 @@ func (t *Trie) SetCacheLimit(l uint16) {
 
 // newFlag returns the cache flag value for a newly created node.
 func (t *Trie) newFlag() nodeFlag {
-	return nodeFlag{dirty: true, gen: t.cachegen}
+	return nodeFlag{new: true, dirty: true, gen: t.cachegen}
 }
 
 // New creates a trie with an existing root node from db.
@@ -142,46 +142,52 @@ func (t *Trie) Get(key []byte) []byte {
 // If a node was not found in the database, a MissingNodeError is returned.
 func (t *Trie) TryGet(key []byte) ([]byte, error) {
 	key = compactHexDecode(key)
-	value, newroot, didResolve, err := t.tryGet(t.root, key, 0)
-	if err == nil && didResolve {
+	value, newroot, err := t.tryGet(t.root, key, 0)
+	if err == nil {
 		t.root = newroot
 	}
 	return value, err
 }
 
-func (t *Trie) tryGet(origNode node, key []byte, pos int) (value []byte, newnode node, didResolve bool, err error) {
+func (t *Trie) tryGet(origNode node, key []byte, pos int) (value []byte, newnode node, err error) {
 	switch n := (origNode).(type) {
 	case nil:
-		return nil, nil, false, nil
+		return nil, nil, nil
 	case valueNode:
-		return n, n, false, nil
+		return n, n, nil
 	case *shortNode:
 		if len(key)-pos < len(n.Key) || !bytes.Equal(n.Key, key[pos:pos+len(n.Key)]) {
 			// key not found in trie
-			return nil, n, false, nil
+			return nil, n, nil
 		}
-		value, newnode, didResolve, err = t.tryGet(n.Val, key, pos+len(n.Key))
-		if err == nil && didResolve {
+		value, newnode, err = t.tryGet(n.Val, key, pos+len(n.Key))
+		if err == nil {
+			if !n.flags.new {
 				n = n.copy()
-			n.Val = newnode
+				n.flags.new = true
+			}
 			n.flags.gen = t.cachegen
+			n.Val = newnode
 		}
-		return value, n, didResolve, err
+		return value, n, err
 	case *fullNode:
-		value, newnode, didResolve, err = t.tryGet(n.Children[key[pos]], key, pos+1)
-		if err == nil && didResolve {
+		value, newnode, err = t.tryGet(n.Children[key[pos]], key, pos+1)
+		if err == nil {
+			if !n.flags.new {
 				n = n.copy()
+				n.flags.new = true
+			}
 			n.flags.gen = t.cachegen
 			n.Children[key[pos]] = newnode
 		}
-		return value, n, didResolve, err
+		return value, n, err
 	case hashNode:
 		child, err := t.resolveHash(n, key[:pos], key[pos:])
 		if err != nil {
-			return nil, n, true, err
+			return nil, n, err
 		}
-		value, newnode, _, err := t.tryGet(child, key, pos)
-		return value, newnode, true, err
+		value, newnode, err := t.tryGet(child, key, pos)
+		return value, newnode, err
 	default:
 		panic(fmt.Sprintf("%T: invalid node: %v", origNode, origNode))
 	}

trie/trie_test.go

@@ -329,8 +329,7 @@ func TestCacheUnload(t *testing.T) {
 	root, _ := trie.Commit()
 
 	// Commit the trie repeatedly and access key1.
-	// The branch containing it is loaded from DB exactly two times:
-	// in the 0th and 6th iteration.
+	// The branch containing it is loaded from DB exactly once.
 	db := &countingDB{Database: trie.db, gets: make(map[string]int)}
 	trie, _ = New(root, db)
 	trie.SetCacheLimit(5)
@@ -339,10 +338,10 @@ func TestCacheUnload(t *testing.T) {
 		trie.Commit()
 	}
 
-	// Check that it got loaded two times.
+	// Check that it got loaded once.
 	for dbkey, count := range db.gets {
-		if count != 2 {
-			t.Errorf("db key %x loaded %d times, want %d times", []byte(dbkey), count, 2)
+		if count != 1 {
+			t.Errorf("db key %x loaded %d times, want %d times", []byte(dbkey), count, 1)
 		}
 	}
 }