encoding: add MarshalAppend to protojson and prototext

Adds MarshalAppend methods to allow for byte slices to be reused.
Copies signature from the binary encoding.

Small changes to internal json and text libraries to use strconv
AppendInt and AppendUint for number encoding.

Change-Id: Ife7c8979c1c153a0a0bf9b70b296b8158d38dffc
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/489615
Reviewed-by: Edward McFarlane <emcfarlane000@gmail.com>
Reviewed-by: Joseph Tsai <joetsai@digital-static.net>
Reviewed-by: Lasse Folger <lassefolger@google.com>
Reviewed-by: Damien Neil <dneil@google.com>
diff --git a/encoding/protojson/encode.go b/encoding/protojson/encode.go
index d09d22e..66b9587 100644
--- a/encoding/protojson/encode.go
+++ b/encoding/protojson/encode.go
@@ -106,13 +106,19 @@
 // MarshalOptions. Do not depend on the output being stable. It may change over
 // time across different versions of the program.
 func (o MarshalOptions) Marshal(m proto.Message) ([]byte, error) {
-	return o.marshal(m)
+	return o.marshal(nil, m)
+}
+
+// MarshalAppend appends the JSON format encoding of m to b,
+// returning the result.
+func (o MarshalOptions) MarshalAppend(b []byte, m proto.Message) ([]byte, error) {
+	return o.marshal(b, m)
 }
 
 // marshal is a centralized function that all marshal operations go through.
 // For profiling purposes, avoid changing the name of this function or
 // introducing other code paths for marshal that do not go through this.
-func (o MarshalOptions) marshal(m proto.Message) ([]byte, error) {
+func (o MarshalOptions) marshal(b []byte, m proto.Message) ([]byte, error) {
 	if o.Multiline && o.Indent == "" {
 		o.Indent = defaultIndent
 	}
@@ -120,7 +126,7 @@
 		o.Resolver = protoregistry.GlobalTypes
 	}
 
-	internalEnc, err := json.NewEncoder(o.Indent)
+	internalEnc, err := json.NewEncoder(b, o.Indent)
 	if err != nil {
 		return nil, err
 	}
@@ -128,7 +134,7 @@
 	// Treat nil message interface as an empty message,
 	// in which case the output in an empty JSON object.
 	if m == nil {
-		return []byte("{}"), nil
+		return append(b, '{', '}'), nil
 	}
 
 	enc := encoder{internalEnc, o}
diff --git a/encoding/protojson/encode_test.go b/encoding/protojson/encode_test.go
index e8db20b..adda076 100644
--- a/encoding/protojson/encode_test.go
+++ b/encoding/protojson/encode_test.go
@@ -2310,3 +2310,44 @@
 		})
 	}
 }
+
+func TestEncodeAppend(t *testing.T) {
+	want := []byte("prefix")
+	got := append([]byte(nil), want...)
+	got, err := protojson.MarshalOptions{}.MarshalAppend(got, &pb3.Scalars{
+		SString: "value",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.HasPrefix(got, want) {
+		t.Fatalf("MarshalAppend modified prefix: got %v, want prefix %v", got, want)
+	}
+}
+
+func TestMarshalAppendAllocations(t *testing.T) {
+	m := &pb3.Scalars{SInt32: 1}
+	const count = 1000
+	size := 12
+	b := make([]byte, size)
+	// AllocsPerRun returns an integral value.
+	marshalAllocs := testing.AllocsPerRun(count, func() {
+		_, err := protojson.MarshalOptions{}.MarshalAppend(b[:0], m)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+	b = nil
+	marshalAppendAllocs := testing.AllocsPerRun(count, func() {
+		var err error
+		b, err = protojson.MarshalOptions{}.MarshalAppend(b, m)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+	if marshalAllocs != marshalAppendAllocs {
+		t.Errorf("%v allocs/op when writing to a preallocated buffer", marshalAllocs)
+		t.Errorf("%v allocs/op when repeatedly appending to a slice", marshalAppendAllocs)
+		t.Errorf("expect amortized allocs/op to be identical")
+	}
+}
diff --git a/encoding/prototext/encode.go b/encoding/prototext/encode.go
index ebf6c65..722a7b4 100644
--- a/encoding/prototext/encode.go
+++ b/encoding/prototext/encode.go
@@ -101,13 +101,19 @@
 // MarshalOptions object. Do not depend on the output being stable. It may
 // change over time across different versions of the program.
 func (o MarshalOptions) Marshal(m proto.Message) ([]byte, error) {
-	return o.marshal(m)
+	return o.marshal(nil, m)
+}
+
+// MarshalAppend appends the textproto format encoding of m to b,
+// returning the result.
+func (o MarshalOptions) MarshalAppend(b []byte, m proto.Message) ([]byte, error) {
+	return o.marshal(b, m)
 }
 
 // marshal is a centralized function that all marshal operations go through.
 // For profiling purposes, avoid changing the name of this function or
 // introducing other code paths for marshal that do not go through this.
-func (o MarshalOptions) marshal(m proto.Message) ([]byte, error) {
+func (o MarshalOptions) marshal(b []byte, m proto.Message) ([]byte, error) {
 	var delims = [2]byte{'{', '}'}
 
 	if o.Multiline && o.Indent == "" {
@@ -117,7 +123,7 @@
 		o.Resolver = protoregistry.GlobalTypes
 	}
 
-	internalEnc, err := text.NewEncoder(o.Indent, delims, o.EmitASCII)
+	internalEnc, err := text.NewEncoder(b, o.Indent, delims, o.EmitASCII)
 	if err != nil {
 		return nil, err
 	}
@@ -125,7 +131,7 @@
 	// Treat nil message interface as an empty message,
 	// in which case there is nothing to output.
 	if m == nil {
-		return []byte{}, nil
+		return b, nil
 	}
 
 	enc := encoder{internalEnc, o}
diff --git a/encoding/prototext/encode_test.go b/encoding/prototext/encode_test.go
index 96510bc..65b93cb 100644
--- a/encoding/prototext/encode_test.go
+++ b/encoding/prototext/encode_test.go
@@ -5,6 +5,7 @@
 package prototext_test
 
 import (
+	"bytes"
 	"math"
 	"testing"
 
@@ -1435,3 +1436,44 @@
 		})
 	}
 }
+
+func TestEncodeAppend(t *testing.T) {
+	want := []byte("prefix")
+	got := append([]byte(nil), want...)
+	got, err := prototext.MarshalOptions{}.MarshalAppend(got, &pb3.Scalars{
+		SString: "value",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.HasPrefix(got, want) {
+		t.Fatalf("MarshalAppend modified prefix: got %v, want prefix %v", got, want)
+	}
+}
+
+func TestMarshalAppendAllocations(t *testing.T) {
+	m := &pb3.Scalars{SInt32: 1}
+	const count = 1000
+	size := 9
+	b := make([]byte, size)
+	// AllocsPerRun returns an integral value.
+	marshalAllocs := testing.AllocsPerRun(count, func() {
+		_, err := prototext.MarshalOptions{}.MarshalAppend(b[:0], m)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+	b = nil
+	marshalAppendAllocs := testing.AllocsPerRun(count, func() {
+		var err error
+		b, err = prototext.MarshalOptions{}.MarshalAppend(b, m)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+	if marshalAllocs != marshalAppendAllocs {
+		t.Errorf("%v allocs/op when writing to a preallocated buffer", marshalAllocs)
+		t.Errorf("%v allocs/op when repeatedly appending to a slice", marshalAppendAllocs)
+		t.Errorf("expect amortized allocs/op to be identical")
+	}
+}
diff --git a/internal/encoding/json/encode.go b/internal/encoding/json/encode.go
index fbdf348..934f2dc 100644
--- a/internal/encoding/json/encode.go
+++ b/internal/encoding/json/encode.go
@@ -41,8 +41,10 @@
 //
 // If indent is a non-empty string, it causes every entry for an Array or Object
 // to be preceded by the indent and trailed by a newline.
-func NewEncoder(indent string) (*Encoder, error) {
-	e := &Encoder{}
+func NewEncoder(buf []byte, indent string) (*Encoder, error) {
+	e := &Encoder{
+		out: buf,
+	}
 	if len(indent) > 0 {
 		if strings.Trim(indent, " \t") != "" {
 			return nil, errors.New("indent may only be composed of space or tab characters")
@@ -176,13 +178,13 @@
 // WriteInt writes out the given signed integer in JSON number value.
 func (e *Encoder) WriteInt(n int64) {
 	e.prepareNext(scalar)
-	e.out = append(e.out, strconv.FormatInt(n, 10)...)
+	e.out = strconv.AppendInt(e.out, n, 10)
 }
 
 // WriteUint writes out the given unsigned integer in JSON number value.
 func (e *Encoder) WriteUint(n uint64) {
 	e.prepareNext(scalar)
-	e.out = append(e.out, strconv.FormatUint(n, 10)...)
+	e.out = strconv.AppendUint(e.out, n, 10)
 }
 
 // StartObject writes out the '{' symbol.
diff --git a/internal/encoding/json/encode_test.go b/internal/encoding/json/encode_test.go
index 5370f8b..c844b55 100644
--- a/internal/encoding/json/encode_test.go
+++ b/internal/encoding/json/encode_test.go
@@ -356,7 +356,7 @@
 	for _, tc := range tests {
 		t.Run(tc.desc, func(t *testing.T) {
 			if tc.wantOut != "" {
-				enc, err := json.NewEncoder("")
+				enc, err := json.NewEncoder(nil, "")
 				if err != nil {
 					t.Fatalf("NewEncoder() returned error: %v", err)
 				}
@@ -367,7 +367,7 @@
 				}
 			}
 			if tc.wantOutIndent != "" {
-				enc, err := json.NewEncoder("\t")
+				enc, err := json.NewEncoder(nil, "\t")
 				if err != nil {
 					t.Fatalf("NewEncoder() returned error: %v", err)
 				}
@@ -387,7 +387,7 @@
 
 	for _, in := range tests {
 		t.Run(in, func(t *testing.T) {
-			enc, err := json.NewEncoder("")
+			enc, err := json.NewEncoder(nil, "")
 			if err != nil {
 				t.Fatalf("NewEncoder() returned error: %v", err)
 			}
diff --git a/internal/encoding/text/encode.go b/internal/encoding/text/encode.go
index da289cc..cf7aed7 100644
--- a/internal/encoding/text/encode.go
+++ b/internal/encoding/text/encode.go
@@ -53,8 +53,10 @@
 // If outputASCII is true, strings will be serialized in such a way that
 // multi-byte UTF-8 sequences are escaped. This property ensures that the
 // overall output is ASCII (as opposed to UTF-8).
-func NewEncoder(indent string, delims [2]byte, outputASCII bool) (*Encoder, error) {
-	e := &Encoder{}
+func NewEncoder(buf []byte, indent string, delims [2]byte, outputASCII bool) (*Encoder, error) {
+	e := &Encoder{
+		encoderState: encoderState{out: buf},
+	}
 	if len(indent) > 0 {
 		if strings.Trim(indent, " \t") != "" {
 			return nil, errors.New("indent may only be composed of space and tab characters")
@@ -195,13 +197,13 @@
 // WriteInt writes out the given signed integer value.
 func (e *Encoder) WriteInt(n int64) {
 	e.prepareNext(scalar)
-	e.out = append(e.out, strconv.FormatInt(n, 10)...)
+	e.out = strconv.AppendInt(e.out, n, 10)
 }
 
 // WriteUint writes out the given unsigned integer value.
 func (e *Encoder) WriteUint(n uint64) {
 	e.prepareNext(scalar)
-	e.out = append(e.out, strconv.FormatUint(n, 10)...)
+	e.out = strconv.AppendUint(e.out, n, 10)
 }
 
 // WriteLiteral writes out the given string as a literal value without quotes.
diff --git a/internal/encoding/text/encode_test.go b/internal/encoding/text/encode_test.go
index af74690..d9c5009 100644
--- a/internal/encoding/text/encode_test.go
+++ b/internal/encoding/text/encode_test.go
@@ -341,7 +341,7 @@
 	t.Helper()
 
 	if tc.wantOut != "" {
-		enc, err := text.NewEncoder("", delims, false)
+		enc, err := text.NewEncoder(nil, "", delims, false)
 		if err != nil {
 			t.Fatalf("NewEncoder returned error: %v", err)
 		}
@@ -352,7 +352,7 @@
 		}
 	}
 	if tc.wantOutIndent != "" {
-		enc, err := text.NewEncoder("\t", delims, false)
+		enc, err := text.NewEncoder(nil, "\t", delims, false)
 		if err != nil {
 			t.Fatalf("NewEncoder returned error: %v", err)
 		}
@@ -520,7 +520,7 @@
 		charType = "ASCII"
 	}
 
-	enc, err := text.NewEncoder("", [2]byte{}, outputASCII)
+	enc, err := text.NewEncoder(nil, "", [2]byte{}, outputASCII)
 	if err != nil {
 		t.Fatalf("[%s] NewEncoder returned error: %v", charType, err)
 	}
@@ -532,7 +532,7 @@
 }
 
 func TestReset(t *testing.T) {
-	enc, err := text.NewEncoder("\t", [2]byte{}, false)
+	enc, err := text.NewEncoder(nil, "\t", [2]byte{}, false)
 	if err != nil {
 		t.Fatalf("NewEncoder returned error: %v", err)
 	}