From fe4c3909d4645ee4841806f14d811a228b9f64fe Mon Sep 17 00:00:00 2001
From: Erik Grinaker <erik@interchain.berlin>
Date: Thu, 30 Jul 2020 11:52:37 +0200
Subject: [PATCH] add InitialVersion option (#299)

---
 CHANGELOG.md         |  6 ++++++
 mutable_tree.go      | 12 ++++++++++++
 mutable_tree_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 nodedb.go            |  6 ++++--
 options.go           | 11 ++++++++---
 5 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 45872736f..87e046315 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+### Improvements
+
+- Added `Options.InitialVersion` to specify the initial version to start new IAVL trees from.
+
 ## 0.14.0 (July 2, 2020)
 
 **Important information:** the pruning functionality introduced with IAVL 0.13.0 via the options 
diff --git a/mutable_tree.go b/mutable_tree.go
index 951378aaf..346683d6a 100644
--- a/mutable_tree.go
+++ b/mutable_tree.go
@@ -328,6 +328,7 @@ func (tree *MutableTree) LoadVersion(targetVersion int64) (int64, error) {
 		return 0, nil
 	}
 
+	firstVersion := int64(0)
 	latestVersion := int64(0)
 
 	var latestRoot []byte
@@ -337,6 +338,9 @@ func (tree *MutableTree) LoadVersion(targetVersion int64) (int64, error) {
 			latestVersion = version
 			latestRoot = r
 		}
+		if firstVersion == 0 || version < firstVersion {
+			firstVersion = version
+		}
 	}
 
 	if !(targetVersion == 0 || latestVersion == targetVersion) {
@@ -344,6 +348,11 @@ func (tree *MutableTree) LoadVersion(targetVersion int64) (int64, error) {
 			targetVersion, latestVersion)
 	}
 
+	if firstVersion > 0 && firstVersion < int64(tree.ndb.opts.InitialVersion) {
+		return latestVersion, fmt.Errorf("initial version set to %v, but found earlier version %v",
+			tree.ndb.opts.InitialVersion, firstVersion)
+	}
+
 	t := &ImmutableTree{
 		ndb:     tree.ndb,
 		version: latestVersion,
@@ -439,6 +448,9 @@ func (tree *MutableTree) GetVersioned(key []byte, version int64) (
 // the tree. Returns the hash and new version number.
 func (tree *MutableTree) SaveVersion() ([]byte, int64, error) {
 	version := tree.version + 1
+	if version == 1 && tree.ndb.opts.InitialVersion > 0 {
+		version = int64(tree.ndb.opts.InitialVersion)
+	}
 
 	if tree.versions[version] {
 		// If the version already exists, return an error as we're attempting to overwrite.
diff --git a/mutable_tree_test.go b/mutable_tree_test.go
index e37c519f7..c7c289f23 100644
--- a/mutable_tree_test.go
+++ b/mutable_tree_test.go
@@ -6,6 +6,7 @@ import (
 	"runtime"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	db "github.com/tendermint/tm-db"
@@ -105,6 +106,47 @@ func TestMutableTree_DeleteVersions(t *testing.T) {
 	}
 }
 
+func TestMutableTree_InitialVersion(t *testing.T) {
+	memDB := db.NewMemDB()
+	tree, err := NewMutableTreeWithOpts(memDB, 0, &Options{InitialVersion: 9})
+	require.NoError(t, err)
+
+	tree.Set([]byte("a"), []byte{0x01})
+	_, version, err := tree.SaveVersion()
+	require.NoError(t, err)
+	assert.EqualValues(t, 9, version)
+
+	tree.Set([]byte("b"), []byte{0x02})
+	_, version, err = tree.SaveVersion()
+	require.NoError(t, err)
+	assert.EqualValues(t, 10, version)
+
+	// Reloading the tree with the same initial version is fine
+	tree, err = NewMutableTreeWithOpts(memDB, 0, &Options{InitialVersion: 9})
+	require.NoError(t, err)
+	version, err = tree.Load()
+	require.NoError(t, err)
+	assert.EqualValues(t, 10, version)
+
+	// Reloading the tree with an initial version beyond the lowest should error
+	tree, err = NewMutableTreeWithOpts(memDB, 0, &Options{InitialVersion: 10})
+	require.NoError(t, err)
+	_, err = tree.Load()
+	require.Error(t, err)
+
+	// Reloading the tree with a lower initial version is fine, and new versions can be produced
+	tree, err = NewMutableTreeWithOpts(memDB, 0, &Options{InitialVersion: 3})
+	require.NoError(t, err)
+	version, err = tree.Load()
+	require.NoError(t, err)
+	assert.EqualValues(t, 10, version)
+
+	tree.Set([]byte("c"), []byte{0x03})
+	_, version, err = tree.SaveVersion()
+	require.NoError(t, err)
+	assert.EqualValues(t, 11, version)
+}
+
 func BenchmarkMutableTree_Set(b *testing.B) {
 	db := db.NewDB("test", db.MemDBBackend, "")
 	t, err := NewMutableTree(db, 100000)
diff --git a/nodedb.go b/nodedb.go
index 32bd22f40..0162f79f2 100644
--- a/nodedb.go
+++ b/nodedb.go
@@ -491,8 +491,10 @@ func (ndb *nodeDB) saveRoot(hash []byte, version int64) error {
 	ndb.mtx.Lock()
 	defer ndb.mtx.Unlock()
 
-	if version != ndb.getLatestVersion()+1 {
-		return fmt.Errorf("must save consecutive versions; expected %d, got %d", ndb.getLatestVersion()+1, version)
+	// We allow the initial version to be arbitrary
+	latest := ndb.getLatestVersion()
+	if latest > 0 && version != latest+1 {
+		return fmt.Errorf("must save consecutive versions; expected %d, got %d", latest+1, version)
 	}
 
 	ndb.batch.Set(ndb.rootKey(version), hash)
diff --git a/options.go b/options.go
index 937e8c94d..61c54740d 100644
--- a/options.go
+++ b/options.go
@@ -2,12 +2,17 @@ package iavl
 
 // Options define tree options.
 type Options struct {
+	// Sync synchronously flushes all writes to storage, using e.g. the fsync syscall.
+	// Disabling this significantly improves performance, but can lose data on e.g. power loss.
 	Sync bool
+
+	// InitialVersion specifies the initial version number. If any versions already exist below
+	// this, an error is returned when loading the tree. Only used for the initial SaveVersion()
+	// call.
+	InitialVersion uint64
 }
 
 // DefaultOptions returns the default options for IAVL.
 func DefaultOptions() *Options {
-	return &Options{
-		Sync: false,
-	}
+	return &Options{}
 }