First attempt

node.go

@@ -100,14 +100,14 @@ func (node *Node) String() string {
 }
 
 // clone creates a shallow copy of a node with its hash set to nil.
-func (node *Node) clone() *Node {
+func (node *Node) clone(version int64) *Node {
 	if node.isLeaf() {
 		cmn.PanicSanity("Attempt to copy a leaf node")
 	}
 	return &Node{
 		key:       node.key,
 		height:    node.height,
-		version:   node.version,
+		version:   version,
 		size:      node.size,
 		hash:      nil,
 		leftHash:  node.leftHash,
@@ -311,7 +311,7 @@ func (node *Node) set(t *Tree, key []byte, value []byte) (
 		}
 	} else {
 		orphaned = append(orphaned, node)
-		node = node.clone()
+		node = node.clone(version)
 
 		if bytes.Compare(key, node.key) < 0 {
 			var leftOrphaned []*Node
@@ -341,6 +341,8 @@ func (node *Node) set(t *Tree, key []byte, value []byte) (
 func (node *Node) remove(t *Tree, key []byte) (
 	newHash []byte, newNode *Node, newKey []byte, value []byte, orphaned []*Node,
 ) {
+	version := t.version + 1
+
 	if node.isLeaf() {
 		if bytes.Equal(key, node.key) {
 			return nil, nil, nil, node.value, []*Node{node}
@@ -362,7 +364,7 @@ func (node *Node) remove(t *Tree, key []byte) (
 		}
 		orphaned = append(orphaned, node)
 
-		newNode := node.clone()
+		newNode := node.clone(version)
 		newNode.leftHash, newNode.leftNode = newLeftHash, newLeftNode
 		newNode.calcHeightAndSize(t)
 		newNode, balanceOrphaned := newNode.balance(t)
@@ -382,7 +384,7 @@ func (node *Node) remove(t *Tree, key []byte) (
 		}
 		orphaned = append(orphaned, node)
 
-		newNode := node.clone()
+		newNode := node.clone(version)
 		newNode.rightHash, newNode.rightNode = newRightHash, newRightNode
 		if newKey != nil {
 			newNode.key = newKey
@@ -410,10 +412,12 @@ func (node *Node) getRightNode(t *Tree) *Node {
 
 // Rotate right and return the new node and orphan.
 func (node *Node) rotateRight(t *Tree) (newNode *Node, orphan *Node) {
+	version := t.version + 1
+
 	// TODO: optimize balance & rotate.
-	node = node.clone()
+	node = node.clone(version)
 	l := node.getLeftNode(t)
-	_l := l.clone()
+	_l := l.clone(version)
 
 	_lrHash, _lrCached := _l.rightHash, _l.rightNode
 	_l.rightHash, _l.rightNode = node.hash, node
@@ -427,10 +431,12 @@ func (node *Node) rotateRight(t *Tree) (newNode *Node, orphan *Node) {
 
 // Rotate left and return the new node and orphan.
 func (node *Node) rotateLeft(t *Tree) (newNode *Node, orphan *Node) {
+	version := t.version + 1
+
 	// TODO: optimize balance & rotate.
-	node = node.clone()
+	node = node.clone(version)
 	r := node.getRightNode(t)
-	_r := r.clone()
+	_r := r.clone(version)
 
 	_rlHash, _rlCached := _r.leftHash, _r.leftNode
 	_r.leftHash, _r.leftNode = node.hash, node

nodedb.go

@@ -24,8 +24,8 @@ var (
 	// which it is expected to exist - which starts out by being the version
 	// of the node being orphaned.
 	orphansPrefix    = "o/"
-	orphansPrefixFmt = "o/%d/"      // o/<version>/
-	orphansKeyFmt    = "o/%d/%d/%x" // o/<version>/<version>/<hash>
+	orphansPrefixFmt = "o/%d/"      // o/<last-version>/
+	orphansKeyFmt    = "o/%d/%d/%x" // o/<last-version>/<first-version>/<hash>
 
 	// These keys are used for the orphan reverse-lookups by node hash.
 	orphansIndexPrefix = "O/"
@@ -183,6 +183,8 @@ func (ndb *nodeDB) Unorphan(hash []byte) {
 }
 
 // Saves orphaned nodes to disk under a special prefix.
+// version: the new version being saved.
+// orphans: the orphan nodes created since version-1
 func (ndb *nodeDB) SaveOrphans(version int64, orphans map[string]int64) {
 	ndb.mtx.Lock()
 	defer ndb.mtx.Unlock()
@@ -197,7 +199,7 @@ func (ndb *nodeDB) SaveOrphans(version int64, orphans map[string]int64) {
 // Saves a single orphan to disk.
 func (ndb *nodeDB) saveOrphan(hash []byte, fromVersion, toVersion int64) {
 	if fromVersion > toVersion {
-		cmn.PanicSanity("Orphan expires before it comes alive")
+		panic(fmt.Sprintf("Orphan expires before it comes alive.  %d > %d", fromVersion, toVersion))
 	}
 	key := ndb.orphanKey(fromVersion, toVersion, hash)
 	ndb.batch.Set(key, hash)
@@ -293,14 +295,13 @@ func (ndb *nodeDB) getPreviousVersion(version int64) int64 {
 
 // deleteRoot deletes the root entry from disk, but not the node it points to.
 func (ndb *nodeDB) deleteRoot(version int64) {
-	key := ndb.rootKey(version)
-	ndb.batch.Delete(key)
-
-	delete(ndb.versionCache, version)
-
 	if version == ndb.getLatestVersion() {
 		cmn.PanicSanity("Tried to delete latest version")
 	}
+
+	key := ndb.rootKey(version)
+	ndb.batch.Delete(key)
+	delete(ndb.versionCache, version)
 }
 
 func (ndb *nodeDB) traverseOrphans(fn func(k, v []byte)) {
@@ -397,7 +398,7 @@ func (ndb *nodeDB) saveRoot(hash []byte, version int64) error {
 	ndb.mtx.Lock()
 	defer ndb.mtx.Unlock()
 
-	if version <= ndb.getLatestVersion() {
+	if version != ndb.getLatestVersion()+1 {
 		return errors.New("can't save root with lower or equal version than latest")
 	}
 

orphaning_tree.go

@@ -1,6 +1,8 @@
 package iavl
 
 import (
+	"fmt"
+
 	cmn "github.com/tendermint/tmlibs/common"
 )
 
@@ -46,10 +48,13 @@ func (tree *orphaningTree) unorphan(hash []byte) {
 // SaveAs saves the underlying Tree and assigns it a new version.
 // Saves orphans too.
 func (tree *orphaningTree) SaveAs(version int64) {
+	if version != tree.version+1 {
+		panic(fmt.Sprintf("Expected to save version %d but tried to save %d", tree.version+1, version))
+	}
 	if tree.root == nil {
 		// There can still be orphans, for example if the root is the node being
 		// removed.
-		tree.ndb.SaveOrphans(tree.version, tree.orphans)
+		tree.ndb.SaveOrphans(version, tree.orphans)
 		tree.ndb.SaveEmptyRoot(version)
 	} else {
 		// Save the current tree. For each saved node, we delete any existing
@@ -60,10 +65,11 @@ func (tree *orphaningTree) SaveAs(version int64) {
 		tree.ndb.SaveBranch(tree.root, func(node *Node) {
 			tree.unorphan(node._hash())
 		})
-		tree.ndb.SaveOrphans(tree.version, tree.orphans)
+		tree.ndb.SaveOrphans(version, tree.orphans)
 		tree.ndb.SaveRoot(tree.root, version)
 	}
 	tree.ndb.Commit()
+	tree.version = version
 }
 
 // Add orphans to the orphan list. Doesn't write to disk.

tree.go

@@ -49,6 +49,11 @@ func (t *Tree) Size() int {
 	return t.root.size
 }
 
+// Version returns the version of the tree.
+func (t *Tree) Version() int64 {
+	return t.version
+}
+
 // Height returns the height of the tree.
 func (t *Tree) Height() int8 {
 	if t.root == nil {
@@ -233,12 +238,13 @@ func (tree *Tree) clone() *Tree {
 
 // Load the tree from disk, from the given root hash, including all orphans.
 // Used internally by VersionedTree.
-func (tree *Tree) load(root []byte) {
+func (tree *Tree) load(version int64, root []byte) {
 	if len(root) == 0 {
 		tree.root = nil
 		return
 	}
 	tree.root = tree.ndb.GetNode(root)
+	tree.version = version
 }
 
 // nodeSize is like Size, but includes inner nodes too.

tree_test.go

@@ -224,7 +224,7 @@ func TestVersionedEmptyTree(t *testing.T) {
 	require.EqualValues(4, v)
 	require.NoError(err)
 
-	require.EqualValues(4, tree.LatestVersion())
+	require.EqualValues(4, tree.Version())
 
 	require.True(tree.VersionExists(1))
 	require.True(tree.VersionExists(3))
@@ -303,7 +303,7 @@ func TestVersionedTree(t *testing.T) {
 	require.NoError(tree.Load())
 
 	require.Len(tree.versions, 2, "wrong number of versions")
-	require.EqualValues(v2, tree.LatestVersion())
+	require.EqualValues(v2, tree.Version())
 
 	// -----1-----
 	// key1 = val0  <orphaned>
@@ -613,14 +613,14 @@ func TestVersionedTreeSaveAndLoad(t *testing.T) {
 	preHash := tree.Hash()
 	require.NotNil(preHash)
 
-	require.Equal(int64(6), tree.LatestVersion())
+	require.Equal(int64(6), tree.Version())
 
 	// Reload the tree, to test that roots and orphans are properly loaded.
 	ntree := NewVersionedTree(0, d)
 	ntree.Load()
 
 	require.False(ntree.IsEmpty())
-	require.Equal(int64(6), ntree.LatestVersion())
+	require.Equal(int64(6), ntree.Version())
 
 	postHash := ntree.Hash()
 	require.Equal(preHash, postHash)
@@ -1042,7 +1042,7 @@ func TestCopyValueSemantics(t *testing.T) {
 	require.Equal([]byte("v2"), val)
 }
 
-func TestResetToLatest(t *testing.T) {
+func TestRollback(t *testing.T) {
 	require := require.New(t)
 
 	tree := NewVersionedTree(0, db.NewMemDB())
@@ -1053,7 +1053,7 @@ func TestResetToLatest(t *testing.T) {
 	tree.Set([]byte("r"), []byte("v"))
 	tree.Set([]byte("s"), []byte("v"))
 
-	tree.ResetToLatest()
+	tree.Rollback()
 
 	tree.Set([]byte("t"), []byte("v"))
 
@@ -1071,42 +1071,6 @@ func TestResetToLatest(t *testing.T) {
 	require.Equal([]byte("v"), val)
 }
 
-func TestLoadVersion(t *testing.T) {
-	require := require.New(t)
-
-	d := db.NewMemDB()
-	tree := NewVersionedTree(0, d)
-
-	tree.Set([]byte("k"), []byte("v"))
-	tree.SaveVersion()
-
-	tree.Set([]byte("r"), []byte("v"))
-	tree.SaveVersion()
-
-	tree.Set([]byte("s"), []byte("v"))
-	tree.SaveVersion()
-
-	tree.Set([]byte("k"), []byte("u"))
-	tree.SaveVersion()
-
-	tree = NewVersionedTree(0, d)
-	tree.LoadVersion(2)
-
-	require.EqualValues(2, tree.LatestVersion())
-	require.Len(tree.versions, 2)
-
-	_, val := tree.Get([]byte("r"))
-	require.Equal([]byte("v"), val)
-
-	_, val = tree.Get([]byte("k"))
-	require.Equal([]byte("v"), val)
-
-	_, val = tree.Get([]byte("s"))
-	require.Nil(val)
-
-	require.Error(tree.LoadVersion(5))
-}
-
 //////////////////////////// BENCHMARKS ///////////////////////////////////////
 
 func BenchmarkTreeLoadAndDelete(b *testing.B) {

versioned_tree.go

@@ -28,11 +28,6 @@ func NewVersionedTree(cacheSize int, db dbm.DB) *VersionedTree {
 	}
 }
 
-// LatestVersion returns the latest saved version of the tree.
-func (tree *VersionedTree) LatestVersion() int64 {
-	return tree.orphaningTree.version
-}
-
 // IsEmpty returns whether or not the tree has any keys. Only trees that are
 // not empty can be saved.
 func (tree *VersionedTree) IsEmpty() bool {
@@ -74,35 +69,6 @@ func (tree *VersionedTree) Remove(key []byte) ([]byte, bool) {
 	return tree.orphaningTree.Remove(key)
 }
 
-// LoadVersion loads the given version and all versions prior.
-func (tree *VersionedTree) LoadVersion(version int64) error {
-	roots, err := tree.ndb.getRoots()
-	if err != nil {
-		return err
-	}
-	if len(roots) == 0 {
-		return nil
-	}
-	if _, ok := roots[version]; !ok {
-		return errors.WithStack(ErrVersionDoesNotExist)
-	}
-
-	// Load all roots from the database up to the given version.
-	for v, root := range roots {
-		if v > version {
-			continue
-		}
-		t := &Tree{ndb: tree.ndb, version: v}
-		t.load(root)
-
-		tree.versions[v] = t
-	}
-	tree.version = version
-	tree.ResetToLatest()
-
-	return nil
-}
-
 // Load a versioned tree from disk. All tree versions are loaded automatically.
 func (tree *VersionedTree) Load() error {
 	roots, err := tree.ndb.getRoots()
@@ -117,7 +83,7 @@ func (tree *VersionedTree) Load() error {
 	latestVersion := int64(0)
 	for version, root := range roots {
 		t := &Tree{ndb: tree.ndb, version: version}
-		t.load(root)
+		t.load(version, root)
 
 		tree.versions[version] = t
 
@@ -133,9 +99,9 @@ func (tree *VersionedTree) Load() error {
 	return nil
 }
 
-// ResetToLatest resets the working tree to the latest saved version, discarding
+// Rollback resets the working tree to the latest saved version, discarding
 // any unsaved modifications.
-func (tree *VersionedTree) ResetToLatest() {
+func (tree *VersionedTree) Rollback() {
 	if tree.version > 0 {
 		tree.orphaningTree = newOrphaningTree(
 			tree.versions[tree.version].clone(),
@@ -164,18 +130,14 @@ func (tree *VersionedTree) SaveVersion() ([]byte, int64, error) {
 		return nil, version, errors.Errorf("version %d was already saved", version)
 	}
 
-	tree.version = version // Sets the inner *Tree version.
+	// Persist version and stash to .versions.
+	tree.orphaningTree.SaveAs(version)
 	tree.versions[version] = tree.orphaningTree.Tree
 
-	tree.orphaningTree.SaveAs(version)
-	tree.orphaningTree = newOrphaningTree(
-		tree.versions[version].clone(),
-	)
+	// Set new working tree.
+	tree.orphaningTree = newOrphaningTree(tree.orphaningTree.clone())
 
-	if tree.root == nil {
-		return nil, version, nil
-	}
-	return tree.root.hash, version, nil
+	return tree.Hash(), version, nil
 }
 
 // DeleteVersion deletes a tree version from disk. The version can then no