From dc0a7a74c15d7372021a5c5da6f2515ee535094d Mon Sep 17 00:00:00 2001
From: 0xbundler <124862913+0xbundler@users.noreply.github.com>
Date: Mon, 30 Oct 2023 16:11:06 +0800
Subject: [PATCH] state/stateobject: add dirtystorage touch & revive logic;
 trie: fix local review leaf expand bug;

---
 core/state/state_expiry.go |  6 ------
 core/state/state_object.go | 18 ++++++++++++++++++
 trie/trie_expiry.go        |  8 +++++++-
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/core/state/state_expiry.go b/core/state/state_expiry.go
index d0742ff65f..75b1b9ec60 100644
--- a/core/state/state_expiry.go
+++ b/core/state/state_expiry.go
@@ -41,9 +41,6 @@ func fetchExpiredStorageFromRemote(meta *stateExpiryMeta, addr common.Address, r
 		// if there need revive expired state, try to revive locally, when the node is not being pruned, just renew the epoch
 		val, err := tr.TryLocalRevive(addr, key.Bytes())
 		//log.Debug("fetchExpiredStorageFromRemote TryLocalRevive", "addr", addr, "key", key, "val", val, "err", err)
-		if _, ok := err.(*trie.MissingNodeError); !ok {
-			return nil, err
-		}
 		switch err.(type) {
 		case *trie.MissingNodeError:
 			// cannot revive locally, request from remote
@@ -86,9 +83,6 @@ func batchFetchExpiredFromRemote(expiryMeta *stateExpiryMeta, addr common.Addres
 		for i, key := range keys {
 			val, err := tr.TryLocalRevive(addr, key.Bytes())
 			//log.Debug("fetchExpiredStorageFromRemote TryLocalRevive", "addr", addr, "key", key, "val", val, "err", err)
-			if _, ok := err.(*trie.MissingNodeError); !ok {
-				return nil, err
-			}
 			switch err.(type) {
 			case *trie.MissingNodeError:
 				expiredKeys = append(expiredKeys, key)
diff --git a/core/state/state_object.go b/core/state/state_object.go
index 317a70495e..e137eb4a36 100644
--- a/core/state/state_object.go
+++ b/core/state/state_object.go
@@ -500,6 +500,24 @@ func (s *stateObject) updateTrie() (Trie, error) {
 				}
 				//log.Debug("updateTrie pendingFutureReviveState", "contract", s.address, "key", key, "epoch", s.db.Epoch(), "tr.epoch", tr.Epoch(), "tr", fmt.Sprintf("%p", tr), "ins", fmt.Sprintf("%p", s))
 			}
+			// TODO(0xbundler): find some trie node with wrong epoch, temporary add get op, will fix later
+			//for key, val := range dirtyStorage {
+			//	_, err = tr.GetStorage(s.address, key.Bytes())
+			//	if err == nil {
+			//		continue
+			//	}
+			//	log.Error("EnableExpire GetStorage error", "addr", s.address, "key", key, "val", val, "origin", s.originStorage[key], "err", err)
+			//	enErr, ok := err.(*trie.ExpiredNodeError)
+			//	if !ok {
+			//		s.db.setError(fmt.Errorf("state object dirtyStorage err, contract: %v, key: %v, err: %v", s.address, key, err))
+			//		continue
+			//	}
+			//	if _, err = fetchExpiredStorageFromRemote(s.db.expiryMeta, s.address, s.data.Root, tr, enErr.Path, key); err != nil {
+			//		log.Error("EnableExpire GetStorage fetchExpiredStorageFromRemote error", "addr", s.address, "key", key, "val", val, "origin", s.originStorage[key], "err", err)
+			//		s.db.setError(fmt.Errorf("state object dirtyStorage fetchExpiredStorageFromRemote err, contract: %v, key: %v, path: %v, err: %v", s.address, key, enErr.Path, err))
+			//	}
+			//	//log.Debug("updateTrie dirtyStorage", "contract", s.address, "key", key, "epoch", s.db.Epoch(), "tr.epoch", tr.Epoch(), "tr", fmt.Sprintf("%p", tr), "ins", fmt.Sprintf("%p", s))
+			//}
 		}
 		for key, value := range dirtyStorage {
 			if len(value) == 0 {
diff --git a/trie/trie_expiry.go b/trie/trie_expiry.go
index 94ee7927ee..4a31086574 100644
--- a/trie/trie_expiry.go
+++ b/trie/trie_expiry.go
@@ -30,7 +30,13 @@ func (t *Trie) tryLocalRevive(origNode node, key []byte, pos int, epoch types.St
 		return n, n, expired, nil
 	case *shortNode:
 		if len(key)-pos < len(n.Key) || !bytes.Equal(n.Key, key[pos:pos+len(n.Key)]) {
-			// key not found in trie
+			// key not found in trie, but just revive for expand
+			if t.renewNode(epoch, false, expired) {
+				n = n.copy()
+				n.setEpoch(t.currentEpoch)
+				n.flags = t.newFlag()
+				return nil, n, true, nil
+			}
 			return nil, n, false, nil
 		}
 		value, newnode, didResolve, err := t.tryLocalRevive(n.Val, key, pos+len(n.Key), epoch)