size.go 3.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  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. if fd.IsPacked() && list.Len() > 0 {
  68. content := 0
  69. for i, llen := 0, list.Len(); i < llen; i++ {
  70. content += o.sizeSingular(num, fd.Kind(), list.Get(i))
  71. }
  72. return protowire.SizeTag(num) + protowire.SizeBytes(content)
  73. }
  74. for i, llen := 0, list.Len(); i < llen; i++ {
  75. size += protowire.SizeTag(num) + o.sizeSingular(num, fd.Kind(), list.Get(i))
  76. }
  77. return size
  78. }
  79. func (o MarshalOptions) sizeMap(num protowire.Number, fd protoreflect.FieldDescriptor, mapv protoreflect.Map) (size int) {
  80. mapv.Range(func(key protoreflect.MapKey, value protoreflect.Value) bool {
  81. size += protowire.SizeTag(num)
  82. size += protowire.SizeBytes(o.sizeField(fd.MapKey(), key.Value()) + o.sizeField(fd.MapValue(), value))
  83. return true
  84. })
  85. return size
  86. }