diff --git a/core/stateless/encoding.go b/core/stateless/encoding.go index 1b20c4cb2a..4ff7acf63b 100644 --- a/core/stateless/encoding.go +++ b/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{}{} diff --git a/core/stateless/encoding_test.go b/core/stateless/encoding_test.go new file mode 100644 index 0000000000..cff67c39a0 --- /dev/null +++ b/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 . + +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) + } + } +} diff --git a/core/stateless/witness.go b/core/stateless/witness.go index f1321699c1..f5164eb73e 100644 --- a/core/stateless/witness.go +++ b/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() { - panic("not yet implemented") +func (w *Witness) AddKey(keys ...[]byte) { + 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, }