From de4694eb2b4953cdff4a4b697d743b151374b933 Mon Sep 17 00:00:00 2001
From: Krish <krish.z@nodereal.io>
Date: Fri, 22 Sep 2023 10:33:35 +0800
Subject: [PATCH] fix: use cache first when get StorageTrie/Trie

---
 core/state/database.go | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/core/state/database.go b/core/state/database.go
index da26719816..5b52682d13 100644
--- a/core/state/database.go
+++ b/core/state/database.go
@@ -238,19 +238,57 @@ func (db *cachingDB) purgeLoop() {
 
 // OpenTrie opens the main account trie at a specific root hash.
 func (db *cachingDB) OpenTrie(root common.Hash) (Trie, error) {
+	//try cache first
+	if db.accountTrieCache != nil {
+		if tr, ok := db.accountTrieCache.Get(root); ok {
+			return tr.(*trie.SecureTrie), nil
+		}
+	}
+
 	tr, err := trie.NewStateTrie(trie.StateTrieID(root), db.triedb)
 	if err != nil {
 		return nil, err
 	}
+
+	//update LRU
+	if db.accountTrieCache != nil {
+		db.accountTrieCache.Add(root, tr.ResetCopy())
+	}
 	return tr, nil
 }
 
 // OpenStorageTrie opens the storage trie of an account.
 func (db *cachingDB) OpenStorageTrie(stateRoot common.Hash, addrHash, root common.Hash) (Trie, error) {
+	// try cache first
+	if db.storageTrieCache != nil {
+		if tries, exist := db.storageTrieCache.Get(addrHash); exist {
+			for _, triePair := range tries {
+				if triePair.root == root {
+					return triePair.trie, nil
+				}
+			}
+		}
+	}
+
 	tr, err := trie.NewStateTrie(trie.StorageTrieID(stateRoot, addrHash, root), db.triedb)
 	if err != nil {
 		return nil, err
 	}
+
+	if db.storageTrieCache != nil {
+		tries, exist := db.storageTrieCache.Get(addrHash)
+		var newTriesArray [3]*triePair
+		if exist {
+			newTriesArray = [3]*triePair{
+				{root: root, trie: tr.ResetCopy()},
+				tries[0],
+				tries[1],
+			}
+		} else {
+			newTriesArray = [3]*triePair{{root: root, trie: tr.ResetCopy()}, nil, nil}
+		}
+		db.storageTrieCache.Add(addrHash, newTriesArray)
+	}
 	return tr, nil
 }