diff --git a/rlp/encbuffer.go b/rlp/encbuffer.go index 61d8bd059c..ca0aa290fe 100644 --- a/rlp/encbuffer.go +++ b/rlp/encbuffer.go @@ -366,6 +366,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 e63ea319b4..8dc9fdaf1f 100644 --- a/rlp/encode_test.go +++ b/rlp/encode_test.go @@ -507,6 +507,39 @@ func TestEncodeToReaderReturnToPool(t *testing.T) { wg.Wait() } +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) {