From 77a3b33a95841c351bdda133c9897f4ceea4ae6c Mon Sep 17 00:00:00 2001
From: Logan <logan@lightscale.io>
Date: Tue, 16 Jan 2024 17:56:02 +0900
Subject: [PATCH] fix(trie): prevent save root hash twice

---
 trie/iterator_test.go                 | 4 ++--
 trie/triedb/hashdb/zktrie_database.go | 3 ++-
 trie/zk_merkle_trie.go                | 7 +++----
 trie/zk_trie.go                       | 7 -------
 4 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/trie/iterator_test.go b/trie/iterator_test.go
index ba46711c3..a5373ba93 100644
--- a/trie/iterator_test.go
+++ b/trie/iterator_test.go
@@ -651,10 +651,10 @@ func TestMerkleTreeIterator(t *testing.T) {
 
 	t.Run("zk merkle tree", func(t *testing.T) {
 		tree, db := makeMerkleTreeWithData(testdata1)
-		expected := db.Len() - 1
+		expected := db.Len()
 		it, _ := tree.NodeIterator(nil)
 		count, leafCount := testIterator(t, db, it)
-		if db.Len()-1 != 0 {
+		if db.Len() != 0 {
 			t.Errorf("db is not empty. remain size %d", db.Len())
 		}
 		if expected != count {
diff --git a/trie/triedb/hashdb/zktrie_database.go b/trie/triedb/hashdb/zktrie_database.go
index 36a4d2e39..fe967759d 100644
--- a/trie/triedb/hashdb/zktrie_database.go
+++ b/trie/triedb/hashdb/zktrie_database.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/VictoriaMetrics/fastcache"
 	zktrie "github.com/kroma-network/zktrie/trie"
+	zkt "github.com/kroma-network/zktrie/types"
 	"github.com/syndtr/goleveldb/leveldb"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -47,7 +48,7 @@ func NewZk(diskdb ethdb.Database, config *Config) *ZktrieDatabase {
 func (db *ZktrieDatabase) Scheme() string { return rawdb.HashScheme }
 
 func (db *ZktrieDatabase) Initialized(genesisRoot common.Hash) bool {
-	return rawdb.HasLegacyTrieNode(db.diskdb, genesisRoot)
+	return rawdb.HasLegacyTrieNode(db.diskdb, common.BytesToHash(zkt.ReverseByteOrder(genesisRoot[:])))
 }
 
 func (db *ZktrieDatabase) Size() (common.StorageSize, common.StorageSize) {
diff --git a/trie/zk_merkle_trie.go b/trie/zk_merkle_trie.go
index 822fe4b8b..d870370d1 100644
--- a/trie/zk_merkle_trie.go
+++ b/trie/zk_merkle_trie.go
@@ -4,7 +4,6 @@ import (
 	zktrie "github.com/kroma-network/zktrie/trie"
 
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/trie/trienode"
@@ -40,13 +39,13 @@ func (z *ZkMerkleTrie) NodeIterator(startKey []byte) (NodeIterator, error) {
 }
 
 func (z *ZkMerkleTrie) Commit(_ bool) (common.Hash, *trienode.NodeSet, error) {
+	if root := z.RootNode().Hash(); root != nil {
+		return common.BytesToHash(root.Bytes()), nil, nil
+	}
 	err := z.ComputeAllNodeHash(func(node zk.TreeNode) error { return z.db.Put(node.Hash()[:], node.CanonicalValue()) })
 	if err != nil {
 		log.Error("Failed to commit zk merkle trie", "err", err)
 	}
-	// There is a bug where root node is saved twice.
-	// It is because of the bottom rawdb.WriteLegacyTrieNode, and we will remove it after checking if it can be removed.
-	rawdb.WriteLegacyTrieNode(z.db.diskdb, common.BytesToHash(z.RootNode().Hash().Bytes()), z.RootNode().CanonicalValue())
 	// Since NodeSet relies directly on mpt, we can't create a NodeSet.
 	// Of course, we might be able to force-fit it by implementing the node interface.
 	// However, NodeSet has been improved in geth, it could be improved to return a NodeSet when a commit is applied.
diff --git a/trie/zk_trie.go b/trie/zk_trie.go
index 47ea5190e..6e43a4c82 100644
--- a/trie/zk_trie.go
+++ b/trie/zk_trie.go
@@ -26,7 +26,6 @@ import (
 	zkt "github.com/kroma-network/zktrie/types"
 
 	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto/poseidon"
 	"github.com/ethereum/go-ethereum/ethdb"
@@ -171,12 +170,6 @@ func (t *ZkTrie) GetKey(kHashBytes []byte) []byte {
 func (t *ZkTrie) Commit(bool) (common.Hash, *trienode.NodeSet, error) {
 	// in current implementation, every update of trie already writes into database
 	// so Commit does nothing
-	node, err := t.Tree().GetNode(t.Tree().Root())
-	if err != nil {
-		return types.GetEmptyRootHash(true), nil, err
-	}
-
-	rawdb.WriteLegacyTrieNode(t.db.diskdb, t.Hash(), node.Value())
 	return t.Hash(), nil, nil
 }