From 6735dde9ed051d1368f048769bfc224c5caf3e32 Mon Sep 17 00:00:00 2001
From: jstr1121 <jun@stridelabs.co>
Date: Thu, 26 Oct 2023 21:01:33 +0800
Subject: [PATCH] add a way to cleanup snaps on state

---
 core/state/statedb.go    |  5 +++++
 internal/ethapi/api.go   |  1 +
 trie/trienode/witness.go | 13 ++++++++++---
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/core/state/statedb.go b/core/state/statedb.go
index df44d25f89..d3f9e2f252 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -180,6 +180,11 @@ func (s *StateDB) GetWitness()  *trienode.Witness {
 	return s.trie.GetWitness()
 }
 
+func (s *StateDB) CleanSnaps() {
+	s.snaps = nil
+	s.snap = nil
+}
+
 // StartPrefetcher initializes a new trie prefetcher to pull in nodes from the
 // state trie concurrently while the state is mutated so that when we reach the
 // commit phase, most of the needed data is already hot.
diff --git a/internal/ethapi/api.go b/internal/ethapi/api.go
index a75b89db71..44d4fc6d34 100644
--- a/internal/ethapi/api.go
+++ b/internal/ethapi/api.go
@@ -972,6 +972,7 @@ func (s *BlockChainAPI) GetRequiredBlockState(ctx context.Context, blockNrOrHash
 	if err != nil {
 		return nil, nil
 	}
+	state.CleanSnaps()
 	s.ProcessBlock(ctx, block, state, vm.Config{})
 	return state.GetWitness(), nil
 	// return s.traceBlock(ctx, block)
diff --git a/trie/trienode/witness.go b/trie/trienode/witness.go
index bbd0c71475..9fb2d0fe10 100644
--- a/trie/trienode/witness.go
+++ b/trie/trienode/witness.go
@@ -32,13 +32,20 @@ type Witness struct {
 }
 
 func (u *Witness) MarshalJSON() ([]byte, error) {
-	nodes := make(map[string]string)
+	type Node struct {
+		Key   string
+		Value string
+	}
+	nodes := []Node{}
 	for key, node := range u.Nodes {
-		nodes[key] = hex.EncodeToString(node)
+		nodes = append(nodes, Node{
+			Key:   hex.EncodeToString([]byte(key)),
+			Value: hex.EncodeToString(node),
+		})
 	}
 	return json.Marshal(&struct {
 		Owner common.Hash
-		Nodes map[string]string
+		Nodes []Node
 	}{
 		Owner: u.Owner,
 		Nodes: nodes,