size.go 3.1 KB

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