From 1c8b8e787e85f78c5dfcdfb83339bd1e0b36939f Mon Sep 17 00:00:00 2001
From: cool-developer <51834436+cool-develope@users.noreply.github.com>
Date: Tue, 7 May 2024 14:25:44 -0400
Subject: [PATCH] fix: working hash with initial version (#943)

---
 mutable_tree.go |  2 +-
 tree_test.go    | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/mutable_tree.go b/mutable_tree.go
index 740a63533..7c73dee52 100644
--- a/mutable_tree.go
+++ b/mutable_tree.go
@@ -140,7 +140,7 @@ func (tree *MutableTree) Hash() []byte {
 
 // WorkingHash returns the hash of the current working tree.
 func (tree *MutableTree) WorkingHash() []byte {
-	return tree.ImmutableTree.Hash()
+	return tree.root.hashWithCount(tree.WorkingVersion())
 }
 
 func (tree *MutableTree) WorkingVersion() int64 {
diff --git a/tree_test.go b/tree_test.go
index ca31acafb..75d3b7b24 100644
--- a/tree_test.go
+++ b/tree_test.go
@@ -1920,3 +1920,37 @@ func TestReferenceRoot(t *testing.T) {
 	_, err = tree.Set([]byte("key1"), []byte("value2"))
 	require.NoError(t, err)
 }
+
+func TestWorkingHashWithInitialVersion(t *testing.T) {
+	db, err := dbm.NewDB("test", "memdb", "")
+	require.NoError(t, err)
+	defer db.Close()
+
+	initialVersion := int64(100)
+	tree := NewMutableTree(db, 0, false, log.NewNopLogger())
+	tree.SetInitialVersion(uint64(initialVersion))
+
+	v := tree.WorkingVersion()
+	require.Equal(t, initialVersion, v)
+
+	_, err = tree.Set([]byte("key1"), []byte("value1"))
+	require.NoError(t, err)
+
+	workingHash := tree.WorkingHash()
+	commitHash, _, err := tree.SaveVersion()
+	require.NoError(t, err)
+	require.Equal(t, commitHash, workingHash)
+
+	db, err = dbm.NewDB("test", "memdb", "")
+	require.NoError(t, err)
+
+	// without WorkingHash
+	tree = NewMutableTree(db, 0, false, log.NewNopLogger(), InitialVersionOption(uint64(initialVersion)))
+
+	_, err = tree.Set([]byte("key1"), []byte("value1"))
+	require.NoError(t, err)
+
+	commitHash1, _, err := tree.SaveVersion()
+	require.NoError(t, err)
+	require.Equal(t, commitHash1, commitHash)
+}