feat: refactor the node key as version + local nonce(seq id)

Makefile

@@ -53,18 +53,18 @@ lint-fix:
 # bench is the basic tests that shouldn't crash an aws instance
 bench:
 	cd benchmarks && \
-		go test $(LDFLAGS) -tags cleveldb,rocksdb,pebbledb -run=NOTEST -bench=Small . && \
-		go test $(LDFLAGS) -tags cleveldb,rocksdb,pebbledb -run=NOTEST -bench=Medium . && \
+		go test $(LDFLAGS) -tags pebbledb -run=NOTEST -bench=Small . && \
+		go test $(LDFLAGS) -tags pebbledb -run=NOTEST -bench=Medium . && \
 		go test $(LDFLAGS) -run=NOTEST -bench=RandomBytes .
 .PHONY: bench
 
 # fullbench is extra tests needing lots of memory and to run locally
 fullbench:
 	cd benchmarks && \
 		go test $(LDFLAGS) -run=NOTEST -bench=RandomBytes . && \
-		go test $(LDFLAGS) -tags cleveldb,rocksdb,pebbledb -run=NOTEST -bench=Small . && \
-		go test $(LDFLAGS) -tags cleveldb,rocksdb,pebbledb -run=NOTEST -bench=Medium . && \
-		go test $(LDFLAGS) -tags cleveldb,rocksdb,pebbledb -run=NOTEST -timeout=30m -bench=Large . && \
+		go test $(LDFLAGS) -tags pebbledb -run=NOTEST -bench=Small . && \
+		go test $(LDFLAGS) -tags pebbledb -run=NOTEST -bench=Medium . && \
+		go test $(LDFLAGS) -tags pebbledb -run=NOTEST -timeout=30m -bench=Large . && \
 		go test $(LDFLAGS) -run=NOTEST -bench=Mem . && \
 		go test $(LDFLAGS) -run=NOTEST -timeout=60m -bench=LevelDB .
 .PHONY: fullbench

basic_test.go

@@ -2,7 +2,6 @@
 package iavl
 
 import (
-	"bytes"
 	"encoding/hex"
 	mrand "math/rand"
 	"sort"
@@ -171,52 +170,28 @@ func TestBasic(t *testing.T) {
 }
 
 func TestUnit(t *testing.T) {
-	expectHash := func(tree *ImmutableTree, hashCount int64) {
-		// ensure number of new hash calculations is as expected.
-		hash, count, err := tree.root.hashWithCount()
-		require.NoError(t, err)
-		if count != hashCount {
-			t.Fatalf("Expected %v new hashes, got %v", hashCount, count)
-		}
-		// nuke hashes and reconstruct hash, ensure it's the same.
-		tree.root.traverse(tree, true, func(node *Node) bool {
-			node.hash = nil
-			return false
-		})
-		// ensure that the new hash after nuking is the same as the old.
-		newHash, _, err := tree.root.hashWithCount()
-		require.NoError(t, err)
-		if !bytes.Equal(hash, newHash) {
-			t.Fatalf("Expected hash %v but got %v after nuking", hash, newHash)
-		}
-	}
-
-	expectSet := func(tree *MutableTree, i int, repr string, hashCount int64) {
-		origNode := tree.root
+	expectSet := func(tree *MutableTree, i int, repr string) {
+		tree.SaveVersion()
 		updated, err := tree.Set(i2b(i), []byte{})
 		require.NoError(t, err)
 		// ensure node was added & structure is as expected.
-		if updated || P(tree.root) != repr {
+		if updated || P(tree.root, tree.ImmutableTree) != repr {
 			t.Fatalf("Adding %v to %v:\nExpected         %v\nUnexpectedly got %v updated:%v",
-				i, P(origNode), repr, P(tree.root), updated)
+				i, P(tree.lastSaved.root, tree.lastSaved), repr, P(tree.root, tree.ImmutableTree), updated)
 		}
-		// ensure hash calculation requirements
-		expectHash(tree.ImmutableTree, hashCount)
-		tree.root = origNode
+		tree.ImmutableTree = tree.lastSaved.clone()
 	}
 
-	expectRemove := func(tree *MutableTree, i int, repr string, hashCount int64) {
-		origNode := tree.root
+	expectRemove := func(tree *MutableTree, i int, repr string) {
+		tree.SaveVersion()
 		value, removed, err := tree.Remove(i2b(i))
 		require.NoError(t, err)
 		// ensure node was added & structure is as expected.
-		if len(value) != 0 || !removed || P(tree.root) != repr {
+		if len(value) != 0 || !removed || P(tree.root, tree.ImmutableTree) != repr {
 			t.Fatalf("Removing %v from %v:\nExpected         %v\nUnexpectedly got %v value:%v removed:%v",
-				i, P(origNode), repr, P(tree.root), value, removed)
+				i, P(tree.lastSaved.root, tree.lastSaved), repr, P(tree.root, tree.ImmutableTree), value, removed)
 		}
-		// ensure hash calculation requirements
-		expectHash(tree.ImmutableTree, hashCount)
-		tree.root = origNode
+		tree.ImmutableTree = tree.lastSaved.clone()
 	}
 
 	// Test Set cases:
@@ -225,40 +200,40 @@ func TestUnit(t *testing.T) {
 	t1, err := T(N(4, 20))
 
 	require.NoError(t, err)
-	expectSet(t1, 8, "((4 8) 20)", 3)
-	expectSet(t1, 25, "(4 (20 25))", 3)
+	expectSet(t1, 8, "((4 8) 20)")
+	expectSet(t1, 25, "(4 (20 25))")
 
 	t2, err := T(N(4, N(20, 25)))
 
 	require.NoError(t, err)
-	expectSet(t2, 8, "((4 8) (20 25))", 3)
-	expectSet(t2, 30, "((4 20) (25 30))", 4)
+	expectSet(t2, 8, "((4 8) (20 25))")
+	expectSet(t2, 30, "((4 20) (25 30))")
 
 	t3, err := T(N(N(1, 2), 6))
 
 	require.NoError(t, err)
-	expectSet(t3, 4, "((1 2) (4 6))", 4)
-	expectSet(t3, 8, "((1 2) (6 8))", 3)
+	expectSet(t3, 4, "((1 2) (4 6))")
+	expectSet(t3, 8, "((1 2) (6 8))")
 
 	t4, err := T(N(N(1, 2), N(N(5, 6), N(7, 9))))
 
 	require.NoError(t, err)
-	expectSet(t4, 8, "(((1 2) (5 6)) ((7 8) 9))", 5)
-	expectSet(t4, 10, "(((1 2) (5 6)) (7 (9 10)))", 5)
+	expectSet(t4, 8, "(((1 2) (5 6)) ((7 8) 9))")
+	expectSet(t4, 10, "(((1 2) (5 6)) (7 (9 10)))")
 
 	// Test Remove cases:
 
 	t10, err := T(N(N(1, 2), 3))
 
 	require.NoError(t, err)
-	expectRemove(t10, 2, "(1 3)", 1)
-	expectRemove(t10, 3, "(1 2)", 0)
+	expectRemove(t10, 2, "(1 3)")
+	expectRemove(t10, 3, "(1 2)")
 
 	t11, err := T(N(N(N(1, 2), 3), N(4, 5)))
 
 	require.NoError(t, err)
-	expectRemove(t11, 4, "((1 2) (3 5))", 2)
-	expectRemove(t11, 3, "((1 2) (4 5))", 1)
+	expectRemove(t11, 4, "((1 2) (3 5))")
+	expectRemove(t11, 3, "((1 2) (4 5))")
 }
 
 func TestRemove(t *testing.T) {

diff.go

@@ -25,7 +25,7 @@ type KVPairReceiver func(pair *KVPair) error
 //
 // The algorithm don't run in constant memory strictly, but it tried the best the only
 // keep minimal intermediate states in memory.
-func (ndb *nodeDB) extractStateChanges(prevVersion int64, prevRoot []byte, root []byte, receiver KVPairReceiver) error {
+func (ndb *nodeDB) extractStateChanges(prevVersion int64, prevRoot *NodeKey, root *NodeKey, receiver KVPairReceiver) error {
 	curIter, err := NewNodeIterator(root, ndb)
 	if err != nil {
 		return err
@@ -70,7 +70,7 @@ func (ndb *nodeDB) extractStateChanges(prevVersion int64, prevRoot []byte, root
 		sharedNode = nil
 		for curIter.Valid() {
 			node := curIter.GetNode()
-			shared := node.version <= prevVersion
+			shared := node.nodeKey.version <= prevVersion
 			curIter.Next(shared)
 			if shared {
 				sharedNode = node

export.go

@@ -65,7 +65,7 @@ func (e *Exporter) export(ctx context.Context) {
 		exportNode := &ExportNode{
 			Key:     node.key,
 			Value:   node.value,
-			Version: node.version,
+			Version: node.nodeKey.version,
 			Height:  node.subtreeHeight,
 		}
 

immutable_tree.go

@@ -149,8 +149,7 @@ func (t *ImmutableTree) Has(key []byte) (bool, error) {
 
 // Hash returns the root hash.
 func (t *ImmutableTree) Hash() ([]byte, error) {
-	hash, _, err := t.root.hashWithCount()
-	return hash, err
+	return t.root.hashWithCount(t.version + 1)
 }
 
 // Export returns an iterator that exports tree nodes as ExportNodes. These nodes can be
@@ -283,7 +282,7 @@ func (t *ImmutableTree) IterateRangeInclusive(start, end []byte, ascending bool,
 	}
 	return t.root.traverseInRange(t, start, end, ascending, true, false, func(node *Node) bool {
 		if node.subtreeHeight == 0 {
-			return fn(node.key, node.value, node.version)
+			return fn(node.key, node.value, node.nodeKey.version)
 		}
 		return false
 	})
@@ -324,12 +323,7 @@ func (t *ImmutableTree) clone() *ImmutableTree {
 //
 
 func (t *ImmutableTree) nodeSize() int {
-	size := 0
-	t.root.traverse(t, true, func(n *Node) bool {
-		size++
-		return false
-	})
-	return size
+	return int(t.root.size*2 - 1)
 }
 
 // TraverseStateChanges iterate the range of versions, compare each version to it's predecessor to extract the state changes of it.

import.go

@@ -27,6 +27,7 @@ type Importer struct {
 	batch     db.Batch
 	batchSize uint32
 	stack     []*Node
+	nonces    []int32
 }
 
 // newImporter creates a new Importer for an empty MutableTree.
@@ -49,9 +50,47 @@ func newImporter(tree *MutableTree, version int64) (*Importer, error) {
 		version: version,
 		batch:   tree.ndb.db.NewBatch(),
 		stack:   make([]*Node, 0, 8),
+		nonces:  make([]int32, version+1),
 	}, nil
 }
 
+// writeNode writes the node content to the storage.
+func (i *Importer) writeNode(node *Node) error {
+	if _, err := node._hash(node.nodeKey.version); err != nil {
+		return err
+	}
+	if err := node.validate(); err != nil {
+		return err
+	}
+
+	buf := bufPool.Get().(*bytes.Buffer)
+	buf.Reset()
+	defer bufPool.Put(buf)
+
+	if err := node.writeBytes(buf); err != nil {
+		return err
+	}
+
+	bytesCopy := make([]byte, buf.Len())
+	copy(bytesCopy, buf.Bytes())
+
+	if err := i.batch.Set(i.tree.ndb.nodeKey(node.nodeKey), bytesCopy); err != nil {
+		return err
+	}
+
+	i.batchSize++
+	if i.batchSize >= maxBatchSize {
+		if err := i.batch.Write(); err != nil {
+			return err
+		}
+		i.batch.Close()
+		i.batch = i.tree.ndb.db.NewBatch()
+		i.batchSize = 0
+	}
+
+	return nil
+}
+
 // Close frees all resources. It is safe to call multiple times. Uncommitted nodes may already have
 // been flushed to the database, but will not be visible.
 func (i *Importer) Close() {
@@ -80,7 +119,6 @@ func (i *Importer) Add(exportNode *ExportNode) error {
 	node := &Node{
 		key:           exportNode.Key,
 		value:         exportNode.Value,
-		version:       exportNode.Version,
 		subtreeHeight: exportNode.Height,
 	}
 
@@ -92,72 +130,31 @@ func (i *Importer) Add(exportNode *ExportNode) error {
 	// We don't modify the stack until we've verified the built node, to avoid leaving the
 	// importer in an inconsistent state when we return an error.
 	stackSize := len(i.stack)
-	switch {
-	case stackSize >= 2 && i.stack[stackSize-1].subtreeHeight < node.subtreeHeight && i.stack[stackSize-2].subtreeHeight < node.subtreeHeight:
-		node.leftNode = i.stack[stackSize-2]
-		node.leftHash = node.leftNode.hash
-		node.rightNode = i.stack[stackSize-1]
-		node.rightHash = node.rightNode.hash
-	case stackSize >= 1 && i.stack[stackSize-1].subtreeHeight < node.subtreeHeight:
-		node.leftNode = i.stack[stackSize-1]
-		node.leftHash = node.leftNode.hash
-	}
-
 	if node.subtreeHeight == 0 {
 		node.size = 1
-	}
-	if node.leftNode != nil {
-		node.size += node.leftNode.size
-	}
-	if node.rightNode != nil {
-		node.size += node.rightNode.size
-	}
-
-	_, err := node._hash()
-	if err != nil {
-		return err
-	}
-
-	err = node.validate()
-	if err != nil {
-		return err
-	}
-
-	buf := bufPool.Get().(*bytes.Buffer)
-	buf.Reset()
-	defer bufPool.Put(buf)
-
-	if err = node.writeBytes(buf); err != nil {
-		return err
-	}
-
-	bytesCopy := make([]byte, buf.Len())
-	copy(bytesCopy, buf.Bytes())
-
-	if err = i.batch.Set(i.tree.ndb.nodeKey(node.hash), bytesCopy); err != nil {
-		return err
-	}
-
-	i.batchSize++
-	if i.batchSize >= maxBatchSize {
-		err = i.batch.Write()
-		if err != nil {
+	} else if stackSize >= 2 && i.stack[stackSize-1].subtreeHeight < node.subtreeHeight && i.stack[stackSize-2].subtreeHeight < node.subtreeHeight {
+		node.leftNode = i.stack[stackSize-2]
+		node.rightNode = i.stack[stackSize-1]
+		node.leftNodeKey = node.leftNode.nodeKey
+		node.rightNodeKey = node.rightNode.nodeKey
+		node.size = node.leftNode.size + node.rightNode.size
+		// Update the stack now.
+		if err := i.writeNode(i.stack[stackSize-2]); err != nil {
+			return err
+		}
+		if err := i.writeNode(i.stack[stackSize-1]); err != nil {
 			return err
 		}
-		i.batch.Close()
-		i.batch = i.tree.ndb.db.NewBatch()
-		i.batchSize = 0
-	}
-
-	// Update the stack now that we know there were no errors
-	switch {
-	case node.leftHash != nil && node.rightHash != nil:
 		i.stack = i.stack[:stackSize-2]
-	case node.leftHash != nil || node.rightHash != nil:
-		i.stack = i.stack[:stackSize-1]
 	}
-	// Only hash\height\size of the node will be used after it be pushed into the stack.
-	i.stack = append(i.stack, &Node{hash: node.hash, subtreeHeight: node.subtreeHeight, size: node.size})
+	i.nonces[exportNode.Version]++
+	node.nodeKey = &NodeKey{
+		version: exportNode.Version,
+		// Nonce is 1-indexed, but start at 2 since the root node having a nonce of 1.
+		nonce: i.nonces[exportNode.Version] + 1,
+	}
+
+	i.stack = append(i.stack, node)
 
 	return nil
 }
@@ -172,11 +169,12 @@ func (i *Importer) Commit() error {
 
 	switch len(i.stack) {
 	case 0:
-		if err := i.batch.Set(i.tree.ndb.rootKey(i.version), []byte{}); err != nil {
+		if err := i.batch.Set(i.tree.ndb.nodeKey(&NodeKey{version: i.version, nonce: 1}), []byte{}); err != nil {
 			return err
 		}
 	case 1:
-		if err := i.batch.Set(i.tree.ndb.rootKey(i.version), i.stack[0].hash); err != nil {
+		i.stack[0].nodeKey.nonce = 1
+		if err := i.writeNode(i.stack[0]); err != nil {
 			return err
 		}
 	default:

import_test.go

@@ -152,6 +152,9 @@ func TestImporter_Add(t *testing.T) {
 			if tc.valid {
 				require.NoError(t, err)
 			} else {
+				if err == nil {
+					err = importer.Commit()
+				}
 				require.Error(t, err)
 			}
 		})