From 2530c4a3e666665a337011a00c38d59d051df340 Mon Sep 17 00:00:00 2001
From: jstr1121 <jun@stridelabs.co>
Date: Tue, 24 Oct 2023 01:04:24 +0800
Subject: [PATCH] expose GetWitness endpoint on Trie objects

---
 core/state/database.go | 2 ++
 core/state/statedb.go  | 8 ++++++++
 internal/ethapi/api.go | 6 ++++--
 light/trie.go          | 4 ++++
 trie/secure_trie.go    | 4 ++++
 trie/trie.go           | 4 ++++
 6 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/core/state/database.go b/core/state/database.go
index 1022cb90dd..6981587994 100644
--- a/core/state/database.go
+++ b/core/state/database.go
@@ -75,6 +75,8 @@ type Trie interface {
 	// a trie.MissingNodeError is returned.
 	GetStorage(addr common.Address, key []byte) ([]byte, error)
 
+	GetWitness() *trienode.Witness
+
 	// GetAccount abstracts an account read from the trie. It retrieves the
 	// account blob from the trie with provided account address and decodes it
 	// with associated decoding algorithm. If the specified account is not in
diff --git a/core/state/statedb.go b/core/state/statedb.go
index 8781144941..df44d25f89 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -172,6 +172,14 @@ func New(root common.Hash, db Database, snaps *snapshot.Tree) (*StateDB, error)
 	return sdb, 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.
+func (s *StateDB) GetWitness()  *trienode.Witness {
+	return s.trie.GetWitness()
+}
+
 // 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 a4956122c0..a75b89db71 100644
--- a/internal/ethapi/api.go
+++ b/internal/ethapi/api.go
@@ -48,6 +48,7 @@ import (
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/rpc"
 	"github.com/ethereum/go-ethereum/trie"
+	"github.com/ethereum/go-ethereum/trie/trienode"
 	"github.com/tyler-smith/go-bip39"
 )
 
@@ -961,7 +962,7 @@ type RecentBlockHash struct {
 }
 
 // GetRequiredBlockState returns all state required to execute a single historical block.
-func (s *BlockChainAPI) GetRequiredBlockState(ctx context.Context, blockNrOrHash rpc.BlockNumberOrHash) ([]*txTraceResult, error) {
+func (s *BlockChainAPI) GetRequiredBlockState(ctx context.Context, blockNrOrHash rpc.BlockNumberOrHash) (*trienode.Witness, error) {
 	block, err := s.b.BlockByNumberOrHash(ctx, blockNrOrHash)
 	if block == nil || err != nil {
 		// When the block doesn't exist, the RPC method should return JSON null
@@ -972,7 +973,8 @@ func (s *BlockChainAPI) GetRequiredBlockState(ctx context.Context, blockNrOrHash
 		return nil, nil
 	}
 	s.ProcessBlock(ctx, block, state, vm.Config{})
-	return s.traceBlock(ctx, block)
+	return state.GetWitness(), nil
+	// return s.traceBlock(ctx, block)
 }
 
 // txTraceResult is the result of a single transaction trace.
diff --git a/light/trie.go b/light/trie.go
index 0a38f41e53..9af9f404a2 100644
--- a/light/trie.go
+++ b/light/trie.go
@@ -106,6 +106,10 @@ type odrTrie struct {
 	trie *trie.Trie
 }
 
+func (t *odrTrie) GetWitness() *trienode.Witness {
+	return t.trie.GetWitness()
+}
+
 func (t *odrTrie) GetStorage(_ common.Address, key []byte) ([]byte, error) {
 	key = crypto.Keccak256(key)
 	var enc []byte
diff --git a/trie/secure_trie.go b/trie/secure_trie.go
index ed6b872fb1..a7f45bc29b 100644
--- a/trie/secure_trie.go
+++ b/trie/secure_trie.go
@@ -81,6 +81,10 @@ func (t *StateTrie) MustGet(key []byte) []byte {
 	return t.trie.MustGet(t.hashKey(key))
 }
 
+func (t *StateTrie) GetWitness() *trienode.Witness {
+	return t.trie.GetWitness()
+}
+
 // GetStorage attempts to retrieve a storage slot with provided account address
 // and slot key. The value bytes must not be modified by the caller.
 // If the specified storage slot is not in the trie, nil will be returned.
diff --git a/trie/trie.go b/trie/trie.go
index 70d09d0c5a..9b7960b29d 100644
--- a/trie/trie.go
+++ b/trie/trie.go
@@ -688,3 +688,7 @@ func (t *Trie) Reset() {
 	t.witness = trienode.NewWitness(common.Hash{})
 	t.committed = false
 }
+
+func (t *Trie) GetWitness() *trienode.Witness {
+	return t.witness
+}