protobuf-go/proto/size.go
Damien Neil b0d217f664 proto, internal/impl: don't create fast path Size for legacy Marshalers
Implementations of the legacy Marshaler type have no way to efficiently
compute the size of the message. Rather than generating an inefficient
fast-path Size method which marshals the message and examines the
length of the result, don't generate a fast-path at all.

Drop the requirement that a fast-path MarshalAppend requires a
corresponding Size.

Avoids O(N^2) behavior when marshaling a legacy Marshaler that
recursively calls proto.Marshal.

Change-Id: I4793cf32275d08f29c8e1a1a44a193d9a5724058
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/213443
Reviewed-by: Joe Tsai <joetsai@google.com>
2020-01-06 22:47:37 +00:00

85 lines
2.5 KiB
Go

// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package proto
import (
"google.golang.org/protobuf/internal/encoding/messageset"
"google.golang.org/protobuf/internal/encoding/wire"
"google.golang.org/protobuf/reflect/protoreflect"
"google.golang.org/protobuf/runtime/protoiface"
)
// Size returns the size in bytes of the wire-format encoding of m.
func Size(m Message) int {
return MarshalOptions{}.Size(m)
}
// Size returns the size in bytes of the wire-format encoding of m.
func (o MarshalOptions) Size(m Message) int {
return sizeMessage(m.ProtoReflect())
}
func sizeMessage(m protoreflect.Message) (size int) {
methods := protoMethods(m)
if methods != nil && methods.Size != nil {
return methods.Size(m, protoiface.MarshalOptions{})
}
if methods != nil && methods.MarshalAppend != nil {
// This is not efficient, but we don't have any choice.
// This case is mainly used for legacy types with a Marshal method.
b, _ := methods.MarshalAppend(nil, m, protoiface.MarshalOptions{})
return len(b)
}
return sizeMessageSlow(m)
}
func sizeMessageSlow(m protoreflect.Message) (size int) {
if messageset.IsMessageSet(m.Descriptor()) {
return sizeMessageSet(m)
}
m.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool {
size += sizeField(fd, v)
return true
})
size += len(m.GetUnknown())
return size
}
func sizeField(fd protoreflect.FieldDescriptor, value protoreflect.Value) (size int) {
num := fd.Number()
switch {
case fd.IsList():
return sizeList(num, fd, value.List())
case fd.IsMap():
return sizeMap(num, fd, value.Map())
default:
return wire.SizeTag(num) + sizeSingular(num, fd.Kind(), value)
}
}
func sizeList(num wire.Number, fd protoreflect.FieldDescriptor, list protoreflect.List) (size int) {
if fd.IsPacked() && list.Len() > 0 {
content := 0
for i, llen := 0, list.Len(); i < llen; i++ {
content += sizeSingular(num, fd.Kind(), list.Get(i))
}
return wire.SizeTag(num) + wire.SizeBytes(content)
}
for i, llen := 0, list.Len(); i < llen; i++ {
size += wire.SizeTag(num) + sizeSingular(num, fd.Kind(), list.Get(i))
}
return size
}
func sizeMap(num wire.Number, fd protoreflect.FieldDescriptor, mapv protoreflect.Map) (size int) {
mapv.Range(func(key protoreflect.MapKey, value protoreflect.Value) bool {
size += wire.SizeTag(num)
size += wire.SizeBytes(sizeField(fd.MapKey(), key.Value()) + sizeField(fd.MapValue(), value))
return true
})
return size
}