From 95927ab58968bf14f03777362eb2fde223c95e26 Mon Sep 17 00:00:00 2001
From: Logan <80692353+jyc228@users.noreply.github.com>
Date: Thu, 11 Apr 2024 15:48:55 +0900
Subject: [PATCH] fix(trie): modify to prevent errors when updating with
 non-existent keys (#96)

---
 trie/merkle_trie.go          |  4 ++++
 trie/zk_merkle_state_trie.go | 15 +++++++++++++--
 trie/zk_merkle_trie_test.go  | 19 +++++++++++++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/trie/merkle_trie.go b/trie/merkle_trie.go
index 46ca9020b..a9c46f880 100644
--- a/trie/merkle_trie.go
+++ b/trie/merkle_trie.go
@@ -74,9 +74,13 @@ type MerkleStateTrie interface {
 	Hash() common.Hash
 	GetNode(path []byte) ([]byte, int, error)
 	MustGet(key []byte) []byte
+	GetAccount(address common.Address) (*types.StateAccount, error)
 	GetAccountByHash(addrHash common.Hash) (*types.StateAccount, error)
 	UpdateAccount(address common.Address, account *types.StateAccount) error
+	DeleteAccount(address common.Address) error
 	MustUpdate(key, value []byte)
+	GetStorage(addr common.Address, key []byte) ([]byte, error)
+	DeleteStorage(address common.Address, key []byte) error
 	MustDelete(key []byte)
 	MustNodeIterator(start []byte) NodeIterator
 	Commit(collectLeaf bool) (common.Hash, *trienode.NodeSet, error)
diff --git a/trie/zk_merkle_state_trie.go b/trie/zk_merkle_state_trie.go
index 00e665260..67f029c8b 100644
--- a/trie/zk_merkle_state_trie.go
+++ b/trie/zk_merkle_state_trie.go
@@ -1,6 +1,9 @@
 package trie
 
 import (
+	"errors"
+
+	"github.com/kroma-network/zktrie/trie"
 	zkt "github.com/kroma-network/zktrie/types"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -92,8 +95,16 @@ func (z *ZkMerkleStateTrie) UpdateContractCode(_ common.Address, _ common.Hash,
 	return nil
 }
 
-func (z *ZkMerkleStateTrie) DeleteStorage(_ common.Address, key []byte) error { return z.Delete(key) }
-func (z *ZkMerkleStateTrie) DeleteAccount(address common.Address) error       { return z.Delete(address[:]) }
+func (z *ZkMerkleStateTrie) DeleteStorage(_ common.Address, key []byte) error { return z.delete(key) }
+func (z *ZkMerkleStateTrie) DeleteAccount(address common.Address) error       { return z.delete(address[:]) }
+
+func (z *ZkMerkleStateTrie) delete(key []byte) error {
+	if err := z.Delete(key); err != nil && errors.Is(err, trie.ErrKeyNotFound) {
+		return nil
+	} else {
+		return err
+	}
+}
 
 func (z *ZkMerkleStateTrie) Prove(key []byte, proofDb ethdb.KeyValueWriter) error {
 	return z.prove(common.ReverseBytes(key), proofDb, func(node zk.TreeNode) error {
diff --git a/trie/zk_merkle_trie_test.go b/trie/zk_merkle_trie_test.go
index 52d48da80..34f91e9b6 100644
--- a/trie/zk_merkle_trie_test.go
+++ b/trie/zk_merkle_trie_test.go
@@ -14,6 +14,25 @@ import (
 	"github.com/ethereum/go-ethereum/trie/zk"
 )
 
+func TestZkTrieNoExistKey(t *testing.T) {
+	testMerkleStateTrieNoExistKey(t, newEmptyZkTrie())
+	testMerkleStateTrieNoExistKey(t, NewEmptyZkMerkleStateTrie(NewZkDatabase(rawdb.NewMemoryDatabase())))
+}
+
+func testMerkleStateTrieNoExistKey(t *testing.T, trie MerkleStateTrie) {
+	address := common.HexToAddress("0xffffffffffffffffffffffffffffffffffffffff")
+	if acc, err := trie.GetAccount(address); acc != nil || err != nil {
+		t.Errorf("GetAccount return acc : %v, err : %v", acc, err)
+	}
+	assert.NoError(t, trie.DeleteAccount(address))
+
+	storageKey := common.LeftPadBytes([]byte{1}, 20)
+	if val, err := trie.GetStorage(address, storageKey); len(val) > 0 || err != nil {
+		t.Errorf("GetStorage return acc : %v, err : %v", val, err)
+	}
+	assert.NoError(t, trie.DeleteStorage(address, storageKey))
+}
+
 func TestProve(t *testing.T) {
 	t.Run("without preimage", func(t *testing.T) {
 		scrolldb, kromadb := NewZkDatabase(rawdb.NewMemoryDatabase()), NewZkDatabase(rawdb.NewMemoryDatabase())