size.go 3.0 KB

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