Pruned node fix

mutable_tree.go

@@ -1035,10 +1035,7 @@ func (tree *MutableTree) saveNewNodes(version int64) error {
 	var recursiveAssignKey func(*Node) ([]byte, error)
 	recursiveAssignKey = func(node *Node) ([]byte, error) {
 		if node.nodeKey != nil {
-			if node.nodeKey.nonce != 0 {
-				return node.nodeKey.GetKey(), nil
-			}
-			return node.hash, nil
+			return node.GetKey(), nil
 		}
 		nonce++
 		node.nodeKey = &NodeKey{

tree_test.go

@@ -1876,12 +1876,47 @@ func TestReferenceRoot(t *testing.T) {
 	_, _, err = tree.SaveVersion()
 	require.NoError(t, err)
 
-	// Load the tree from disk.
+	// load the tree from disk
 	tree = NewMutableTree(db, 0, false, log.NewNopLogger())
 	_, err = tree.Load()
 	require.NoError(t, err)
 	require.Equal(t, int64(2), tree.Version())
 	// check the root of version 2 is the leaf node of key2
 	require.Equal(t, tree.root.GetKey(), (&NodeKey{version: 1, nonce: 3}).GetKey())
 	require.Equal(t, tree.root.key, []byte("key2"))
+
+	// test the reference root when pruning
+	db, err = dbm.NewDB("test", "memdb", "")
+	require.NoError(t, err)
+	tree = NewMutableTree(db, 0, false, log.NewNopLogger())
+
+	_, err = tree.Set([]byte("key1"), []byte("value1"))
+	require.NoError(t, err)
+
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	_, _, err = tree.SaveVersion() // empty version
+	require.NoError(t, err)
+
+	require.NoError(t, tree.DeleteVersionsTo(1))
+	_, _, err = tree.SaveVersion() // empty version
+	require.NoError(t, err)
+
+	// load the tree from disk
+	tree = NewMutableTree(db, 0, false, log.NewNopLogger())
+	_, err = tree.Load()
+	require.NoError(t, err)
+
+	_, err = tree.Set([]byte("key2"), []byte("value2"))
+	require.NoError(t, err)
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	// load the tree from disk to check if the reference root is loaded correctly
+	tree = NewMutableTree(db, 0, false, log.NewNopLogger())
+	_, err = tree.Load()
+	require.NoError(t, err)
+	_, err = tree.Set([]byte("key1"), []byte("value2"))
+	require.NoError(t, err)
 }