fix: Make GetVersioned() available after doing LazyLoadVersion

mutable_tree.go

@@ -27,6 +27,7 @@ type MutableTree struct {
 	lastSaved      *ImmutableTree   // The most recently saved tree.
 	orphans        map[string]int64 // Nodes removed by changes to working tree.
 	versions       map[int64]bool   // The previous, saved versions of the tree.
+	allRootLoaded  bool             // Whether all roots are loaded or not(by LazyLoadVersion)
 	ndb            *nodeDB
 }
 
@@ -45,6 +46,7 @@ func NewMutableTreeWithOpts(db dbm.DB, cacheSize int, opts *Options) (*MutableTr
 		lastSaved:     head.clone(),
 		orphans:       map[string]int64{},
 		versions:      map[int64]bool{},
+		allRootLoaded: false,
 		ndb:           ndb,
 	}, nil
 }
@@ -57,7 +59,17 @@ func (tree *MutableTree) IsEmpty() bool {
 
 // VersionExists returns whether or not a version exists.
 func (tree *MutableTree) VersionExists(version int64) bool {
-	return tree.versions[version]
+	if tree.allRootLoaded {
+		return tree.versions[version]
+	}
+
+	has, ok := tree.versions[version]
+	if ok {
+		return has
+	}
+	has, _ = tree.ndb.HasRoot(version)
+	tree.versions[version] = has
+	return has
 }
 
 // AvailableVersions returns all available versions in ascending order
@@ -311,7 +323,11 @@ func (tree *MutableTree) LazyLoadVersion(targetVersion int64) (int64, error) {
 	iTree := &ImmutableTree{
 		ndb:     tree.ndb,
 		version: targetVersion,
-		root:    tree.ndb.GetNode(rootHash),
+	}
+	if len(rootHash) > 0 {
+		// If rootHash is empty then root of tree should be nil
+		// This makes `LazyLoadVersion` to do the same thing as `LoadVersion`
+		iTree.root = tree.ndb.GetNode(rootHash)
 	}
 
 	tree.orphans = map[string]int64{}
@@ -372,6 +388,7 @@ func (tree *MutableTree) LoadVersion(targetVersion int64) (int64, error) {
 	tree.orphans = map[string]int64{}
 	tree.ImmutableTree = t
 	tree.lastSaved = t.clone()
+	tree.allRootLoaded = true
 
 	return latestVersion, nil
 }
@@ -413,11 +430,13 @@ func (tree *MutableTree) GetImmutable(version int64) (*ImmutableTree, error) {
 	if rootHash == nil {
 		return nil, ErrVersionDoesNotExist
 	} else if len(rootHash) == 0 {
+		tree.versions[version] = true
 		return &ImmutableTree{
 			ndb:     tree.ndb,
 			version: version,
 		}, nil
 	}
+	tree.versions[version] = true
 	return &ImmutableTree{
 		root:    tree.ndb.GetNode(rootHash),
 		ndb:     tree.ndb,
@@ -441,7 +460,7 @@ func (tree *MutableTree) Rollback() {
 func (tree *MutableTree) GetVersioned(key []byte, version int64) (
 	index int64, value []byte,
 ) {
-	if tree.versions[version] {
+	if tree.VersionExists(version) {
 		t, err := tree.GetImmutable(version)
 		if err != nil {
 			return -1, nil
@@ -459,7 +478,7 @@ func (tree *MutableTree) SaveVersion() ([]byte, int64, error) {
 		version = int64(tree.ndb.opts.InitialVersion)
 	}
 
-	if tree.versions[version] {
+	if tree.VersionExists(version) {
 		// If the version already exists, return an error as we're attempting to overwrite.
 		// However, the same hash means idempotent (i.e. no-op).
 		existingHash, err := tree.ndb.getRoot(version)
@@ -525,10 +544,9 @@ func (tree *MutableTree) deleteVersion(version int64) error {
 	if version == tree.version {
 		return errors.Errorf("cannot delete latest saved version (%d)", version)
 	}
-	if _, ok := tree.versions[version]; !ok {
+	if !tree.VersionExists(version) {
 		return errors.Wrap(ErrVersionDoesNotExist, "")
 	}
-
 	if err := tree.ndb.DeleteVersion(version, true); err != nil {
 		return err
 	}

mutable_tree_test.go

@@ -288,3 +288,95 @@ func BenchmarkMutableTree_Set(b *testing.B) {
 		t.Set(randBytes(10), []byte{})
 	}
 }
+
+func prepareTree(t *testing.T) *MutableTree {
+	mdb := db.NewMemDB()
+	tree, err := NewMutableTree(mdb, 1000)
+	require.NoError(t, err)
+	for i := 0; i < 100; i++ {
+		tree.Set([]byte{byte(i)}, []byte("a"))
+	}
+	_, ver, err := tree.SaveVersion()
+	require.True(t, ver == 1)
+	require.NoError(t, err)
+	for i := 0; i < 100; i++ {
+		tree.Set([]byte{byte(i)}, []byte("b"))
+	}
+	_, ver, err = tree.SaveVersion()
+	require.True(t, ver == 2)
+	require.NoError(t, err)
+	newTree, err := NewMutableTree(mdb, 1000)
+	require.NoError(t, err)
+
+	return newTree
+}
+
+func TestMutableTree_VersionExists(t *testing.T) {
+	tree := prepareTree(t)
+	require.True(t, tree.VersionExists(1))
+	require.True(t, tree.VersionExists(2))
+	require.False(t, tree.VersionExists(3))
+}
+
+func checkGetVersioned(t *testing.T, tree *MutableTree, version, index int64, key, value []byte) {
+	idx, val := tree.GetVersioned(key, version)
+	require.True(t, idx == index)
+	require.True(t, bytes.Equal(val, value))
+}
+
+func TestMutableTree_GetVersioned(t *testing.T) {
+	tree := prepareTree(t)
+	ver, err := tree.LazyLoadVersion(1)
+	require.True(t, ver == 1)
+	require.NoError(t, err)
+	// check key of unloaded version
+	checkGetVersioned(t, tree, 1, 1, []byte{1}, []byte("a"))
+	checkGetVersioned(t, tree, 2, 1, []byte{1}, []byte("b"))
+	checkGetVersioned(t, tree, 3, -1, []byte{1}, nil)
+
+	tree = prepareTree(t)
+	ver, err = tree.LazyLoadVersion(2)
+	require.True(t, ver == 2)
+	require.NoError(t, err)
+	checkGetVersioned(t, tree, 1, 1, []byte{1}, []byte("a"))
+	checkGetVersioned(t, tree, 2, 1, []byte{1}, []byte("b"))
+	checkGetVersioned(t, tree, 3, -1, []byte{1}, nil)
+}
+
+func TestMutableTree_DeleteVersion(t *testing.T) {
+	tree := prepareTree(t)
+	ver, err := tree.LazyLoadVersion(2)
+	require.True(t, ver == 2)
+	require.NoError(t, err)
+
+	require.NoError(t, tree.DeleteVersion(1))
+
+	require.False(t, tree.VersionExists(1))
+	require.True(t, tree.VersionExists(2))
+	require.False(t, tree.VersionExists(3))
+
+	// cannot delete latest version
+	require.Error(t, tree.DeleteVersion(2))
+}
+
+func TestMutableTree_LazyLoadVersionWithEmptyTree(t *testing.T) {
+	mdb := db.NewMemDB()
+	tree, err := NewMutableTree(mdb, 1000)
+	require.NoError(t, err)
+	_, v1, err := tree.SaveVersion()
+	require.NoError(t, err)
+
+	newTree1, err := NewMutableTree(mdb, 1000)
+	require.NoError(t, err)
+	v2, err := newTree1.LazyLoadVersion(1)
+	require.NoError(t, err)
+	require.True(t, v1 == v2)
+
+	newTree2, err := NewMutableTree(mdb, 1000)
+	require.NoError(t, err)
+	v2, err = newTree1.LoadVersion(1)
+	require.NoError(t, err)
+	require.True(t, v1 == v2)
+
+	require.True(t, newTree1.root == newTree2.root)
+}

nodedb.go

@@ -543,6 +543,10 @@ func (ndb *nodeDB) Commit() error {
 	return nil
 }
 
+func (ndb *nodeDB) HasRoot(version int64) (bool, error) {
+	return ndb.db.Has(ndb.rootKey(version))
+}
+
 func (ndb *nodeDB) getRoot(version int64) ([]byte, error) {
 	return ndb.db.Get(ndb.rootKey(version))
 }

proof_range.go

@@ -541,7 +541,7 @@ func (t *ImmutableTree) GetRangeWithProof(startKey []byte, endKey []byte, limit
 // GetVersionedWithProof gets the value under the key at the specified version
 // if it exists, or returns nil.
 func (tree *MutableTree) GetVersionedWithProof(key []byte, version int64) ([]byte, *RangeProof, error) {
-	if tree.versions[version] {
+	if tree.VersionExists(version) {
 		t, err := tree.GetImmutable(version)
 		if err != nil {
 			return nil, nil, err
@@ -557,7 +557,7 @@ func (tree *MutableTree) GetVersionedWithProof(key []byte, version int64) ([]byt
 func (tree *MutableTree) GetVersionedRangeWithProof(startKey, endKey []byte, limit int, version int64) (
 	keys, values [][]byte, proof *RangeProof, err error) {
 
-	if tree.versions[version] {
+	if tree.VersionExists(version) {
 		t, err := tree.GetImmutable(version)
 		if err != nil {
 			return nil, nil, nil, err