fix: Add refnode at InitialVersion pointing at version 1 root

.github/workflows/ci.yml

@@ -3,6 +3,7 @@ on:
   push:
     branches:
       - master
+      - release/*
   pull_request:
 
 jobs:

mutable_tree.go

@@ -137,7 +137,9 @@ func (tree *MutableTree) Hash() []byte {
 
 // WorkingHash returns the hash of the current working tree.
 func (tree *MutableTree) WorkingHash() []byte {
-	return tree.root.hashWithCount(tree.WorkingVersion())
+	// Use tree.version instead of tree.WorkingVersion() because tree.version
+	// is always the latest version
+	return tree.root.hashWithCount(tree.version + 1)
 }
 
 func (tree *MutableTree) WorkingVersion() int64 {
@@ -746,6 +748,7 @@ func (tree *MutableTree) SaveVersion() ([]byte, int64, error) {
 			return nil, 0, err
 		}
 	} else {
+		// nodeKey already set, no changes to tree
 		if tree.root.nodeKey != nil {
 			// it means there are no updated nodes
 			if err := tree.ndb.SaveRoot(version, tree.root.nodeKey); err != nil {
@@ -761,12 +764,28 @@ func (tree *MutableTree) SaveVersion() ([]byte, int64, error) {
 				}
 			}
 		} else {
+			// No nodeKey, assign new nodeKeys and save nodes. This happens when
+			// any node in the tree has been updated as setting a key will
+			// reset all nodeKeys to nil
 			if err := tree.saveNewNodes(); err != nil {
 				return nil, 0, err
 			}
 		}
 	}
 
+	// Addition to node version patch for backwards compatibility:
+	// This is to resolve queries at InitialVersion, loading the tree at
+	// InitialVersion when the root node has version 1.
+	// Note: version != nodekey.Version due to patch.
+	if tree.ndb.opts.InitialVersion > 0 && version == int64(tree.ndb.opts.InitialVersion) {
+		// SaveRoot is meant for saving a tree when there are no updates, which
+		// simply creates a reference node to the root node. We reuse this to
+		// create a reference node from InitialVersion -> 1 (root node).
+		if err := tree.ndb.SaveRoot(version, tree.root.nodeKey); err != nil {
+			return nil, 0, err
+		}
+	}
+
 	if err := tree.ndb.Commit(); err != nil {
 		return nil, version, err
 	}

mutable_tree_test.go

@@ -1447,37 +1447,125 @@ func TestMutableTree_InitialVersion_FirstVersion(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, initialVersion, version)
 
+	err = tree.ndb.batch.WriteSync()
+	require.NoError(t, err)
+
+	// Save 2 empty versions after the initial version, should produce 2 refnodes
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	_, latestVersion, err := tree.SaveVersion()
+	require.NoError(t, err)
+
+	// ------------------------------
+	// Verify
+	// 1. Node exists with Version 1, not InitialVersion
+	// 2. VersionExists(InitialVersion) - should return true
+	// 3. AvailableVersions() - should include InitialVersion and NOT version 1
+	// 4. ndb.GetRoot(InitialVersion) - should return the root node
+	// 5. GetImmutable(InitialVersion) - should return the immutable tree
+	// 6. LoadVersion(InitialVersion) - loads tree even if node contains version 1. This is last since it runs some previous methods internally
+
 	// the nodes created at the first version are not assigned with the `InitialVersion`
-	rootKey := GetRootKey(1)
-	node, err := tree.ndb.GetNode(rootKey)
+	t.Run("1. node exists with root key version 1", func(t *testing.T) {
+		rootKey := GetRootKey(1)
+		node, err := tree.ndb.GetNode(rootKey)
+		require.NoError(t, err)
+		require.Equal(t, int64(1), node.nodeKey.version, "new nodes on new tree should be version 1")
+
+		// Check fast node version
+		fastNode, err := tree.ndb.GetFastNode([]byte("hello"))
+		require.NoError(t, err)
+		require.Equal(t, int64(1), fastNode.GetVersionLastUpdatedAt(), "fast nodes be version 1")
+	})
+
+	t.Run("2. VersionExists(InitialVersion)", func(t *testing.T) {
+		hasVersion := tree.VersionExists(initialVersion)
+		require.True(t, hasVersion, "initial version should still be found")
+	})
+
+	t.Run("3. AvailableVersions()", func(t *testing.T) {
+		versions := tree.AvailableVersions()
+		require.Contains(t, versions, int(initialVersion), "initial version should be found")
+		require.NotContains(t, versions, int64(1), "version 1 should not be found")
+	})
+
+	t.Run("4. GetRoot(InitialVersion)", func(t *testing.T) {
+		root, err := tree.ndb.GetRoot(initialVersion)
+
+		// Internally, this should use a refnode
+		require.NoError(t, err)
+		require.NotNil(t, root, "root node should be found at initial version")
+	})
+
+	t.Run("5. GetImmutable(InitialVersion)", func(t *testing.T) {
+		immutableTree, err := tree.GetImmutable(initialVersion)
+		require.NoError(t, err)
+		require.NotNil(t, immutableTree, "immutable tree should be found at initial version")
+	})
+
+	t.Run("6. LoadVersion(InitialVersion)", func(t *testing.T) {
+		// Final check - runs previous methods internally as well
+		_, err = tree.LoadVersion(initialVersion)
+		require.NoError(t, err)
+	})
+
+	// Internal ndb methods for good measure
+	firstVersion, err := tree.ndb.getFirstVersion()
 	require.NoError(t, err)
-	require.Equal(t, int64(1), node.nodeKey.version, "new nodes on new tree should be version 1")
+	require.Equal(t, initialVersion, firstVersion, "first version should be the initialVersion")
 
-	// Check fast node version
-	fastNode, err := tree.ndb.GetFastNode([]byte("hello"))
+	latest, err := tree.ndb.getLatestVersion()
+	require.NoError(t, err)
+	require.Equal(t, latestVersion, latest, "latest version should be updated")
+
+	// Check reference nodes point at the correct node key
+	err = tree.ndb.traversePrefix(nodeKeyFormat.Prefix(), func(key, value []byte) error {
+		if isRef, _ := isReferenceRoot(value); isRef {
+			// key still contains prefix, so we need to remove first byte
+			refnodeKey := GetNodeKey(key[1:])
+
+			nk := GetNodeKey(value[1:])
+
+			t.Logf("ref node %s -> %s", refnodeKey.String(), nk.String())
+
+			val, err := tree.ndb.GetNode(nk.GetKey())
+			require.NoError(t, err)
+			require.NotNil(t, val, "reference node should point to a valid node")
+		}
+
+		return nil
+	})
 	require.NoError(t, err)
-	require.Equal(t, int64(1), fastNode.GetVersionLastUpdatedAt(), "fast nodes be version 1")
 
 	// ------------------------------
 	// Writes on existing tree
 
+	// Load the latest version again for writing
+	_, err = tree.LoadVersion(latestVersion)
+	require.NoError(t, err)
+
 	_, err = tree.Set([]byte("hello"), []byte("world1"))
 	require.NoError(t, err)
 
 	_, version, err = tree.SaveVersion()
 	require.NoError(t, err)
-	require.Equal(t, initialVersion+1, version, "new version should be initialVersion+1")
+	require.Equal(t, latestVersion+1, version, "new version should be initialVersion+1")
 
-	rootKey = GetRootKey(version)
+	rootKey := GetRootKey(version)
 	// the following versions behaves normally
-	node, err = tree.ndb.GetNode(rootKey)
+	node, err := tree.ndb.GetNode(rootKey)
 	require.NoError(t, err)
-	require.Equal(t, initialVersion+1, node.nodeKey.version, "new nodes on existing tree should use initialVersion")
+	require.Equal(t, latestVersion+1, node.nodeKey.version, "new nodes on existing tree should use initialVersion")
 
 	// Check fast node version
-	fastNode, err = tree.ndb.GetFastNode([]byte("hello"))
+	fastNode, err := tree.ndb.GetFastNode([]byte("hello"))
 	require.NoError(t, err)
-	require.Equal(t, initialVersion+1, fastNode.GetVersionLastUpdatedAt(), "fast nodes should have the same version as the tree")
+	require.Equal(
+		t, latestVersion+1,
+		fastNode.GetVersionLastUpdatedAt(),
+		"fast nodes should have the same version as the tree",
+	)
 }
 
 func TestMutableTreeClose(t *testing.T) {

tree_test.go

@@ -1936,6 +1936,9 @@ func TestWorkingHashWithInitialVersion(t *testing.T) {
 	_, err = tree.Set([]byte("key1"), []byte("value1"))
 	require.NoError(t, err)
 
+	// WorkingHash sets the hashes in the nodes, SaveVersion doesn't recompute it.
+	// This test ensures WorkingHash and SaveVersion are consistent, using the
+	// same version number.
 	workingHash := tree.WorkingHash()
 	commitHash, _, err := tree.SaveVersion()
 	require.NoError(t, err)
@@ -1952,5 +1955,5 @@ func TestWorkingHashWithInitialVersion(t *testing.T) {
 
 	commitHash1, _, err := tree.SaveVersion()
 	require.NoError(t, err)
-	require.Equal(t, commitHash1, commitHash)
+	assert.Equal(t, hex.EncodeToString(commitHash1), hex.EncodeToString(commitHash))
 }