amit.ghosh | 258d14c | 2020-10-02 15:13:38 +0200 | [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 sizeMessage(m.ProtoReflect()) |
| 27 | } |
| 28 | |
| 29 | func sizeMessage(m protoreflect.Message) (size int) { |
| 30 | methods := protoMethods(m) |
| 31 | if methods != nil && methods.Size != nil { |
| 32 | out := methods.Size(protoiface.SizeInput{ |
| 33 | Message: m, |
| 34 | }) |
| 35 | return out.Size |
| 36 | } |
| 37 | if methods != nil && methods.Marshal != nil { |
| 38 | // This is not efficient, but we don't have any choice. |
| 39 | // This case is mainly used for legacy types with a Marshal method. |
| 40 | out, _ := methods.Marshal(protoiface.MarshalInput{ |
| 41 | Message: m, |
| 42 | }) |
| 43 | return len(out.Buf) |
| 44 | } |
| 45 | return sizeMessageSlow(m) |
| 46 | } |
| 47 | |
| 48 | func sizeMessageSlow(m protoreflect.Message) (size int) { |
| 49 | if messageset.IsMessageSet(m.Descriptor()) { |
| 50 | return sizeMessageSet(m) |
| 51 | } |
| 52 | m.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool { |
| 53 | size += sizeField(fd, v) |
| 54 | return true |
| 55 | }) |
| 56 | size += len(m.GetUnknown()) |
| 57 | return size |
| 58 | } |
| 59 | |
| 60 | func sizeField(fd protoreflect.FieldDescriptor, value protoreflect.Value) (size int) { |
| 61 | num := fd.Number() |
| 62 | switch { |
| 63 | case fd.IsList(): |
| 64 | return sizeList(num, fd, value.List()) |
| 65 | case fd.IsMap(): |
| 66 | return sizeMap(num, fd, value.Map()) |
| 67 | default: |
| 68 | return protowire.SizeTag(num) + sizeSingular(num, fd.Kind(), value) |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | func sizeList(num protowire.Number, fd protoreflect.FieldDescriptor, list protoreflect.List) (size int) { |
| 73 | if fd.IsPacked() && list.Len() > 0 { |
| 74 | content := 0 |
| 75 | for i, llen := 0, list.Len(); i < llen; i++ { |
| 76 | content += sizeSingular(num, fd.Kind(), list.Get(i)) |
| 77 | } |
| 78 | return protowire.SizeTag(num) + protowire.SizeBytes(content) |
| 79 | } |
| 80 | |
| 81 | for i, llen := 0, list.Len(); i < llen; i++ { |
| 82 | size += protowire.SizeTag(num) + sizeSingular(num, fd.Kind(), list.Get(i)) |
| 83 | } |
| 84 | return size |
| 85 | } |
| 86 | |
| 87 | func sizeMap(num protowire.Number, fd protoreflect.FieldDescriptor, mapv protoreflect.Map) (size int) { |
| 88 | mapv.Range(func(key protoreflect.MapKey, value protoreflect.Value) bool { |
| 89 | size += protowire.SizeTag(num) |
| 90 | size += protowire.SizeBytes(sizeField(fd.MapKey(), key.Value()) + sizeField(fd.MapValue(), value)) |
| 91 | return true |
| 92 | }) |
| 93 | return size |
| 94 | } |