feat: add post-order export & restoring

export.go

@@ -2,31 +2,70 @@ package iavl
 
 import "fmt"
 
+// TraverseOrderType is the type of the order in which the tree is traversed.
+type TraverseOrderType uint8
+
+const (
+	PreOrder TraverseOrderType = iota
+	PostOrder
+)
+
 type Exporter struct {
 	tree  *Tree
 	out   chan *Node
 	errCh chan error
 }
 
-func (tree *Tree) ExportPreOrder() *Exporter {
+func (tree *Tree) Export(order TraverseOrderType) *Exporter {
 	exporter := &Exporter{
 		tree:  tree,
 		out:   make(chan *Node),
 		errCh: make(chan error),
 	}
-	go func() {
+
+	go func(traverseOrder TraverseOrderType) {
 		defer close(exporter.out)
 		defer close(exporter.errCh)
-		exporter.preOrderNext(tree.root)
-	}()
+
+		if traverseOrder == PostOrder {
+			exporter.postOrderNext(tree.root)
+		} else if traverseOrder == PreOrder {
+			exporter.preOrderNext(tree.root)
+		}
+	}(order)
+
 	return exporter
 }
 
+func (e *Exporter) postOrderNext(node *Node) {
+	if node.isLeaf() {
+		e.out <- node
+		return
+	}
+
+	left, err := node.getLeftNode(e.tree)
+	if err != nil {
+		e.errCh <- err
+		return
+	}
+	e.postOrderNext(left)
+
+	right, err := node.getRightNode(e.tree)
+	if err != nil {
+		e.errCh <- err
+		return
+	}
+	e.postOrderNext(right)
+
+	e.out <- node
+}
+
 func (e *Exporter) preOrderNext(node *Node) {
 	e.out <- node
 	if node.isLeaf() {
 		return
 	}
+
 	left, err := node.getLeftNode(e.tree)
 	if err != nil {
 		e.errCh <- err

snapshot.go

@@ -73,6 +73,7 @@ func (sql *SqliteDb) Snapshot(ctx context.Context, tree *Tree) error {
 type SnapshotOptions struct {
 	StoreLeafValues bool
 	SaveTree        bool
+	TraverseOrder   TraverseOrderType
 }
 
 func NewIngestSnapshotConnection(snapshotDbPath string) (*sqlite3.Conn, error) {
@@ -252,93 +253,14 @@ func (sql *SqliteDb) WriteSnapshot(
 	}
 
 	var (
-		step           func() (*Node, error)
-		maybeFlush     func() error
-		count          int
-		uniqueVersions = make(map[int64]struct{})
+		root           *Node
+		uniqueVersions map[int64]struct{}
 	)
-	maybeFlush = func() error {
-		count++
-		if count%snap.batchSize == 0 {
-			if err = snap.flush(); err != nil {
-				return err
-			}
-			if err = snap.prepareWrite(); err != nil {
-				return err
-			}
-		}
-		return nil
-	}
-
-	step = func() (*Node, error) {
-		snapshotNode, err := nextFn()
-		if err != nil {
-			return nil, err
-		}
-		ordinal := snap.ordinal
-		snap.ordinal++
-
-		node := &Node{
-			key:           snapshotNode.Key,
-			subtreeHeight: snapshotNode.Height,
-			nodeKey:       NewNodeKey(snapshotNode.Version, uint32(ordinal)),
-		}
-		if node.subtreeHeight == 0 {
-			node.value = snapshotNode.Value
-			node.size = 1
-			node._hash(snapshotNode.Version)
-			if !opts.StoreLeafValues {
-				node.value = nil
-			}
-			nodeBz, err := node.Bytes()
-			if err != nil {
-				return nil, err
-			}
-			if err = snap.snapshotInsert.Exec(ordinal, snapshotNode.Version, ordinal, nodeBz); err != nil {
-				return nil, err
-			}
-			if err = snap.leafInsert.Exec(snapshotNode.Version, ordinal, nodeBz); err != nil {
-				return nil, err
-			}
-			if err = maybeFlush(); err != nil {
-				return nil, err
-			}
-			return node, nil
-		}
-
-		node.leftNode, err = step()
-		if err != nil {
-			return nil, err
-		}
-		node.leftNodeKey = node.leftNode.nodeKey
-		node.rightNode, err = step()
-		if err != nil {
-			return nil, err
-		}
-		node.rightNodeKey = node.rightNode.nodeKey
-
-		node.size = node.leftNode.size + node.rightNode.size
-		node._hash(snapshotNode.Version)
-		node.leftNode = nil
-		node.rightNode = nil
-
-		nodeBz, err := node.Bytes()
-		if err != nil {
-			return nil, err
-		}
-		if err = snap.snapshotInsert.Exec(ordinal, snapshotNode.Version, ordinal, nodeBz); err != nil {
-			return nil, err
-		}
-		if err = snap.treeInsert.Exec(snapshotNode.Version, ordinal, nodeBz); err != nil {
-			return nil, err
-		}
-		uniqueVersions[snapshotNode.Version] = struct{}{}
-		if err = maybeFlush(); err != nil {
-			return nil, err
-		}
-		return node, nil
+	if opts.TraverseOrder == PostOrder {
+		root, uniqueVersions, err = snap.restorePostOrderStep(nextFn, opts.StoreLeafValues)
+	} else if opts.TraverseOrder == PreOrder {
+		root, uniqueVersions, err = snap.restorePreOrderStep(nextFn, opts.StoreLeafValues)
 	}
-	root, err := step()
 	if err != nil {
 		return nil, err
 	}
@@ -592,6 +514,165 @@ func (snap *sqliteSnapshot) prepareWrite() error {
 	return err
 }
 
+func (snap *sqliteSnapshot) restorePostOrderStep(nextFn func() (*SnapshotNode, error), isStoreLeafValues bool) (*Node, map[int64]struct{}, error) {
+	var (
+		snapshotNode   *SnapshotNode
+		err            error
+		count          int
+		stack          []*Node
+		uniqueVersions = make(map[int64]struct{})
+	)
+
+	for {
+		snapshotNode, err = nextFn()
+		if err != nil || snapshotNode == nil {
+			break
+		}
+
+		uniqueVersions[snapshotNode.Version] = struct{}{}
+		node := &Node{
+			key:           snapshotNode.Key,
+			subtreeHeight: snapshotNode.Height,
+			nodeKey:       NewNodeKey(snapshotNode.Version, uint32(snap.ordinal)),
+		}
+
+		stackSize := len(stack)
+		if node.isLeaf() {
+			node.value = snapshotNode.Value
+			node.size = 1
+			node._hash(snapshotNode.Version)
+			if !isStoreLeafValues {
+				node.value = nil
+			}
+
+			count++
+			if err := snap.writeSnapNode(node, snapshotNode.Version, count); err != nil {
+				return nil, nil, err
+			}
+		} else if stackSize >= 2 && stack[stackSize-1].subtreeHeight < node.subtreeHeight && stack[stackSize-2].subtreeHeight < node.subtreeHeight {
+			node.leftNode = stack[stackSize-2]
+			node.leftNodeKey = node.leftNode.nodeKey
+			node.rightNode = stack[stackSize-1]
+			node.rightNodeKey = node.rightNode.nodeKey
+			node.size = node.leftNode.size + node.rightNode.size
+			node._hash(snapshotNode.Version)
+			stack = stack[:stackSize-2]
+
+			node.leftNode = nil
+			node.rightNode = nil
+
+			count++
+			if err := snap.writeSnapNode(node, snapshotNode.Version, count); err != nil {
+				return nil, nil, err
+			}
+		}
+
+		stack = append(stack, node)
+		snap.ordinal++
+	}
+
+	if err != nil && !errors.Is(err, ErrorExportDone) {
+		return nil, nil, err
+	}
+
+	if len(stack) != 1 {
+		return nil, nil, fmt.Errorf("expected stack size 1, got %d", len(stack))
+	}
+
+	return stack[0], uniqueVersions, nil
+}
+
+func (snap *sqliteSnapshot) restorePreOrderStep(nextFn func() (*SnapshotNode, error), isStoreLeafValues bool) (*Node, map[int64]struct{}, error) {
+	var (
+		count          int
+		step           func() (*Node, error)
+		uniqueVersions = make(map[int64]struct{})
+	)
+
+	step = func() (*Node, error) {
+		snapshotNode, err := nextFn()
+		if err != nil {
+			return nil, err
+		}
+
+		node := &Node{
+			key:           snapshotNode.Key,
+			subtreeHeight: snapshotNode.Height,
+			nodeKey:       NewNodeKey(snapshotNode.Version, uint32(snap.ordinal)),
+		}
+
+		if node.isLeaf() {
+			node.value = snapshotNode.Value
+			node.size = 1
+			node._hash(snapshotNode.Version)
+			if !isStoreLeafValues {
+				node.value = nil
+			}
+		} else {
+			node.leftNode, err = step()
+			if err != nil {
+				return nil, err
+			}
+			node.leftNodeKey = node.leftNode.nodeKey
+			node.rightNode, err = step()
+			if err != nil {
+				return nil, err
+			}
+			node.rightNodeKey = node.rightNode.nodeKey
+
+			node.size = node.leftNode.size + node.rightNode.size
+			node._hash(snapshotNode.Version)
+			node.leftNode = nil
+			node.rightNode = nil
+			uniqueVersions[snapshotNode.Version] = struct{}{}
+		}
+
+		count++
+		if err := snap.writeSnapNode(node, snapshotNode.Version, count); err != nil {
+			return nil, err
+		}
+		snap.ordinal++
+
+		return node, nil
+	}
+
+	node, err := step()
+
+	return node, uniqueVersions, err
+}
+
+func (snap *sqliteSnapshot) writeSnapNode(node *Node, version int64, count int) error {
+	ordinal := snap.ordinal
+
+	nodeBz, err := node.Bytes()
+	if err != nil {
+		return err
+	}
+	if err = snap.snapshotInsert.Exec(ordinal, version, ordinal, nodeBz); err != nil {
+		return err
+	}
+	if node.isLeaf() {
+		if err = snap.leafInsert.Exec(version, ordinal, nodeBz); err != nil {
+			return err
+		}
+	} else {
+		if err = snap.treeInsert.Exec(version, ordinal, nodeBz); err != nil {
+			return err
+		}
+	}
+
+	if count%snap.batchSize == 0 {
+		if err := snap.flush(); err != nil {
+			return err
+		}
+		if err := snap.prepareWrite(); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func rehashTree(node *Node) {
 	if node.isLeaf() {
 		return

tree_test.go

@@ -159,8 +159,7 @@ func TestTree_Hash(t *testing.T) {
 	opts.Until = 100
 	opts.UntilHash = "0101e1d6f3158dcb7221acd7ed36ce19f2ef26847ffea7ce69232e362539e5cf"
 
-	ctx := context.Background()
-	ctx, cancel := context.WithCancel(ctx)
+	_, cancel := context.WithCancel(context.Background())
 
 	testStart := time.Now()
 	multiTree := NewMultiTree(tmpDir, TreeOptions{CheckpointInterval: 10})
@@ -205,13 +204,14 @@ func TestTree_Build_Load(t *testing.T) {
 
 	// export the tree at version 12,000 and import it into a sql db
 	ctx := context.Background()
+	traverseOrder := PreOrder
 	restoreMt := NewMultiTree(t.TempDir(), TreeOptions{CheckpointInterval: 4000})
 	for sk, tree := range multiTree.Trees {
 		require.NoError(t, restoreMt.MountTree(sk))
-		exporter := tree.ExportPreOrder()
+		exporter := tree.Export(traverseOrder)
 
 		restoreTree := restoreMt.Trees[sk]
-		_, err := restoreTree.sql.WriteSnapshot(ctx, tree.Version(), exporter.Next, SnapshotOptions{SaveTree: true})
+		_, err := restoreTree.sql.WriteSnapshot(ctx, tree.Version(), exporter.Next, SnapshotOptions{SaveTree: true, TraverseOrder: traverseOrder})
 		require.NoError(t, err)
 		require.NoError(t, restoreTree.LoadSnapshot(tree.Version()))
 	}