Make Encode methods accept []byte

Taskfile.yml

@@ -88,7 +88,8 @@ tasks:
     desc: "Run benchmarks"
     cmds:
       - go test -list='Benchmark.*' ./...
-      - go test -count=10 -bench=BenchmarkDocument -benchtime={{.BENCH_TIME}} ./wirebson | tee -a new.txt
+      # -timeout is needed due to https://github.com/golang/go/issues/69181
+      - go test -count=10 -bench=BenchmarkDocument -benchtime={{.BENCH_TIME}} -timeout=60m ./wirebson | tee -a new.txt
       - bin/benchstat old.txt new.txt
 
   fuzz:

op_msg.go

@@ -38,7 +38,8 @@ type OpMsg struct {
 
 // NewOpMsg creates a message with a single section of kind 0 with a single document.
 func NewOpMsg(doc wirebson.AnyDocument) (*OpMsg, error) {
-	raw, err := doc.Encode()
+	raw := make([]byte, wirebson.Size(doc))
+	err := doc.Encode(raw)
 	if err != nil {
 		return nil, lazyerrors.Error(err)
 	}

op_query.go

@@ -37,7 +37,8 @@ type OpQuery struct {
 
 // NewOpQuery creates a new OpQuery message.
 func NewOpQuery(doc wirebson.AnyDocument) (*OpQuery, error) {
-	raw, err := doc.Encode()
+	raw := make([]byte, wirebson.Size(doc))
+	err := doc.Encode(raw)
 	if err != nil {
 		return nil, lazyerrors.Error(err)
 	}

op_reply.go

@@ -37,7 +37,8 @@ type OpReply struct {
 
 // NewOpReply creates a new OpReply message.
 func NewOpReply(doc wirebson.AnyDocument) (*OpReply, error) {
-	raw, err := doc.Encode()
+	raw := make([]byte, wirebson.Size(doc))
+	err := doc.Encode(raw)
 	if err != nil {
 		return nil, lazyerrors.Error(err)
 	}
@@ -132,7 +133,9 @@ func (reply *OpReply) RawDocument() wirebson.RawDocument {
 // SetDocument sets reply document.
 func (reply *OpReply) SetDocument(doc *wirebson.Document) {
 	var err error
-	reply.document, err = doc.Encode()
+
+	reply.document = make([]byte, wirebson.Size(doc))
+	err = doc.Encode(reply.document)
 	if err != nil {
 		panic(err)
 	}

wire_test.go

@@ -40,7 +40,9 @@ func TestMain(m *testing.M) {
 func makeRawDocument(pairs ...any) wirebson.RawDocument {
 	d := wirebson.MustDocument(pairs...)
 
-	raw, err := d.Encode()
+	raw := make([]byte, wirebson.Size(d))
+
+	err := d.Encode(raw)
 	if err != nil {
 		panic(err)
 	}

wirebson/array.go

@@ -15,7 +15,6 @@
 package wirebson
 
 import (
-	"bytes"
 	"encoding/binary"
 	"iter"
 	"log/slog"
@@ -149,32 +148,30 @@ func (arr *Array) SortInterface(less func(a, b any) bool) sort.Interface {
 	}
 }
 
-// Encode encodes non-nil Array.
+// Encode encodes Array v into raw.
 //
-// TODO https://github.com/FerretDB/wire/issues/21
-// This method should accept a slice of bytes, not return it.
-// That would allow to avoid unnecessary allocations.
-func (arr *Array) Encode() (RawArray, error) {
+// raw must be at least Size(arr) bytes long; otherwise, Encode will panic.
+// Only raw[0:Size(arr)] bytes are modified.
+func (arr *Array) Encode(raw RawArray) error {
 	must.NotBeZero(arr)
 
-	size := sizeArray(arr)
-	buf := bytes.NewBuffer(make([]byte, 0, size))
+	// ensure raw length early
+	s := sizeArray(arr)
+	raw[s-1] = 0
 
-	if err := binary.Write(buf, binary.LittleEndian, uint32(size)); err != nil {
-		return nil, lazyerrors.Error(err)
-	}
+	binary.LittleEndian.PutUint32(raw, uint32(s))
 
-	for i, v := range arr.values {
-		if err := encodeField(buf, strconv.Itoa(i), v); err != nil {
-			return nil, lazyerrors.Error(err)
+	i := 4
+	for n, v := range arr.values {
+		w, err := encodeField(raw[i:], strconv.Itoa(n), v)
+		if err != nil {
+			return lazyerrors.Error(err)
 		}
-	}
 
-	if err := binary.Write(buf, binary.LittleEndian, byte(0)); err != nil {
-		return nil, lazyerrors.Error(err)
+		i += w
 	}
 
-	return buf.Bytes(), nil
+	return nil
 }
 
 // Decode returns itself to implement [AnyArray].

wirebson/bson.go

@@ -70,7 +70,7 @@ type ScalarType interface {
 // Note that the Encode and Decode methods could return the receiver itself,
 // so care must be taken when results are modified.
 type AnyDocument interface {
-	Encode() (RawDocument, error)
+	Encode(RawDocument) error
 	Decode() (*Document, error)
 	LogMessage() string
 	LogMessageIndent() string
@@ -81,7 +81,7 @@ type AnyDocument interface {
 // Note that the Encode and Decode methods could return the receiver itself,
 // so care must be taken when results are modified.
 type AnyArray interface {
-	Encode() (RawArray, error)
+	Encode(RawArray) error
 	Decode() (*Array, error)
 	LogMessage() string
 	LogMessageIndent() string

wirebson/bson_test.go

@@ -767,7 +767,8 @@ func TestNormal(t *testing.T) {
 				assert.NotEmpty(t, doc.LogMessage())
 				assert.NotEmpty(t, doc.LogMessageIndent())
 
-				raw, err := doc.Encode()
+				raw := make(RawDocument, Size(doc))
+				err = doc.Encode(raw)
 				require.NoError(t, err)
 				assert.Equal(t, tc.raw, raw)
 			})
@@ -783,7 +784,8 @@ func TestNormal(t *testing.T) {
 				assert.NotEmpty(t, doc.LogMessage())
 				assert.Equal(t, testutil.Unindent(tc.b), doc.LogMessageIndent())
 
-				raw, err := doc.Encode()
+				raw := make(RawDocument, Size(doc))
+				err = doc.Encode(raw)
 				require.NoError(t, err)
 				assert.Equal(t, tc.raw, raw)
 			})
@@ -866,8 +868,9 @@ func BenchmarkDocument(b *testing.B) {
 				b.ReportAllocs()
 				b.ResetTimer()
 
+				raw = make(RawDocument, Size(doc))
 				for range b.N {
-					raw, err = doc.Encode()
+					err = doc.Encode(raw)
 				}
 
 				b.StopTimer()
@@ -946,8 +949,9 @@ func BenchmarkDocument(b *testing.B) {
 				b.ReportAllocs()
 				b.ResetTimer()
 
+				raw = make(RawDocument, Size(doc))
 				for range b.N {
-					raw, err = doc.Encode()
+					err = doc.Encode(raw)
 				}
 
 				b.StopTimer()
@@ -1041,7 +1045,8 @@ func testRawDocument(t *testing.T, rawDoc RawDocument) {
 		assert.NotEmpty(t, doc.LogMessage())
 		assert.NotEmpty(t, doc.LogMessageIndent())
 
-		raw, err := doc.Encode()
+		raw := make(RawDocument, Size(doc))
+		err = doc.Encode(raw)
 		if err == nil {
 			assert.Equal(t, rawDoc, raw)
 		}
@@ -1060,7 +1065,8 @@ func testRawDocument(t *testing.T, rawDoc RawDocument) {
 		assert.NotEmpty(t, doc.LogMessage())
 		assert.NotEmpty(t, doc.LogMessageIndent())
 
-		raw, err := doc.Encode()
+		raw := make(RawDocument, Size(doc))
+		err = doc.Encode(raw)
 		require.NoError(t, err)
 		assert.Equal(t, rawDoc, raw)
 	})

wirebson/document.go

@@ -15,7 +15,6 @@
 package wirebson
 
 import (
-	"bytes"
 	"encoding/binary"
 	"iter"
 	"log/slog"
@@ -227,30 +226,28 @@ func (doc *Document) Command() string {
 
 // Encode encodes non-nil Document.
 //
-// TODO https://github.com/FerretDB/wire/issues/21
-// This method should accept a slice of bytes, not return it.
-// That would allow to avoid unnecessary allocations.
-func (doc *Document) Encode() (RawDocument, error) {
+// raw must be at least Size(doc) bytes long; otherwise, Encode will panic.
+// Only raw[0:Size(doc)] bytes are modified.
+func (doc *Document) Encode(raw RawDocument) error {
 	must.NotBeZero(doc)
 
-	size := sizeDocument(doc)
-	buf := bytes.NewBuffer(make([]byte, 0, size))
+	// ensure raw length early
+	s := sizeDocument(doc)
+	raw[s-1] = 0
 
-	if err := binary.Write(buf, binary.LittleEndian, uint32(size)); err != nil {
-		return nil, lazyerrors.Error(err)
-	}
+	binary.LittleEndian.PutUint32(raw, uint32(s))
 
+	i := 4
 	for _, f := range doc.fields {
-		if err := encodeField(buf, f.name, f.value); err != nil {
-			return nil, lazyerrors.Error(err)
+		w, err := encodeField(raw[i:], f.name, f.value)
+		if err != nil {
+			return lazyerrors.Error(err)
 		}
-	}
 
-	if err := binary.Write(buf, binary.LittleEndian, byte(0)); err != nil {
-		return nil, lazyerrors.Error(err)
+		i += w
 	}
 
-	return buf.Bytes(), nil
+	return nil
 }
 
 // Decode returns itself to implement [AnyDocument].