From 999f09f8af4d3e8ea3547fc788389b37940d770e Mon Sep 17 00:00:00 2001
From: Delweng <delweng@gmail.com>
Date: Fri, 13 Jun 2025 15:04:24 +0800
Subject: [PATCH] trie: no need to store preimage if not enabled (#32012)

As the preimage will only be stored if `t.preimages != nil`, so no need
to save them into local cache if not enabled. This will reduce the memory
wasted to copy the bytes

---------

Signed-off-by: jsvisa <delweng@gmail.com>
---
 trie/database_test.go |  4 ++++
 trie/secure_trie.go   | 36 +++++++++++++++++++++++++-----------
 triedb/database.go    |  5 +++++
 3 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/trie/database_test.go b/trie/database_test.go
index 535f0d61b2..98fd2372ef 100644
--- a/trie/database_test.go
+++ b/trie/database_test.go
@@ -86,6 +86,10 @@ func (db *testDb) InsertPreimage(preimages map[common.Hash][]byte) {
 	rawdb.WritePreimages(db.disk, preimages)
 }
 
+func (db *testDb) PreimageEnabled() bool {
+	return true
+}
+
 func (db *testDb) Scheme() string { return db.scheme }
 
 func (db *testDb) Update(root common.Hash, parent common.Hash, nodes *trienode.MergedNodeSet) error {
diff --git a/trie/secure_trie.go b/trie/secure_trie.go
index 7852e16619..ac20f961e5 100644
--- a/trie/secure_trie.go
+++ b/trie/secure_trie.go
@@ -33,6 +33,9 @@ type preimageStore interface {
 
 	// InsertPreimage commits a set of preimages along with their hashes.
 	InsertPreimage(preimages map[common.Hash][]byte)
+
+	// PreimageEnabled returns true if the preimage store is enabled.
+	PreimageEnabled() bool
 }
 
 // SecureTrie is the old name of StateTrie.
@@ -84,8 +87,7 @@ func NewStateTrie(id *ID, db database.NodeDatabase) (*StateTrie, error) {
 	tr := &StateTrie{trie: *trie, db: db}
 
 	// link the preimage store if it's supported
-	preimages, ok := db.(preimageStore)
-	if ok {
+	if preimages, ok := db.(preimageStore); ok && preimages.PreimageEnabled() {
 		tr.preimages = preimages
 	}
 	return tr, nil
@@ -159,7 +161,9 @@ func (t *StateTrie) GetNode(path []byte) ([]byte, int, error) {
 func (t *StateTrie) MustUpdate(key, value []byte) {
 	hk := crypto.Keccak256(key)
 	t.trie.MustUpdate(hk, value)
-	t.getSecKeyCache()[common.Hash(hk)] = common.CopyBytes(key)
+	if t.preimages != nil {
+		t.getSecKeyCache()[common.Hash(hk)] = common.CopyBytes(key)
+	}
 }
 
 // UpdateStorage associates key with value in the trie. Subsequent calls to
@@ -177,7 +181,9 @@ func (t *StateTrie) UpdateStorage(_ common.Address, key, value []byte) error {
 	if err != nil {
 		return err
 	}
-	t.getSecKeyCache()[common.Hash(hk)] = common.CopyBytes(key)
+	if t.preimages != nil {
+		t.getSecKeyCache()[common.Hash(hk)] = common.CopyBytes(key)
+	}
 	return nil
 }
 
@@ -191,7 +197,9 @@ func (t *StateTrie) UpdateAccount(address common.Address, acc *types.StateAccoun
 	if err := t.trie.Update(hk, data); err != nil {
 		return err
 	}
-	t.getSecKeyCache()[common.Hash(hk)] = address.Bytes()
+	if t.preimages != nil {
+		t.getSecKeyCache()[common.Hash(hk)] = address.Bytes()
+	}
 	return nil
 }
 
@@ -203,7 +211,9 @@ func (t *StateTrie) UpdateContractCode(_ common.Address, _ common.Hash, _ []byte
 // will omit any encountered error but just print out an error message.
 func (t *StateTrie) MustDelete(key []byte) {
 	hk := crypto.Keccak256(key)
-	delete(t.getSecKeyCache(), common.Hash(hk))
+	if t.preimages != nil {
+		delete(t.getSecKeyCache(), common.Hash(hk))
+	}
 	t.trie.MustDelete(hk)
 }
 
@@ -212,26 +222,30 @@ func (t *StateTrie) MustDelete(key []byte) {
 // If a node is not found in the database, a MissingNodeError is returned.
 func (t *StateTrie) DeleteStorage(_ common.Address, key []byte) error {
 	hk := crypto.Keccak256(key)
-	delete(t.getSecKeyCache(), common.Hash(hk))
+	if t.preimages != nil {
+		delete(t.getSecKeyCache(), common.Hash(hk))
+	}
 	return t.trie.Delete(hk)
 }
 
 // DeleteAccount abstracts an account deletion from the trie.
 func (t *StateTrie) DeleteAccount(address common.Address) error {
 	hk := crypto.Keccak256(address.Bytes())
-	delete(t.getSecKeyCache(), common.Hash(hk))
+	if t.preimages != nil {
+		delete(t.getSecKeyCache(), common.Hash(hk))
+	}
 	return t.trie.Delete(hk)
 }
 
 // GetKey returns the sha3 preimage of a hashed key that was
 // previously used to store a value.
 func (t *StateTrie) GetKey(shaKey []byte) []byte {
-	if key, ok := t.getSecKeyCache()[common.BytesToHash(shaKey)]; ok {
-		return key
-	}
 	if t.preimages == nil {
 		return nil
 	}
+	if key, ok := t.getSecKeyCache()[common.BytesToHash(shaKey)]; ok {
+		return key
+	}
 	return t.preimages.Preimage(common.BytesToHash(shaKey))
 }
 
diff --git a/triedb/database.go b/triedb/database.go
index 18e24cd176..1c9d49b245 100644
--- a/triedb/database.go
+++ b/triedb/database.go
@@ -213,6 +213,11 @@ func (db *Database) InsertPreimage(preimages map[common.Hash][]byte) {
 	db.preimages.insertPreimage(preimages)
 }
 
+// PreimageEnabled returns the indicator if the pre-image store is enabled.
+func (db *Database) PreimageEnabled() bool {
+	return db.preimages != nil
+}
+
 // Cap iteratively flushes old but still referenced trie nodes until the total
 // memory usage goes below the given threshold. The held pre-images accumulated
 // up to this point will be flushed in case the size exceeds the threshold.