Merge pull request #2678 from CortexFoundation/dev

rlp: add Size method to EncoderBuffer

Author: ucwong

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)

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++ {