diff --git a/proto/encode.go b/proto/encode.go index a57133c7..756082ca 100644 --- a/proto/encode.go +++ b/proto/encode.go @@ -101,7 +101,7 @@ func (o MarshalOptions) marshalMessage(b []byte, m protoreflect.Message) ([]byte !(o.Deterministic && methods.Flags&protoiface.SupportMarshalDeterministic == 0) { sz := methods.Size(m, protoiface.MarshalOptions(o)) if cap(b) < len(b)+sz { - x := make([]byte, len(b), len(b)+sz) + x := make([]byte, len(b), growcap(cap(b), len(b)+sz)) copy(x, b) b = x } @@ -111,6 +111,32 @@ func (o MarshalOptions) marshalMessage(b []byte, m protoreflect.Message) ([]byte return o.marshalMessageSlow(b, m) } +// growcap scales up the capacity of a slice. +// +// Given a slice with a current capacity of oldcap and a desired +// capacity of wantcap, growcap returns a new capacity >= wantcap. +// +// The algorithm is mostly identical to the one used by append as of Go 1.14. +func growcap(oldcap, wantcap int) (newcap int) { + if wantcap > oldcap*2 { + newcap = wantcap + } else if oldcap < 1024 { + // The Go 1.14 runtime takes this case when len(s) < 1024, + // not when cap(s) < 1024. The difference doesn't seem + // significant here. + newcap = oldcap * 2 + } else { + newcap = oldcap + for 0 < newcap && newcap < wantcap { + newcap += newcap / 4 + } + if newcap <= 0 { + newcap = wantcap + } + } + return newcap +} + func (o MarshalOptions) marshalMessageSlow(b []byte, m protoreflect.Message) ([]byte, error) { if messageset.IsMessageSet(m.Descriptor()) { return marshalMessageSet(b, m, o) diff --git a/proto/encode_test.go b/proto/encode_test.go index b5af4617..663612db 100644 --- a/proto/encode_test.go +++ b/proto/encode_test.go @@ -155,3 +155,30 @@ func TestEncodeOneofNilWrapper(t *testing.T) { t.Errorf("Marshal return non-empty, want empty") } } + +func TestMarshalAppendAllocations(t *testing.T) { + m := &test3pb.TestAllTypes{OptionalInt32: 1} + size := proto.Size(m) + const count = 1000 + b := make([]byte, size) + // AllocsPerRun returns an integral value. + marshalAllocs := testing.AllocsPerRun(count, func() { + _, err := proto.MarshalOptions{}.MarshalAppend(b[:0], m) + if err != nil { + t.Fatal(err) + } + }) + b = nil + marshalAppendAllocs := testing.AllocsPerRun(count, func() { + var err error + b, err = proto.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") + } +}