protobuf-go/proto/size.go
Damien Neil d30e561d9e proto: add MarshalState, UnmarshalState
Add functions to the proto package which plumb through the fast-path state.

As a sample use case: A followup CL adds an Initialized field to
protoiface.UnmarshalOutput, permitting the unmarshaller to report back
when it can confirm that a message is fully initialized. We want to
preserve that information when an unmarshal operation threads through
the proto package (such as when unmarshaling extensions).

To allow these functions to be added as methods of MarshalOptions and
UnmarshalOptions rather than top-level functions, separate the options
from the input structs.

Also update options passed to fast-path methods to set AllowPartial and
Merge to reflect the expected behavior of those methods. (Always allow
partial, never merge.)

Change-Id: I482477b0c9340793be533e75a86d0bb88708716a
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/215877
Reviewed-by: Joe Tsai <joetsai@google.com>
2020-01-22 20:52:17 +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.Marshal != nil {
// This is not efficient, but we don't have any choice.
// This case is mainly used for legacy types with a Marshal method.
out, _ := methods.Marshal(m, protoiface.MarshalInput{}, protoiface.MarshalOptions{})
return len(out.Buf)
}
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
}