feat: refactor the export traversal order as pre-order

CHANGELOG.md

@@ -13,6 +13,7 @@
 - [#636](https://github.com/cosmos/iavl/pull/636) Speed up rollback method: `LoadVersionForOverwriting`.
 - [#654](https://github.com/cosmos/iavl/pull/654) Add API `TraverseStateChanges` to extract state changes from iavl versions.
 - [#638](https://github.com/cosmos/iavl/pull/638) Make LazyLoadVersion check the opts.InitialVersion, add API `LazyLoadVersionForOverwriting`.
+- [#662](https://github.com/cosmos/iavl/pull/662) Refactor `Export` and `Import` to provide both `post-order` and `pre-order`.
 
 ## 0.19.4 (October 28, 2022)
 

benchmarks/cosmos-exim/main.go

@@ -127,7 +127,7 @@ func runExport(dbPath string) (int64, map[string][]*iavl.ExportNode, error) {
 			return 0, nil, err
 		}
 		start := time.Now().UTC()
-		exporter, err := itree.Export()
+		exporter, err := itree.Export(iavl.PreOrderTraverse)
 		if err != nil {
 			return 0, nil, err
 		}
@@ -177,7 +177,7 @@ func runImport(version int64, exports map[string][]*iavl.ExportNode) error {
 		if err != nil {
 			return err
 		}
-		importer, err := newTree.Import(version)
+		importer, err := newTree.Import(version, iavl.PreOrderTraverse)
 		if err != nil {
 			return err
 		}

export.go

@@ -6,6 +6,14 @@ import (
 	"fmt"
 )
 
+type OrderType int
+
+// OrderTraverse is the type of traversal order to use when exporting and importing.
+const (
+	PreOrderTraverse OrderType = iota
+	PostOrderTraverse
+)
+
 // exportBufferSize is the number of nodes to buffer in the exporter. It improves throughput by
 // processing multiple nodes per context switch, but take care to avoid excessive memory usage,
 // especially since callers may export several IAVL stores in parallel (e.g. the Cosmos SDK).
@@ -27,17 +35,17 @@ type ExportNode struct {
 
 // Exporter exports nodes from an ImmutableTree. It is created by ImmutableTree.Export().
 //
-// Exported nodes can be imported into an empty tree with MutableTree.Import(). Nodes are exported
-// depth-first post-order (LRN), this order must be preserved when importing in order to recreate
-// the same tree structure.
+// Exported nodes can be imported into an empty tree with MutableTree.Import(). Nodes are exported in traverseOrder,
+// this order must be preserved when importing in order to recreate the same tree structure.
 type Exporter struct {
-	tree   *ImmutableTree
-	ch     chan *ExportNode
-	cancel context.CancelFunc
+	tree          *ImmutableTree
+	traverseOrder OrderType
+	ch            chan *ExportNode
+	cancel        context.CancelFunc
 }
 
 // NewExporter creates a new Exporter. Callers must call Close() when done.
-func newExporter(tree *ImmutableTree) (*Exporter, error) {
+func newExporter(tree *ImmutableTree, traverseOrder OrderType) (*Exporter, error) {
 	if tree == nil {
 		return nil, fmt.Errorf("tree is nil: %w", ErrNotInitalizedTree)
 	}
@@ -48,9 +56,10 @@ func newExporter(tree *ImmutableTree) (*Exporter, error) {
 
 	ctx, cancel := context.WithCancel(context.Background())
 	exporter := &Exporter{
-		tree:   tree,
-		ch:     make(chan *ExportNode, exportBufferSize),
-		cancel: cancel,
+		tree:          tree,
+		traverseOrder: traverseOrder,
+		ch:            make(chan *ExportNode, exportBufferSize),
+		cancel:        cancel,
 	}
 
 	tree.ndb.incrVersionReaders(tree.version)
@@ -61,22 +70,41 @@ func newExporter(tree *ImmutableTree) (*Exporter, error) {
 
 // export exports nodes
 func (e *Exporter) export(ctx context.Context) {
-	e.tree.root.traversePost(e.tree, true, func(node *Node) bool {
-		exportNode := &ExportNode{
-			Key:     node.key,
-			Value:   node.value,
-			Version: node.version,
-			Height:  node.subtreeHeight,
-		}
-
-		select {
-		case e.ch <- exportNode:
-			return false
-		case <-ctx.Done():
-			return true
-		}
-	})
-	close(e.ch)
+	defer close(e.ch)
+	switch e.traverseOrder {
+	case PreOrderTraverse:
+		e.tree.root.traverse(e.tree, true, func(node *Node) bool {
+			exportNode := &ExportNode{
+				Key:     node.key,
+				Value:   node.value,
+				Version: node.version,
+				Height:  node.subtreeHeight,
+			}
+
+			select {
+			case e.ch <- exportNode:
+				return false
+			case <-ctx.Done():
+				return true
+			}
+		})
+	case PostOrderTraverse:
+		e.tree.root.traversePost(e.tree, true, func(node *Node) bool {
+			exportNode := &ExportNode{
+				Key:     node.key,
+				Value:   node.value,
+				Version: node.version,
+				Height:  node.subtreeHeight,
+			}
+
+			select {
+			case e.ch <- exportNode:
+				return false
+			case <-ctx.Done():
+				return true
+			}
+		})
+	}
 }
 
 // Next fetches the next exported node, or returns ExportDone when done.

export_test.go

@@ -160,32 +160,46 @@ func setupExportTreeSized(t require.TestingT, treeSize int) *ImmutableTree { //n
 func TestExporter(t *testing.T) {
 	tree := setupExportTreeBasic(t)
 
-	expect := []*ExportNode{
-		{Key: []byte("a"), Value: []byte{1}, Version: 1, Height: 0},
-		{Key: []byte("b"), Value: []byte{2}, Version: 3, Height: 0},
-		{Key: []byte("b"), Value: nil, Version: 3, Height: 1},
-		{Key: []byte("c"), Value: []byte{3}, Version: 3, Height: 0},
-		{Key: []byte("c"), Value: nil, Version: 3, Height: 2},
-		{Key: []byte("d"), Value: []byte{4}, Version: 2, Height: 0},
-		{Key: []byte("e"), Value: []byte{5}, Version: 3, Height: 0},
-		{Key: []byte("e"), Value: nil, Version: 3, Height: 1},
-		{Key: []byte("d"), Value: nil, Version: 3, Height: 3},
+	expects := map[OrderType][]*ExportNode{
+		PreOrderTraverse: {
+			{Key: []byte("d"), Value: nil, Version: 3, Height: 3},
+			{Key: []byte("c"), Value: nil, Version: 3, Height: 2},
+			{Key: []byte("b"), Value: nil, Version: 3, Height: 1},
+			{Key: []byte("a"), Value: []byte{1}, Version: 1, Height: 0},
+			{Key: []byte("b"), Value: []byte{2}, Version: 3, Height: 0},
+			{Key: []byte("c"), Value: []byte{3}, Version: 3, Height: 0},
+			{Key: []byte("e"), Value: nil, Version: 3, Height: 1},
+			{Key: []byte("d"), Value: []byte{4}, Version: 2, Height: 0},
+			{Key: []byte("e"), Value: []byte{5}, Version: 3, Height: 0},
+		},
+		PostOrderTraverse: {
+			{Key: []byte("a"), Value: []byte{1}, Version: 1, Height: 0},
+			{Key: []byte("b"), Value: []byte{2}, Version: 3, Height: 0},
+			{Key: []byte("b"), Value: nil, Version: 3, Height: 1},
+			{Key: []byte("c"), Value: []byte{3}, Version: 3, Height: 0},
+			{Key: []byte("c"), Value: nil, Version: 3, Height: 2},
+			{Key: []byte("d"), Value: []byte{4}, Version: 2, Height: 0},
+			{Key: []byte("e"), Value: []byte{5}, Version: 3, Height: 0},
+			{Key: []byte("e"), Value: nil, Version: 3, Height: 1},
+			{Key: []byte("d"), Value: nil, Version: 3, Height: 3},
+		},
 	}
-
-	actual := make([]*ExportNode, 0, len(expect))
-	exporter, err := tree.Export()
-	require.NoError(t, err)
-	defer exporter.Close()
-	for {
-		node, err := exporter.Next()
-		if err == ErrorExportDone {
-			break
-		}
+	for orderType, expect := range expects {
+		actual := make([]*ExportNode, 0, len(expect))
+		exporter, err := tree.Export(orderType)
 		require.NoError(t, err)
-		actual = append(actual, node)
-	}
+		defer exporter.Close()
+		for {
+			node, err := exporter.Next()
+			if err == ErrorExportDone {
+				break
+			}
+			require.NoError(t, err)
+			actual = append(actual, node)
+		}
 
-	assert.Equal(t, expect, actual)
+		assert.Equal(t, expect, actual)
+	}
 }
 
 func TestExporter_Import(t *testing.T) {
@@ -197,59 +211,60 @@ func TestExporter_Import(t *testing.T) {
 		testcases["sized tree"] = setupExportTreeSized(t, 4096)
 		testcases["random tree"] = setupExportTreeRandom(t)
 	}
+	for _, orderType := range []OrderType{PostOrderTraverse, PreOrderTraverse} {
+		for desc, tree := range testcases {
+			tree := tree
+			t.Run(desc, func(t *testing.T) {
+				t.Parallel()
 
-	for desc, tree := range testcases {
-		tree := tree
-		t.Run(desc, func(t *testing.T) {
-			t.Parallel()
-
-			exporter, err := tree.Export()
-			require.NoError(t, err)
-			defer exporter.Close()
-
-			newTree, err := NewMutableTree(db.NewMemDB(), 0, false)
-			require.NoError(t, err)
-			importer, err := newTree.Import(tree.Version())
-			require.NoError(t, err)
-			defer importer.Close()
+				exporter, err := tree.Export(orderType)
+				require.NoError(t, err)
+				defer exporter.Close()
 
-			for {
-				item, err := exporter.Next()
-				if err == ErrorExportDone {
-					err = importer.Commit()
+				newTree, err := NewMutableTree(db.NewMemDB(), 0, false)
+				require.NoError(t, err)
+				importer, err := newTree.Import(tree.Version(), orderType)
+				require.NoError(t, err)
+				defer importer.Close()
+
+				for {
+					item, err := exporter.Next()
+					if err == ErrorExportDone {
+						err = importer.Commit()
+						require.NoError(t, err)
+						break
+					}
+					require.NoError(t, err)
+					err = importer.Add(item)
 					require.NoError(t, err)
-					break
 				}
+
+				treeHash, err := tree.Hash()
 				require.NoError(t, err)
-				err = importer.Add(item)
+				newTreeHash, err := newTree.Hash()
 				require.NoError(t, err)
-			}
-
-			treeHash, err := tree.Hash()
-			require.NoError(t, err)
-			newTreeHash, err := newTree.Hash()
-			require.NoError(t, err)
 
-			require.Equal(t, treeHash, newTreeHash, "Tree hash mismatch")
-			require.Equal(t, tree.Size(), newTree.Size(), "Tree size mismatch")
-			require.Equal(t, tree.Version(), newTree.Version(), "Tree version mismatch")
+				require.Equal(t, treeHash, newTreeHash, "Tree hash mismatch")
+				require.Equal(t, tree.Size(), newTree.Size(), "Tree size mismatch")
+				require.Equal(t, tree.Version(), newTree.Version(), "Tree version mismatch")
 
-			tree.Iterate(func(key, value []byte) bool { //nolint:errcheck
-				index, _, err := tree.GetWithIndex(key)
-				require.NoError(t, err)
-				newIndex, newValue, err := newTree.GetWithIndex(key)
-				require.NoError(t, err)
-				require.Equal(t, index, newIndex, "Index mismatch for key %v", key)
-				require.Equal(t, value, newValue, "Value mismatch for key %v", key)
-				return false
+				tree.Iterate(func(key, value []byte) bool { //nolint:errcheck
+					index, _, err := tree.GetWithIndex(key)
+					require.NoError(t, err)
+					newIndex, newValue, err := newTree.GetWithIndex(key)
+					require.NoError(t, err)
+					require.Equal(t, index, newIndex, "Index mismatch for key %v", key)
+					require.Equal(t, value, newValue, "Value mismatch for key %v", key)
+					return false
+				})
 			})
-		})
+		}
 	}
 }
 
 func TestExporter_Close(t *testing.T) {
 	tree := setupExportTreeSized(t, 4096)
-	exporter, err := tree.Export()
+	exporter, err := tree.Export(PreOrderTraverse)
 	require.NoError(t, err)
 
 	node, err := exporter.Next()
@@ -292,7 +307,7 @@ func TestExporter_DeleteVersionErrors(t *testing.T) {
 
 	itree, err := tree.GetImmutable(2)
 	require.NoError(t, err)
-	exporter, err := itree.Export()
+	exporter, err := itree.Export(PostOrderTraverse)
 	require.NoError(t, err)
 	defer exporter.Close()
 
@@ -311,17 +326,34 @@ func BenchmarkExport(b *testing.B) {
 	b.StopTimer()
 	tree := setupExportTreeSized(b, 4096)
 	b.StartTimer()
-	for n := 0; n < b.N; n++ {
-		exporter, err := tree.Export()
-		require.NoError(b, err)
-		for {
-			_, err := exporter.Next()
-			if err == ErrorExportDone {
-				break
-			} else if err != nil {
-				b.Error(err)
+	b.Run("post order export", func(sub *testing.B) {
+		for n := 0; n < sub.N; n++ {
+			exporter, err := tree.Export(PostOrderTraverse)
+			require.NoError(sub, err)
+			for {
+				_, err := exporter.Next()
+				if err == ErrorExportDone {
+					break
+				} else if err != nil {
+					sub.Error(err)
+				}
 			}
+			exporter.Close()
 		}
-		exporter.Close()
-	}
+	})
+	b.Run("pre order export", func(sub *testing.B) {
+		for n := 0; n < sub.N; n++ {
+			exporter, err := tree.Export(PreOrderTraverse)
+			require.NoError(sub, err)
+			for {
+				_, err := exporter.Next()
+				if err == ErrorExportDone {
+					break
+				} else if err != nil {
+					sub.Error(err)
+				}
+			}
+			exporter.Close()
+		}
+	})
 }

immutable_tree.go

@@ -155,8 +155,8 @@ func (t *ImmutableTree) Hash() ([]byte, error) {
 
 // Export returns an iterator that exports tree nodes as ExportNodes. These nodes can be
 // imported with MutableTree.Import() to recreate an identical tree.
-func (t *ImmutableTree) Export() (*Exporter, error) {
-	return newExporter(t)
+func (t *ImmutableTree) Export(traverseOrder OrderType) (*Exporter, error) {
+	return newExporter(t, traverseOrder)
 }
 
 // GetWithIndex returns the index and value of the specified key if it exists, or nil and the next index