David K. Bainbridge | bd6b288 | 2021-08-26 13:31:02 +0000 | [diff] [blame] | 1 | // Copyright 2019 The Go Authors. All rights reserved. |
| 2 | // Use of this source code is governed by a BSD-style |
| 3 | // license that can be found in the LICENSE file. |
| 4 | |
| 5 | package proto |
| 6 | |
| 7 | import ( |
| 8 | "google.golang.org/protobuf/encoding/protowire" |
| 9 | "google.golang.org/protobuf/internal/encoding/messageset" |
| 10 | "google.golang.org/protobuf/reflect/protoreflect" |
| 11 | "google.golang.org/protobuf/runtime/protoiface" |
| 12 | ) |
| 13 | |
| 14 | // Size returns the size in bytes of the wire-format encoding of m. |
| 15 | func Size(m Message) int { |
| 16 | return MarshalOptions{}.Size(m) |
| 17 | } |
| 18 | |
| 19 | // Size returns the size in bytes of the wire-format encoding of m. |
| 20 | func (o MarshalOptions) Size(m Message) int { |
| 21 | // Treat a nil message interface as an empty message; nothing to output. |
| 22 | if m == nil { |
| 23 | return 0 |
| 24 | } |
| 25 | |
| 26 | return o.size(m.ProtoReflect()) |
| 27 | } |
| 28 | |
| 29 | // size is a centralized function that all size operations go through. |
| 30 | // For profiling purposes, avoid changing the name of this function or |
| 31 | // introducing other code paths for size that do not go through this. |
| 32 | func (o MarshalOptions) size(m protoreflect.Message) (size int) { |
| 33 | methods := protoMethods(m) |
| 34 | if methods != nil && methods.Size != nil { |
| 35 | out := methods.Size(protoiface.SizeInput{ |
| 36 | Message: m, |
| 37 | }) |
| 38 | return out.Size |
| 39 | } |
| 40 | if methods != nil && methods.Marshal != nil { |
| 41 | // This is not efficient, but we don't have any choice. |
| 42 | // This case is mainly used for legacy types with a Marshal method. |
| 43 | out, _ := methods.Marshal(protoiface.MarshalInput{ |
| 44 | Message: m, |
| 45 | }) |
| 46 | return len(out.Buf) |
| 47 | } |
| 48 | return o.sizeMessageSlow(m) |
| 49 | } |
| 50 | |
| 51 | func (o MarshalOptions) sizeMessageSlow(m protoreflect.Message) (size int) { |
| 52 | if messageset.IsMessageSet(m.Descriptor()) { |
| 53 | return o.sizeMessageSet(m) |
| 54 | } |
| 55 | m.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool { |
| 56 | size += o.sizeField(fd, v) |
| 57 | return true |
| 58 | }) |
| 59 | size += len(m.GetUnknown()) |
| 60 | return size |
| 61 | } |
| 62 | |
| 63 | func (o MarshalOptions) sizeField(fd protoreflect.FieldDescriptor, value protoreflect.Value) (size int) { |
| 64 | num := fd.Number() |
| 65 | switch { |
| 66 | case fd.IsList(): |
| 67 | return o.sizeList(num, fd, value.List()) |
| 68 | case fd.IsMap(): |
| 69 | return o.sizeMap(num, fd, value.Map()) |
| 70 | default: |
| 71 | return protowire.SizeTag(num) + o.sizeSingular(num, fd.Kind(), value) |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | func (o MarshalOptions) sizeList(num protowire.Number, fd protoreflect.FieldDescriptor, list protoreflect.List) (size int) { |
| 76 | if fd.IsPacked() && list.Len() > 0 { |
| 77 | content := 0 |
| 78 | for i, llen := 0, list.Len(); i < llen; i++ { |
| 79 | content += o.sizeSingular(num, fd.Kind(), list.Get(i)) |
| 80 | } |
| 81 | return protowire.SizeTag(num) + protowire.SizeBytes(content) |
| 82 | } |
| 83 | |
| 84 | for i, llen := 0, list.Len(); i < llen; i++ { |
| 85 | size += protowire.SizeTag(num) + o.sizeSingular(num, fd.Kind(), list.Get(i)) |
| 86 | } |
| 87 | return size |
| 88 | } |
| 89 | |
| 90 | func (o MarshalOptions) sizeMap(num protowire.Number, fd protoreflect.FieldDescriptor, mapv protoreflect.Map) (size int) { |
| 91 | mapv.Range(func(key protoreflect.MapKey, value protoreflect.Value) bool { |
| 92 | size += protowire.SizeTag(num) |
| 93 | size += protowire.SizeBytes(o.sizeField(fd.MapKey(), key.Value()) + o.sizeField(fd.MapValue(), value)) |
| 94 | return true |
| 95 | }) |
| 96 | return size |
| 97 | } |