feat(trie): track committed state in witness

core/state/statedb.go

@@ -1256,9 +1256,29 @@ func (s *StateDB) commit(deleteEmptyObjects bool, noStorageWiping bool) (*stateU
 	origin := s.originalRoot
 	s.originalRoot = root
 
+	// --- Start fork code ---
+	// Track the committed nodes
+	if s.witness != nil {
+		s.witness.AddCommitted(nodes)
+	}
+	// --- End fork code ---
+
 	return newStateUpdate(noStorageWiping, origin, root, deletes, updates, nodes), nil
 }
 
+// --- Start fork code ---
+// CommitWithoutFlush is a wrapper of commit which does not commit the state mutations
+// to the configured data stores.
+func (s *StateDB) CommitWithoutFlush(deleteEmptyObjects bool, noStorageWiping bool) (common.Hash, error) {
+	ret, err := s.commit(deleteEmptyObjects, noStorageWiping)
+	if err != nil {
+		return common.Hash{}, err
+	}
+	return ret.root, nil
+}
+
+// --- End fork code ---
+
 // commitAndFlush is a wrapper of commit which also commits the state mutations
 // to the configured data stores.
 func (s *StateDB) commitAndFlush(block uint64, deleteEmptyObjects bool, noStorageWiping bool) (*stateUpdate, error) {

core/stateless/witness.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/trie/trienode"
 )
 
 // HeaderReader is an interface to pull in headers in place of block hashes for
@@ -42,6 +43,10 @@ type Witness struct {
 	Codes   map[string]struct{} // Set of bytecodes ran or accessed
 	State   map[string]struct{} // Set of MPT state trie nodes (account and storage together)
 
+	// --- Start fork code ---
+	Committed map[string]struct{} // Set of MPT state trie nodes (account and storage together) that have been committed
+	// --- End fork code ---
+
 	chain HeaderReader // Chain reader to convert block hash ops to header proofs
 	lock  sync.Mutex   // Lock to allow concurrent state insertions
 }
@@ -64,7 +69,10 @@ func NewWitness(context *types.Header, chain HeaderReader) (*Witness, error) {
 		Headers: headers,
 		Codes:   make(map[string]struct{}),
 		State:   make(map[string]struct{}),
-		chain:   chain,
+		// --- Start fork code ---
+		Committed: make(map[string]struct{}),
+		// --- End fork code ---
+		chain: chain,
 	}, nil
 }
 
@@ -98,6 +106,17 @@ func (w *Witness) AddState(nodes map[string]struct{}) {
 	maps.Copy(w.State, nodes)
 }
 
+// AddCommitted inserts a batch of MPT trie nodes into the witness.
+func (w *Witness) AddCommitted(nodes *trienode.MergedNodeSet) {
+	for _, set := range nodes.Sets {
+		for _, node := range set.Nodes {
+			if len(node.Blob) > 0 {
+				w.Committed[string(node.Blob)] = struct{}{}
+			}
+		}
+	}
+}
+
 // Copy deep-copies the witness object.  Witness.Block isn't deep-copied as it
 // is never mutated by Witness
 func (w *Witness) Copy() *Witness {