trie: don't copy when we are the exclusive owner of the trie

Currently, to make the trie copy perform well, we only copy the trie's node on
write. However, there is no tracking on whether the trie is copied or not, so we
always perform copy-on-write even though we are the only owner of the trie. This
commit adds a shared bool to trie, this is set when trie copy is performed. When
the trie is not shared, it is safe to do all the modification in-place without
the need of creating a new copy.

> go test -test.v -test.run=^$ -test.bench=BenchmarkUpdate

           │    old.txt    │               new.txt                │
           │    sec/op     │    sec/op     vs base                │
UpdateBE-8   1247.0n ± 36%   378.4n ± 34%  -69.66% (p=0.000 n=10)
UpdateLE-8   1675.5n ±  1%   430.8n ± 31%  -74.29% (p=0.000 n=10)
geomean       1.445µ         403.7n        -72.07%

           │   old.txt    │              new.txt               │
           │     B/op     │    B/op     vs base                │
UpdateBE-8    1796.0 ± 0%   226.0 ± 0%  -87.42% (p=0.000 n=10)
UpdateLE-8    1848.0 ± 0%   228.5 ± 1%  -87.64% (p=0.000 n=10)
geomean      1.779Ki        227.2       -87.53%

           │  old.txt   │              new.txt               │
           │ allocs/op  │ allocs/op   vs base                │
UpdateBE-8   9.000 ± 0%   4.000 ± 0%  -55.56% (p=0.000 n=10)
UpdateLE-8   9.000 ± 0%   4.000 ± 0%  -55.56% (p=0.000 n=10)
geomean      9.000        4.000       -55.56%

trie/secure_trie.go

@@ -64,6 +64,12 @@ func NewSecure(root common.Hash, db *Database) (*SecureTrie, error) {
 	return &SecureTrie{trie: *trie}, nil
 }
 
+// UnshareTrie marks this trie as unshared, the trie exclusively
+// owns all of its nodes.
+func (t *SecureTrie) UnshareTrie() {
+	t.trie.shared = false
+}
+
 // Get returns the value for key stored in the trie.
 // The value bytes must not be modified by the caller.
 func (t *SecureTrie) Get(key []byte) []byte {
@@ -185,6 +191,8 @@ func (t *SecureTrie) Hash() common.Hash {
 
 // Copy returns a copy of SecureTrie.
 func (t *SecureTrie) Copy() *SecureTrie {
+	// Mark both trie as shared
+	t.trie.shared = true
 	cpy := *t
 	return &cpy
 }

trie/trie.go

@@ -66,6 +66,11 @@ type Trie struct {
 	// hashing operation. This number will not directly map to the number of
 	// actually unhashed nodes
 	unhashed int
+
+	// The trie is already copied so we might share the trie node with
+	// other tries. If shared is false, we are the exclusive owner of
+	// the trie so we can modify it in place without copying.
+	shared bool
 }
 
 // newFlag returns the cache flag value for a newly created node.
@@ -136,14 +141,22 @@ func (t *Trie) tryGet(origNode node, key []byte, pos int) (value []byte, newnode
 		}
 		value, newnode, didResolve, err = t.tryGet(n.Val, key, pos+len(n.Key))
 		if err == nil && didResolve {
-			n = n.copy()
+			// This node might be shared with other tries, we must copy before
+			// modifying
+			if t.shared {
+				n = n.copy()
+			}
 			n.Val = newnode
 		}
 		return value, n, didResolve, err
 	case *fullNode:
 		value, newnode, didResolve, err = t.tryGet(n.Children[key[pos]], key, pos+1)
 		if err == nil && didResolve {
-			n = n.copy()
+			// This node might be shared with other tries, we must copy before
+			// modifying
+			if t.shared {
+				n = n.copy()
+			}
 			n.Children[key[pos]] = newnode
 		}
 		return value, n, didResolve, err
@@ -210,15 +223,23 @@ func (t *Trie) tryGetNode(origNode node, path []byte, pos int) (item []byte, new
 		}
 		item, newnode, resolved, err = t.tryGetNode(n.Val, path, pos+len(n.Key))
 		if err == nil && resolved > 0 {
-			n = n.copy()
+			// This node might be shared with other tries, we must copy before
+			// modifying
+			if t.shared {
+				n = n.copy()
+			}
 			n.Val = newnode
 		}
 		return item, n, resolved, err
 
 	case *fullNode:
 		item, newnode, resolved, err = t.tryGetNode(n.Children[path[pos]], path, pos+1)
 		if err == nil && resolved > 0 {
-			n = n.copy()
+			// This node might be shared with other tries, we must copy before
+			// modifying
+			if t.shared {
+				n = n.copy()
+			}
 			n.Children[path[pos]] = newnode
 		}
 		return item, n, resolved, err
@@ -300,7 +321,14 @@ func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error
 			if !dirty || err != nil {
 				return false, n, err
 			}
-			return true, &shortNode{n.Key, nn, t.newFlag()}, nil
+			if t.shared {
+				// This node might be shared with other tries, we must create a new node
+				return true, &shortNode{n.Key, nn, t.newFlag()}, nil
+			} else {
+				n.Val = nn
+				n.flags = t.newFlag()
+				return true, n, nil
+			}
 		}
 		// Otherwise branch out at the index where they differ.
 		branch := &fullNode{flags: t.newFlag()}
@@ -317,15 +345,28 @@ func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error
 		if matchlen == 0 {
 			return true, branch, nil
 		}
-		// Otherwise, replace it with a short node leading up to the branch.
-		return true, &shortNode{key[:matchlen], branch, t.newFlag()}, nil
+		// Otherwise, replace it with a new short node leading up to the branch if
+		// trie is shared; otherwise, modify the current short node to point to the
+		// new branch.
+		if t.shared {
+			return true, &shortNode{key[:matchlen], branch, t.newFlag()}, nil
+		} else {
+			n.Key = key[:matchlen]
+			n.Val = branch
+			n.flags = t.newFlag()
+			return true, n, nil
+		}
 
 	case *fullNode:
 		dirty, nn, err := t.insert(n.Children[key[0]], append(prefix, key[0]), key[1:], value)
 		if !dirty || err != nil {
 			return false, n, err
 		}
-		n = n.copy()
+		// This node might be shared with other tries, we must copy before
+		// modifying
+		if t.shared {
+			n = n.copy()
+		}
 		n.flags = t.newFlag()
 		n.Children[key[0]] = nn
 		return true, n, nil
@@ -401,17 +442,36 @@ func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) {
 			// always creates a new slice) instead of append to
 			// avoid modifying n.Key since it might be shared with
 			// other nodes.
-			return true, &shortNode{concat(n.Key, child.Key...), child.Val, t.newFlag()}, nil
+			if t.shared {
+				// This node might be shared with other tries, we must create a new node
+				return true, &shortNode{concat(n.Key, child.Key...), child.Val, t.newFlag()}, nil
+			} else {
+				n.Key = concat(n.Key, child.Key...)
+				n.Val = child.Val
+				n.flags = t.newFlag()
+				return true, n, nil
+			}
 		default:
-			return true, &shortNode{n.Key, child, t.newFlag()}, nil
+			if t.shared {
+				// This node might be shared with other tries, we must create a new node
+				return true, &shortNode{n.Key, child, t.newFlag()}, nil
+			} else {
+				n.Val = child
+				n.flags = t.newFlag()
+				return true, n, nil
+			}
 		}
 
 	case *fullNode:
 		dirty, nn, err := t.delete(n.Children[key[0]], append(prefix, key[0]), key[1:])
 		if !dirty || err != nil {
 			return false, n, err
 		}
-		n = n.copy()
+		// This node might be shared with other tries, we must copy before
+		// modifying
+		if t.shared {
+			n = n.copy()
+		}
 		n.flags = t.newFlag()
 		n.Children[key[0]] = nn