diff --git a/rlp/encbuffer.go b/rlp/encbuffer.go index cd3f63b99c..0c783e3965 100644 --- a/rlp/encbuffer.go +++ b/rlp/encbuffer.go @@ -350,6 +350,16 @@ func (w *EncoderBuffer) AppendToBytes(dst []byte) []byte { return out } +// Size returns the total size of the content that was encoded up to this point. +// Note this does not count the size of any lists which are still 'open' (i.e. for +// which ListEnd has not been called yet). +func (w EncoderBuffer) Size() int { + if w.buf == nil { + return 0 + } + return w.buf.size() +} + // Write appends b directly to the encoder output. func (w EncoderBuffer) Write(b []byte) (int, error) { return w.buf.Write(b) diff --git a/rlp/encode_test.go b/rlp/encode_test.go index 83c8b692d1..10d1232631 100644 --- a/rlp/encode_test.go +++ b/rlp/encode_test.go @@ -503,7 +503,40 @@ func TestEncodeToReaderReturnToPool(t *testing.T) { wg.Wait() } -var sink any +func TestEncoderBufferSize(t *testing.T) { + var output bytes.Buffer + eb := NewEncoderBuffer(&output) + + assertSize := func(state string, expectedSize int) { + t.Helper() + if s := eb.Size(); s != expectedSize { + t.Fatalf("wrong size %s: %d", state, s) + } + } + + assertSize("empty buffer", 0) + outerList := eb.List() + assertSize("after outer List()", 0) + eb.WriteString("abc") + assertSize("after string write", 4) + innerList := eb.List() + assertSize("after inner List()", 4) + eb.WriteUint64(1) + eb.WriteUint64(2) + assertSize("after inner list writes", 6) + eb.ListEnd(innerList) + assertSize("after end of inner list", 7) + eb.ListEnd(outerList) + assertSize("after end of outer list", 8) + eb.Flush() + assertSize("after Flush()", 0) + + if output.Len() != 8 { + t.Fatalf("wrong final output size %d", output.Len()) + } +} + +var sink interface{} func BenchmarkIntsize(b *testing.B) { for i := 0; i < b.N; i++ {