Problem: memiavl import API not implemented

memiavl/import.go

@@ -0,0 +1,183 @@
+package memiavl
+
+import (
+	"bufio"
+	"encoding/binary"
+	"fmt"
+	"math"
+	"os"
+	"path/filepath"
+
+	"github.com/cosmos/iavl"
+	"github.com/pkg/errors"
+)
+
+// Import a stream of `iavl.ExportNode`s into a new snapshot.
+func Import(dir string, version int64, nodes chan *iavl.ExportNode, writeHashIndex bool) (returnErr error) {
+	if version > int64(math.MaxUint32) {
+		return errors.New("version overflows uint32")
+	}
+
+	nodesFile := filepath.Join(dir, FileNameNodes)
+	kvsFile := filepath.Join(dir, FileNameKVs)
+	kvsIndexFile := filepath.Join(dir, FileNameKVIndex)
+
+	fpNodes, err := createFile(nodesFile)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err := fpNodes.Close(); returnErr == nil {
+			returnErr = err
+		}
+	}()
+
+	fpKVs, err := createFile(kvsFile)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err := fpKVs.Close(); returnErr == nil {
+			returnErr = err
+		}
+	}()
+
+	nodesWriter := bufio.NewWriter(fpNodes)
+	kvsWriter := bufio.NewWriter(fpKVs)
+
+	i := &importer{
+		snapshotWriter: *newSnapshotWriter(nodesWriter, kvsWriter),
+	}
+
+	var counter int
+	for node := range nodes {
+		if err := i.Add(node); err != nil {
+			return err
+		}
+		counter++
+	}
+
+	if len(i.nodeStack) > 1 {
+		return errors.Errorf("invalid node structure, found stack size %v after imported", len(i.nodeStack))
+	}
+
+	var rootIndex uint32
+	if len(i.nodeStack) == 0 {
+		rootIndex = EmptyRootNodeIndex
+	} else {
+		rootIndex = uint32(counter - 1)
+
+		if err := nodesWriter.Flush(); err != nil {
+			return err
+		}
+		if err := kvsWriter.Flush(); err != nil {
+			return err
+		}
+
+		if err := fpKVs.Sync(); err != nil {
+			return err
+		}
+		if err := fpNodes.Sync(); err != nil {
+			return err
+		}
+
+		if writeHashIndex {
+			// re-open kvs file for reading
+			input, err := os.Open(kvsFile)
+			if err != nil {
+				return err
+			}
+			defer func() {
+				if err := input.Close(); returnErr == nil {
+					returnErr = err
+				}
+			}()
+			// N = 2L-1
+			leaves := (counter + 1) / 2
+			if err := buildIndex(input, kvsIndexFile, dir, leaves); err != nil {
+				return fmt.Errorf("build MPHF index failed: %w", err)
+			}
+		}
+	}
+
+	// write metadata
+	var metadataBuf [SizeMetadata]byte
+	binary.LittleEndian.PutUint32(metadataBuf[:], SnapshotFileMagic)
+	binary.LittleEndian.PutUint32(metadataBuf[4:], SnapshotFormat)
+	binary.LittleEndian.PutUint32(metadataBuf[8:], uint32(version))
+	binary.LittleEndian.PutUint32(metadataBuf[12:], rootIndex)
+
+	metadataFile := filepath.Join(dir, FileNameMetadata)
+	fpMetadata, err := createFile(metadataFile)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if err := fpMetadata.Close(); returnErr == nil {
+			returnErr = err
+		}
+	}()
+
+	if _, err := fpMetadata.Write(metadataBuf[:]); err != nil {
+		return err
+	}
+
+	return fpMetadata.Sync()
+}
+
+type importer struct {
+	snapshotWriter
+
+	indexStack []uint32
+	nodeStack  []*MemNode
+}
+
+func (i *importer) Add(n *iavl.ExportNode) error {
+	if n.Version > int64(math.MaxUint32) {
+		return errors.New("version overflows uint32")
+	}
+
+	if n.Height == 0 {
+		node := &MemNode{
+			height:  uint8(n.Height),
+			size:    1,
+			version: uint32(n.Version),
+			key:     n.Key,
+			value:   n.Value,
+		}
+		nodeHash := node.Hash()
+		idx, err := i.writeLeaf(node.version, node.key, node.value, nodeHash)
+		if err != nil {
+			return err
+		}
+		i.indexStack = append(i.indexStack, idx)
+		i.nodeStack = append(i.nodeStack, node)
+		return nil
+	}
+
+	// branch node
+	leftIndex := i.indexStack[len(i.indexStack)-2]
+	leftNode := i.nodeStack[len(i.nodeStack)-2]
+	rightNode := i.nodeStack[len(i.nodeStack)-1]
+
+	node := &MemNode{
+		height:  uint8(n.Height),
+		size:    leftNode.size + rightNode.size,
+		version: uint32(n.Version),
+		key:     n.Key,
+		left:    leftNode,
+		right:   rightNode,
+	}
+	nodeHash := node.Hash()
+	idx, err := i.writeBranch(node.version, uint32(node.size), node.height, leftIndex+1, nodeHash)
+	if err != nil {
+		return err
+	}
+
+	i.indexStack = i.indexStack[:len(i.indexStack)-2]
+	i.indexStack = append(i.indexStack, idx)
+
+	i.nodeStack = i.nodeStack[:len(i.nodeStack)-2]
+	i.nodeStack = append(i.nodeStack, node)
+	return nil
+}

memiavl/snapshot.go

@@ -404,45 +404,56 @@ func (w *snapshotWriter) writeKeyValue(key, value []byte) error {
 	return nil
 }
 
-// writeRecursive write the node recursively in depth-first post-order,
-// returns `(nodeIndex, err)`.
-func (w *snapshotWriter) writeRecursive(node Node) (uint32, error) {
+func (w *snapshotWriter) writeNode(node, hash []byte) (uint32, error) {
+	if _, err := w.nodesWriter.Write(node); err != nil {
+		return 0, err
+	}
+	if _, err := w.nodesWriter.Write(hash); err != nil {
+		return 0, err
+	}
+
+	i := w.nodeIndex
+	w.nodeIndex++
+	return i, nil
+}
+
+func (w *snapshotWriter) writeLeaf(version uint32, key, value, hash []byte) (uint32, error) {
 	var buf [SizeNodeWithoutHash]byte
+	binary.LittleEndian.PutUint32(buf[OffsetVersion:], version)
+	keyOffset := w.kvsOffset
+	if err := w.writeKeyValue(key, value); err != nil {
+		return 0, err
+	}
+	binary.LittleEndian.PutUint64(buf[OffsetKeyOffset:], keyOffset)
 
-	buf[OffsetHeight] = node.Height()
-	binary.LittleEndian.PutUint32(buf[OffsetVersion:], node.Version())
+	return w.writeNode(buf[:], hash)
+}
 
-	if isLeaf(node) {
-		keyOffset := w.kvsOffset
-		if err := w.writeKeyValue(node.Key(), node.Value()); err != nil {
-			return 0, err
-		}
+func (w *snapshotWriter) writeBranch(version, size uint32, height uint8, keyNode uint32, hash []byte) (uint32, error) {
+	var buf [SizeNodeWithoutHash]byte
+	buf[OffsetHeight] = height
+	binary.LittleEndian.PutUint32(buf[OffsetVersion:], version)
+	binary.LittleEndian.PutUint32(buf[OffsetSize:], size)
+	binary.LittleEndian.PutUint32(buf[OffsetKeyNode:], keyNode)
 
-		binary.LittleEndian.PutUint64(buf[OffsetKeyOffset:], keyOffset)
-	} else {
-		binary.LittleEndian.PutUint32(buf[OffsetSize:], uint32(node.Size()))
+	return w.writeNode(buf[:], hash)
+}
 
-		// store the minimal key from right subtree, but propagate the one from left subtree
-		leftIndex, err := w.writeRecursive(node.Left())
-		if err != nil {
-			return 0, err
-		}
-		if _, err = w.writeRecursive(node.Right()); err != nil {
-			return 0, err
-		}
-		binary.LittleEndian.PutUint32(buf[OffsetKeyNode:], leftIndex+1)
+// writeRecursive write the node recursively in depth-first post-order,
+// returns `(nodeIndex, err)`.
+func (w *snapshotWriter) writeRecursive(node Node) (uint32, error) {
+	if isLeaf(node) {
+		return w.writeLeaf(node.Version(), node.Key(), node.Value(), node.Hash())
 	}
 
-	if _, err := w.nodesWriter.Write(buf[:]); err != nil {
+	leftIndex, err := w.writeRecursive(node.Left())
+	if err != nil {
 		return 0, err
 	}
-	if _, err := w.nodesWriter.Write(node.Hash()); err != nil {
+	if _, err := w.writeRecursive(node.Right()); err != nil {
 		return 0, err
 	}
-
-	i := w.nodeIndex
-	w.nodeIndex++
-	return i, nil
+	return w.writeBranch(node.Version(), uint32(node.Size()), node.Height(), leftIndex+1, node.Hash())
 }
 
 // buildIndex build MPHF index for the kvs file.

memiavl/snapshot_test.go

@@ -85,3 +85,48 @@ func TestSnapshotExport(t *testing.T) {
 
 	require.Equal(t, expNodes, nodes)
 }
+
+func TestSnapshotImportExport(t *testing.T) {
+	// setup test tree
+	tree := NewEmptyTree(0)
+	for _, changes := range ChangeSets {
+		applyChangeSet(tree, changes)
+		_, _, err := tree.SaveVersion(true)
+		require.NoError(t, err)
+	}
+
+	snapshotDir := t.TempDir()
+	require.NoError(t, tree.WriteSnapshot(snapshotDir, true))
+	snapshot, err := OpenSnapshot(snapshotDir)
+	require.NoError(t, err)
+
+	ch := make(chan *iavl.ExportNode)
+
+	go func() {
+		defer close(ch)
+
+		exporter := snapshot.Export()
+		for {
+			node, err := exporter.Next()
+			if err == iavl.ExportDone {
+				break
+			}
+			require.NoError(t, err)
+			ch <- node
+		}
+	}()
+
+	snapshotDir2 := t.TempDir()
+	err = Import(snapshotDir2, tree.Version(), ch, true)
+	require.NoError(t, err)
+
+	snapshot2, err := OpenSnapshot(snapshotDir2)
+	require.NoError(t, err)
+	require.Equal(t, snapshot.RootNode().Hash(), snapshot2.RootNode().Hash())
+
+	// verify all the node hashes in snapshot
+	for i := 0; i < snapshot2.nodesLen(); i++ {
+		node := snapshot2.Node(uint32(i))
+		require.Equal(t, node.Hash(), HashNode(node))
+	}
+}