feat: implement the lazy_set feature

cmd/dbgenerator/README.md

@@ -0,0 +1,16 @@
+# DBGenerator
+
+`dbgenerator` is a command line tool to generate a `iavl` tree based on the legacy format of the node key.
+This tool is used for testing the `lazy loading and set` feature of the `iavl` tree.
+
+## Usage
+
+It takes 4 arguments:
+    - dbtype: the type of database to use. 
+    - dbdir: the directory to store the database.
+    - `random` or `sequential`: The `sequential` option will generate the tree from `1` to `version` in order. The `random` option will remove half of the versions randomly from the `sequential` mode.
+    - version: the number of versions to generate.
+
+```shell
+go run main.go <dbtype> <dbdir> <random|sequential> <version>
+```

cmd/dbgenerator/go.mod

@@ -0,0 +1,31 @@
+module dbgenerator
+
+go 1.20
+
+require (
+	github.com/cometbft/cometbft-db v0.7.0
+	github.com/cosmos/iavl v0.20.0
+)
+
+require (
+	github.com/cespare/xxhash v1.1.0 // indirect
+	github.com/confio/ics23/go v0.9.0 // indirect
+	github.com/dgraph-io/badger/v2 v2.2007.4 // indirect
+	github.com/dgraph-io/ristretto v0.0.3-0.20200630154024-f66de99634de // indirect
+	github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 // indirect
+	github.com/dustin/go-humanize v1.0.0 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/golang/snappy v0.0.4 // indirect
+	github.com/google/btree v1.1.2 // indirect
+	github.com/jmhodges/levigo v1.0.0 // indirect
+	github.com/klauspost/compress v1.15.9 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/syndtr/goleveldb v1.0.1-0.20200815110645-5c35d600f0ca // indirect
+	github.com/tecbot/gorocksdb v0.0.0-20191217155057-f0fad39f321c // indirect
+	go.etcd.io/bbolt v1.3.6 // indirect
+	golang.org/x/crypto v0.4.0 // indirect
+	golang.org/x/net v0.4.0 // indirect
+	golang.org/x/sys v0.6.0 // indirect
+	google.golang.org/protobuf v1.28.1 // indirect
+)

cmd/dbgenerator/main.go

@@ -0,0 +1,116 @@
+package main
+
+import (
+	"fmt"
+	"math/rand"
+	"os"
+	"path/filepath"
+	"strconv"
+
+	dbm "github.com/cometbft/cometbft-db"
+	"github.com/cosmos/iavl"
+)
+
+const (
+	DefaultCacheSize = 1000
+)
+
+func main() {
+	args := os.Args[1:]
+	if len(args) < 4 {
+		fmt.Fprintln(os.Stderr, "Usage: dbgenerator <dbtype> <dbdir> <random|sequential> <version>")
+		os.Exit(1)
+	}
+
+	version, err := strconv.Atoi(args[3])
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Invalid version number: %s\n", err)
+		os.Exit(1)
+	}
+
+	if err = generateTree(args[0], args[1], args[2], version); err != nil {
+		fmt.Fprintf(os.Stderr, "Error generating tree: %s\n", err)
+	}
+}
+
+func openDB(dbType, dbDir string) (dbm.DB, error) {
+	dir, err := filepath.Abs(dbDir)
+	if err != nil {
+		return nil, err
+	}
+
+	db, err := dbm.NewDB("test", dbm.BackendType(dbType), dir)
+	if err != nil {
+		return nil, err
+	}
+	return db, nil
+}
+
+func generateTree(dbType, dbDir, mode string, version int) error {
+	db, err := openDB(dbType, dbDir)
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+
+	switch mode {
+	case "random":
+		return generateRandomTree(db, version)
+	case "sequential":
+		_, err = generateSequentialTree(db, version)
+		return err
+	default:
+		return fmt.Errorf("invalid mode: %s", mode)
+	}
+}
+
+func generateRandomTree(db dbm.DB, version int) error {
+	t, err := generateSequentialTree(db, version)
+	if err != nil {
+		return err
+	}
+
+	// delete the half of the versions
+	versions := make([]int64, version)
+	for i := 0; i < version; i++ {
+		versions[i] = int64(i + 1)
+	}
+
+	// make sure the latest version is not deleted
+	for i := 1; i <= version/2; i++ {
+		index := rand.Intn(version - i)
+		if err := t.DeleteVersion(versions[index]); err != nil {
+			return err
+		}
+		versions[index], versions[version-i-1] = versions[version-i-1], versions[index]
+	}
+
+	return nil
+}
+
+func generateSequentialTree(db dbm.DB, version int) (*iavl.MutableTree, error) {
+	t, err := iavl.NewMutableTreeWithOpts(db, DefaultCacheSize, nil, false)
+	if err != nil {
+		return nil, err
+	}
+
+	for i := 0; i < version; i++ {
+		leafCount := rand.Int31n(50)
+		for j := int32(0); j < leafCount; j++ {
+			t.Set(randBytes(32), randBytes(32))
+		}
+		if _, _, err = t.SaveVersion(); err != nil {
+			return nil, err
+		}
+	}
+
+	return t, err
+}
+
+func randBytes(length int) []byte {
+	key := make([]byte, length)
+	// math.rand.Read always returns err=nil
+	// we do not need cryptographic randomness for this test:
+	rand.Read(key)
+	return key
+}

diff.go

@@ -21,7 +21,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 *NodeKey, root *NodeKey, receiver KVPairReceiver) error {
+func (ndb *nodeDB) extractStateChanges(prevVersion int64, prevRoot, root []byte, receiver KVPairReceiver) error {
 	curIter, err := NewNodeIterator(root, ndb)
 	if err != nil {
 		return err

import.go

@@ -74,7 +74,7 @@ func (i *Importer) writeNode(node *Node) error {
 	bytesCopy := make([]byte, buf.Len())
 	copy(bytesCopy, buf.Bytes())
 
-	if err := i.batch.Set(i.tree.ndb.nodeKey(node.nodeKey), bytesCopy); err != nil {
+	if err := i.batch.Set(i.tree.ndb.nodeKey(node.GetKey()), bytesCopy); err != nil {
 		return err
 	}
 
@@ -135,8 +135,8 @@ func (i *Importer) Add(exportNode *ExportNode) error {
 	} 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.leftNodeKey = node.leftNode.GetKey()
+		node.rightNodeKey = node.rightNode.GetKey()
 		node.size = node.leftNode.size + node.rightNode.size
 		// Update the stack now.
 		if err := i.writeNode(i.stack[stackSize-2]); err != nil {
@@ -169,7 +169,7 @@ func (i *Importer) Commit() error {
 
 	switch len(i.stack) {
 	case 0:
-		if err := i.batch.Set(i.tree.ndb.nodeKey(&NodeKey{version: i.version, nonce: 1}), []byte{}); err != nil {
+		if err := i.batch.Set(i.tree.ndb.nodeKey(GetRootKey(i.version)), []byte{}); err != nil {
 			return err
 		}
 	case 1:
@@ -178,7 +178,7 @@ func (i *Importer) Commit() error {
 			return err
 		}
 		if i.stack[0].nodeKey.version < i.version { // it means there is no update in the given version
-			if err := i.batch.Set(i.tree.ndb.nodeKey(&NodeKey{version: i.version, nonce: 1}), i.tree.ndb.nodeKey(i.stack[0].nodeKey)); err != nil {
+			if err := i.batch.Set(i.tree.ndb.nodeKey(GetRootKey(i.version)), i.tree.ndb.nodeKey(i.stack[0].GetKey())); err != nil {
 				return err
 			}
 		}

iterator.go

@@ -268,7 +268,7 @@ type NodeIterator struct {
 }
 
 // NewNodeIterator returns a new NodeIterator to traverse the tree of the root node.
-func NewNodeIterator(rootKey *NodeKey, ndb *nodeDB) (*NodeIterator, error) {
+func NewNodeIterator(rootKey []byte, ndb *nodeDB) (*NodeIterator, error) {
 	if rootKey == nil {
 		return &NodeIterator{
 			nodesToVisit: []*Node{},

iterator_test.go

@@ -344,7 +344,7 @@ func TestNodeIterator_Success(t *testing.T) {
 	require.NoError(t, err)
 
 	// check if the iterating count is same with the entire node count of the tree
-	itr, err := NewNodeIterator(tree.root.nodeKey, tree.ndb)
+	itr, err := NewNodeIterator(tree.root.GetKey(), tree.ndb)
 	require.NoError(t, err)
 	nodeCount := 0
 	for ; itr.Valid(); itr.Next(false) {
@@ -353,7 +353,7 @@ func TestNodeIterator_Success(t *testing.T) {
 	require.Equal(t, int64(nodeCount), tree.Size()*2-1)
 
 	// check if the skipped node count is right
-	itr, err = NewNodeIterator(tree.root.nodeKey, tree.ndb)
+	itr, err = NewNodeIterator(tree.root.GetKey(), tree.ndb)
 	require.NoError(t, err)
 	updateCount := 0
 	skipCount := 0

keyformat/key_format.go

@@ -139,6 +139,11 @@ func (kf *KeyFormat) Scan(key []byte, args ...interface{}) {
 	}
 }
 
+// Length of the key format.
+func (kf *KeyFormat) Length() int {
+	return kf.length
+}
+
 // Return the prefix as a string.
 func (kf *KeyFormat) Prefix() string {
 	return string([]byte{kf.prefix})