encode.go 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201
  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 impl
  5. import (
  6. "math"
  7. "sort"
  8. "sync/atomic"
  9. "google.golang.org/protobuf/internal/flags"
  10. proto "google.golang.org/protobuf/proto"
  11. piface "google.golang.org/protobuf/runtime/protoiface"
  12. )
  13. type marshalOptions struct {
  14. flags piface.MarshalInputFlags
  15. }
  16. func (o marshalOptions) Options() proto.MarshalOptions {
  17. return proto.MarshalOptions{
  18. AllowPartial: true,
  19. Deterministic: o.Deterministic(),
  20. UseCachedSize: o.UseCachedSize(),
  21. }
  22. }
  23. func (o marshalOptions) Deterministic() bool { return o.flags&piface.MarshalDeterministic != 0 }
  24. func (o marshalOptions) UseCachedSize() bool { return o.flags&piface.MarshalUseCachedSize != 0 }
  25. // size is protoreflect.Methods.Size.
  26. func (mi *MessageInfo) size(in piface.SizeInput) piface.SizeOutput {
  27. var p pointer
  28. if ms, ok := in.Message.(*messageState); ok {
  29. p = ms.pointer()
  30. } else {
  31. p = in.Message.(*messageReflectWrapper).pointer()
  32. }
  33. size := mi.sizePointer(p, marshalOptions{
  34. flags: in.Flags,
  35. })
  36. return piface.SizeOutput{Size: size}
  37. }
  38. func (mi *MessageInfo) sizePointer(p pointer, opts marshalOptions) (size int) {
  39. mi.init()
  40. if p.IsNil() {
  41. return 0
  42. }
  43. if opts.UseCachedSize() && mi.sizecacheOffset.IsValid() {
  44. if size := atomic.LoadInt32(p.Apply(mi.sizecacheOffset).Int32()); size >= 0 {
  45. return int(size)
  46. }
  47. }
  48. return mi.sizePointerSlow(p, opts)
  49. }
  50. func (mi *MessageInfo) sizePointerSlow(p pointer, opts marshalOptions) (size int) {
  51. if flags.ProtoLegacy && mi.isMessageSet {
  52. size = sizeMessageSet(mi, p, opts)
  53. if mi.sizecacheOffset.IsValid() {
  54. atomic.StoreInt32(p.Apply(mi.sizecacheOffset).Int32(), int32(size))
  55. }
  56. return size
  57. }
  58. if mi.extensionOffset.IsValid() {
  59. e := p.Apply(mi.extensionOffset).Extensions()
  60. size += mi.sizeExtensions(e, opts)
  61. }
  62. for _, f := range mi.orderedCoderFields {
  63. if f.funcs.size == nil {
  64. continue
  65. }
  66. fptr := p.Apply(f.offset)
  67. if f.isPointer && fptr.Elem().IsNil() {
  68. continue
  69. }
  70. size += f.funcs.size(fptr, f, opts)
  71. }
  72. if mi.unknownOffset.IsValid() {
  73. if u := mi.getUnknownBytes(p); u != nil {
  74. size += len(*u)
  75. }
  76. }
  77. if mi.sizecacheOffset.IsValid() {
  78. if size > math.MaxInt32 {
  79. // The size is too large for the int32 sizecache field.
  80. // We will need to recompute the size when encoding;
  81. // unfortunately expensive, but better than invalid output.
  82. atomic.StoreInt32(p.Apply(mi.sizecacheOffset).Int32(), -1)
  83. } else {
  84. atomic.StoreInt32(p.Apply(mi.sizecacheOffset).Int32(), int32(size))
  85. }
  86. }
  87. return size
  88. }
  89. // marshal is protoreflect.Methods.Marshal.
  90. func (mi *MessageInfo) marshal(in piface.MarshalInput) (out piface.MarshalOutput, err error) {
  91. var p pointer
  92. if ms, ok := in.Message.(*messageState); ok {
  93. p = ms.pointer()
  94. } else {
  95. p = in.Message.(*messageReflectWrapper).pointer()
  96. }
  97. b, err := mi.marshalAppendPointer(in.Buf, p, marshalOptions{
  98. flags: in.Flags,
  99. })
  100. return piface.MarshalOutput{Buf: b}, err
  101. }
  102. func (mi *MessageInfo) marshalAppendPointer(b []byte, p pointer, opts marshalOptions) ([]byte, error) {
  103. mi.init()
  104. if p.IsNil() {
  105. return b, nil
  106. }
  107. if flags.ProtoLegacy && mi.isMessageSet {
  108. return marshalMessageSet(mi, b, p, opts)
  109. }
  110. var err error
  111. // The old marshaler encodes extensions at beginning.
  112. if mi.extensionOffset.IsValid() {
  113. e := p.Apply(mi.extensionOffset).Extensions()
  114. // TODO: Special handling for MessageSet?
  115. b, err = mi.appendExtensions(b, e, opts)
  116. if err != nil {
  117. return b, err
  118. }
  119. }
  120. for _, f := range mi.orderedCoderFields {
  121. if f.funcs.marshal == nil {
  122. continue
  123. }
  124. fptr := p.Apply(f.offset)
  125. if f.isPointer && fptr.Elem().IsNil() {
  126. continue
  127. }
  128. b, err = f.funcs.marshal(b, fptr, f, opts)
  129. if err != nil {
  130. return b, err
  131. }
  132. }
  133. if mi.unknownOffset.IsValid() && !mi.isMessageSet {
  134. if u := mi.getUnknownBytes(p); u != nil {
  135. b = append(b, (*u)...)
  136. }
  137. }
  138. return b, nil
  139. }
  140. func (mi *MessageInfo) sizeExtensions(ext *map[int32]ExtensionField, opts marshalOptions) (n int) {
  141. if ext == nil {
  142. return 0
  143. }
  144. for _, x := range *ext {
  145. xi := getExtensionFieldInfo(x.Type())
  146. if xi.funcs.size == nil {
  147. continue
  148. }
  149. n += xi.funcs.size(x.Value(), xi.tagsize, opts)
  150. }
  151. return n
  152. }
  153. func (mi *MessageInfo) appendExtensions(b []byte, ext *map[int32]ExtensionField, opts marshalOptions) ([]byte, error) {
  154. if ext == nil {
  155. return b, nil
  156. }
  157. switch len(*ext) {
  158. case 0:
  159. return b, nil
  160. case 1:
  161. // Fast-path for one extension: Don't bother sorting the keys.
  162. var err error
  163. for _, x := range *ext {
  164. xi := getExtensionFieldInfo(x.Type())
  165. b, err = xi.funcs.marshal(b, x.Value(), xi.wiretag, opts)
  166. }
  167. return b, err
  168. default:
  169. // Sort the keys to provide a deterministic encoding.
  170. // Not sure this is required, but the old code does it.
  171. keys := make([]int, 0, len(*ext))
  172. for k := range *ext {
  173. keys = append(keys, int(k))
  174. }
  175. sort.Ints(keys)
  176. var err error
  177. for _, k := range keys {
  178. x := (*ext)[int32(k)]
  179. xi := getExtensionFieldInfo(x.Type())
  180. b, err = xi.funcs.marshal(b, x.Value(), xi.wiretag, opts)
  181. if err != nil {
  182. return b, err
  183. }
  184. }
  185. return b, nil
  186. }
  187. }