core/stateless: preserve witness keys across encoding

core/stateless/encoding.go

@@ -34,6 +34,10 @@ func (w *Witness) ToExtWitness() *ExtWitness {
 	for code := range w.Codes {
 		ext.Codes = append(ext.Codes, []byte(code))
 	}
+	ext.Keys = make([]hexutil.Bytes, 0, len(w.Keys))
+	for key := range w.Keys {
+		ext.Keys = append(ext.Keys, []byte(key))
+	}
 	ext.State = make([]hexutil.Bytes, 0, len(w.State))
 	for node := range w.State {
 		ext.State = append(ext.State, []byte(node))
@@ -52,6 +56,10 @@ func (w *Witness) FromExtWitness(ext *ExtWitness) error {
 	for _, code := range ext.Codes {
 		w.Codes[string(code)] = struct{}{}
 	}
+	w.Keys = make(map[string]struct{}, len(ext.Keys))
+	for _, key := range ext.Keys {
+		w.Keys[string(key)] = struct{}{}
+	}
 	w.State = make(map[string]struct{}, len(ext.State))
 	for _, node := range ext.State {
 		w.State[string(node)] = struct{}{}

core/stateless/encoding_test.go

@@ -0,0 +1,73 @@
+// Copyright 2026 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package stateless
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common/hexutil"
+	"github.com/ethereum/go-ethereum/core/types"
+)
+
+func TestWitnessExtRoundTripPreservesKeys(t *testing.T) {
+	ext := &ExtWitness{
+		Headers: []*types.Header{{Number: big.NewInt(1)}},
+		Codes:   []hexutil.Bytes{[]byte("code")},
+		Keys:    []hexutil.Bytes{[]byte("key-a"), []byte("key-b")},
+		State:   []hexutil.Bytes{[]byte("state")},
+	}
+	var witness Witness
+	if err := witness.FromExtWitness(ext); err != nil {
+		t.Fatalf("FromExtWitness error: %v", err)
+	}
+	if len(witness.Keys) != len(ext.Keys) {
+		t.Fatalf("stored %d keys, want %d", len(witness.Keys), len(ext.Keys))
+	}
+	for _, key := range ext.Keys {
+		if _, ok := witness.Keys[string(key)]; !ok {
+			t.Fatalf("missing key %q after FromExtWitness", []byte(key))
+		}
+	}
+	roundTrip := witness.ToExtWitness()
+	if len(roundTrip.Keys) != len(ext.Keys) {
+		t.Fatalf("encoded %d keys, want %d", len(roundTrip.Keys), len(ext.Keys))
+	}
+	encoded := make(map[string]struct{}, len(roundTrip.Keys))
+	for _, key := range roundTrip.Keys {
+		encoded[string(key)] = struct{}{}
+	}
+	for _, key := range ext.Keys {
+		if _, ok := encoded[string(key)]; !ok {
+			t.Fatalf("missing key %q after ToExtWitness", []byte(key))
+		}
+	}
+}
+
+func TestWitnessAddKey(t *testing.T) {
+	witness := &Witness{}
+	witness.AddKey([]byte("key-a"), nil, []byte("key-a"), []byte("key-b"))
+
+	if len(witness.Keys) != 2 {
+		t.Fatalf("stored %d keys, want 2", len(witness.Keys))
+	}
+	for _, key := range []string{"key-a", "key-b"} {
+		if _, ok := witness.Keys[key]; !ok {
+			t.Fatalf("missing key %q", key)
+		}
+	}
+}

core/stateless/witness.go

@@ -40,6 +40,7 @@ type Witness struct {
 
 	Headers []*types.Header     // Past headers in reverse order (0=parent, 1=parent's-parent, etc). First *must* be set.
 	Codes   map[string]struct{} // Set of bytecodes ran or accessed
+	Keys    map[string]struct{} // Set of accessed trie keys
 	State   map[string]struct{} // Set of MPT state trie nodes (account and storage together)
 
 	chain HeaderReader  // Chain reader to convert block hash ops to header proofs
@@ -64,6 +65,7 @@ func NewWitness(context *types.Header, chain HeaderReader, enableStats bool) (*W
 		context: context,
 		Headers: headers,
 		Codes:   make(map[string]struct{}),
+		Keys:    make(map[string]struct{}),
 		State:   make(map[string]struct{}),
 		chain:   chain,
 	}
@@ -119,8 +121,16 @@ func (w *Witness) ReportMetrics(blockNumber uint64) {
 	w.stats.ReportMetrics(blockNumber)
 }
 
-func (w *Witness) AddKey() {
+func (w *Witness) AddKey(keys ...[]byte) {
-	panic("not yet implemented")
+	if w.Keys == nil {
+		w.Keys = make(map[string]struct{})
+	}
+	for _, key := range keys {
+		if len(key) == 0 {
+			continue
+		}
+		w.Keys[string(key)] = struct{}{}
+	}
 }
 
 // Copy deep-copies the witness object.  Witness.Block isn't deep-copied as it
@@ -129,6 +139,7 @@ func (w *Witness) Copy() *Witness {
 	cpy := &Witness{
 		Headers: slices.Clone(w.Headers),
 		Codes:   maps.Clone(w.Codes),
+		Keys:    maps.Clone(w.Keys),
 		State:   maps.Clone(w.State),
 		chain:   w.chain,
 	}