size.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674
  1. // Protocol Buffers for Go with Gadgets
  2. //
  3. // Copyright (c) 2013, The GoGo Authors. All rights reserved.
  4. // http://github.com/gogo/protobuf
  5. //
  6. // Redistribution and use in source and binary forms, with or without
  7. // modification, are permitted provided that the following conditions are
  8. // met:
  9. //
  10. // * Redistributions of source code must retain the above copyright
  11. // notice, this list of conditions and the following disclaimer.
  12. // * Redistributions in binary form must reproduce the above
  13. // copyright notice, this list of conditions and the following disclaimer
  14. // in the documentation and/or other materials provided with the
  15. // distribution.
  16. //
  17. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  18. // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  19. // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  20. // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  21. // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  22. // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  23. // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  24. // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  25. // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  26. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  27. // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  28. /*
  29. The size plugin generates a Size or ProtoSize method for each message.
  30. This is useful with the MarshalTo method generated by the marshalto plugin and the
  31. gogoproto.marshaler and gogoproto.marshaler_all extensions.
  32. It is enabled by the following extensions:
  33. - sizer
  34. - sizer_all
  35. - protosizer
  36. - protosizer_all
  37. The size plugin also generates a test given it is enabled using one of the following extensions:
  38. - testgen
  39. - testgen_all
  40. And a benchmark given it is enabled using one of the following extensions:
  41. - benchgen
  42. - benchgen_all
  43. Let us look at:
  44. github.com/gogo/protobuf/test/example/example.proto
  45. Btw all the output can be seen at:
  46. github.com/gogo/protobuf/test/example/*
  47. The following message:
  48. option (gogoproto.sizer_all) = true;
  49. message B {
  50. option (gogoproto.description) = true;
  51. optional A A = 1 [(gogoproto.nullable) = false, (gogoproto.embed) = true];
  52. repeated bytes G = 2 [(gogoproto.customtype) = "github.com/gogo/protobuf/test/custom.Uint128", (gogoproto.nullable) = false];
  53. }
  54. given to the size plugin, will generate the following code:
  55. func (m *B) Size() (n int) {
  56. var l int
  57. _ = l
  58. l = m.A.Size()
  59. n += 1 + l + sovExample(uint64(l))
  60. if len(m.G) > 0 {
  61. for _, e := range m.G {
  62. l = e.Size()
  63. n += 1 + l + sovExample(uint64(l))
  64. }
  65. }
  66. if m.XXX_unrecognized != nil {
  67. n += len(m.XXX_unrecognized)
  68. }
  69. return n
  70. }
  71. and the following test code:
  72. func TestBSize(t *testing5.T) {
  73. popr := math_rand5.New(math_rand5.NewSource(time5.Now().UnixNano()))
  74. p := NewPopulatedB(popr, true)
  75. dAtA, err := github_com_gogo_protobuf_proto2.Marshal(p)
  76. if err != nil {
  77. panic(err)
  78. }
  79. size := p.Size()
  80. if len(dAtA) != size {
  81. t.Fatalf("size %v != marshalled size %v", size, len(dAtA))
  82. }
  83. }
  84. func BenchmarkBSize(b *testing5.B) {
  85. popr := math_rand5.New(math_rand5.NewSource(616))
  86. total := 0
  87. pops := make([]*B, 1000)
  88. for i := 0; i < 1000; i++ {
  89. pops[i] = NewPopulatedB(popr, false)
  90. }
  91. b.ResetTimer()
  92. for i := 0; i < b.N; i++ {
  93. total += pops[i%1000].Size()
  94. }
  95. b.SetBytes(int64(total / b.N))
  96. }
  97. The sovExample function is a size of varint function for the example.pb.go file.
  98. */
  99. package size
  100. import (
  101. "fmt"
  102. "os"
  103. "strconv"
  104. "strings"
  105. "github.com/gogo/protobuf/gogoproto"
  106. "github.com/gogo/protobuf/proto"
  107. descriptor "github.com/gogo/protobuf/protoc-gen-gogo/descriptor"
  108. "github.com/gogo/protobuf/protoc-gen-gogo/generator"
  109. "github.com/gogo/protobuf/vanity"
  110. )
  111. type size struct {
  112. *generator.Generator
  113. generator.PluginImports
  114. atleastOne bool
  115. localName string
  116. typesPkg generator.Single
  117. }
  118. func NewSize() *size {
  119. return &size{}
  120. }
  121. func (p *size) Name() string {
  122. return "size"
  123. }
  124. func (p *size) Init(g *generator.Generator) {
  125. p.Generator = g
  126. }
  127. func wireToType(wire string) int {
  128. switch wire {
  129. case "fixed64":
  130. return proto.WireFixed64
  131. case "fixed32":
  132. return proto.WireFixed32
  133. case "varint":
  134. return proto.WireVarint
  135. case "bytes":
  136. return proto.WireBytes
  137. case "group":
  138. return proto.WireBytes
  139. case "zigzag32":
  140. return proto.WireVarint
  141. case "zigzag64":
  142. return proto.WireVarint
  143. }
  144. panic("unreachable")
  145. }
  146. func keySize(fieldNumber int32, wireType int) int {
  147. x := uint32(fieldNumber)<<3 | uint32(wireType)
  148. size := 0
  149. for size = 0; x > 127; size++ {
  150. x >>= 7
  151. }
  152. size++
  153. return size
  154. }
  155. func (p *size) sizeVarint() {
  156. p.P(`
  157. func sov`, p.localName, `(x uint64) (n int) {
  158. for {
  159. n++
  160. x >>= 7
  161. if x == 0 {
  162. break
  163. }
  164. }
  165. return n
  166. }`)
  167. }
  168. func (p *size) sizeZigZag() {
  169. p.P(`func soz`, p.localName, `(x uint64) (n int) {
  170. return sov`, p.localName, `(uint64((x << 1) ^ uint64((int64(x) >> 63))))
  171. }`)
  172. }
  173. func (p *size) std(field *descriptor.FieldDescriptorProto, name string) (string, bool) {
  174. if gogoproto.IsStdTime(field) {
  175. if gogoproto.IsNullable(field) {
  176. return p.typesPkg.Use() + `.SizeOfStdTime(*` + name + `)`, true
  177. } else {
  178. return p.typesPkg.Use() + `.SizeOfStdTime(` + name + `)`, true
  179. }
  180. } else if gogoproto.IsStdDuration(field) {
  181. if gogoproto.IsNullable(field) {
  182. return p.typesPkg.Use() + `.SizeOfStdDuration(*` + name + `)`, true
  183. } else {
  184. return p.typesPkg.Use() + `.SizeOfStdDuration(` + name + `)`, true
  185. }
  186. }
  187. return "", false
  188. }
  189. func (p *size) generateField(proto3 bool, file *generator.FileDescriptor, message *generator.Descriptor, field *descriptor.FieldDescriptorProto, sizeName string) {
  190. fieldname := p.GetOneOfFieldName(message, field)
  191. nullable := gogoproto.IsNullable(field)
  192. repeated := field.IsRepeated()
  193. doNilCheck := gogoproto.NeedsNilCheck(proto3, field)
  194. if repeated {
  195. p.P(`if len(m.`, fieldname, `) > 0 {`)
  196. p.In()
  197. } else if doNilCheck {
  198. p.P(`if m.`, fieldname, ` != nil {`)
  199. p.In()
  200. }
  201. packed := field.IsPacked() || (proto3 && field.IsPacked3())
  202. _, wire := p.GoType(message, field)
  203. wireType := wireToType(wire)
  204. fieldNumber := field.GetNumber()
  205. if packed {
  206. wireType = proto.WireBytes
  207. }
  208. key := keySize(fieldNumber, wireType)
  209. switch *field.Type {
  210. case descriptor.FieldDescriptorProto_TYPE_DOUBLE,
  211. descriptor.FieldDescriptorProto_TYPE_FIXED64,
  212. descriptor.FieldDescriptorProto_TYPE_SFIXED64:
  213. if packed {
  214. p.P(`n+=`, strconv.Itoa(key), `+sov`, p.localName, `(uint64(len(m.`, fieldname, `)*8))`, `+len(m.`, fieldname, `)*8`)
  215. } else if repeated {
  216. p.P(`n+=`, strconv.Itoa(key+8), `*len(m.`, fieldname, `)`)
  217. } else if proto3 {
  218. p.P(`if m.`, fieldname, ` != 0 {`)
  219. p.In()
  220. p.P(`n+=`, strconv.Itoa(key+8))
  221. p.Out()
  222. p.P(`}`)
  223. } else if nullable {
  224. p.P(`n+=`, strconv.Itoa(key+8))
  225. } else {
  226. p.P(`n+=`, strconv.Itoa(key+8))
  227. }
  228. case descriptor.FieldDescriptorProto_TYPE_FLOAT,
  229. descriptor.FieldDescriptorProto_TYPE_FIXED32,
  230. descriptor.FieldDescriptorProto_TYPE_SFIXED32:
  231. if packed {
  232. p.P(`n+=`, strconv.Itoa(key), `+sov`, p.localName, `(uint64(len(m.`, fieldname, `)*4))`, `+len(m.`, fieldname, `)*4`)
  233. } else if repeated {
  234. p.P(`n+=`, strconv.Itoa(key+4), `*len(m.`, fieldname, `)`)
  235. } else if proto3 {
  236. p.P(`if m.`, fieldname, ` != 0 {`)
  237. p.In()
  238. p.P(`n+=`, strconv.Itoa(key+4))
  239. p.Out()
  240. p.P(`}`)
  241. } else if nullable {
  242. p.P(`n+=`, strconv.Itoa(key+4))
  243. } else {
  244. p.P(`n+=`, strconv.Itoa(key+4))
  245. }
  246. case descriptor.FieldDescriptorProto_TYPE_INT64,
  247. descriptor.FieldDescriptorProto_TYPE_UINT64,
  248. descriptor.FieldDescriptorProto_TYPE_UINT32,
  249. descriptor.FieldDescriptorProto_TYPE_ENUM,
  250. descriptor.FieldDescriptorProto_TYPE_INT32:
  251. if packed {
  252. p.P(`l = 0`)
  253. p.P(`for _, e := range m.`, fieldname, ` {`)
  254. p.In()
  255. p.P(`l+=sov`, p.localName, `(uint64(e))`)
  256. p.Out()
  257. p.P(`}`)
  258. p.P(`n+=`, strconv.Itoa(key), `+sov`, p.localName, `(uint64(l))+l`)
  259. } else if repeated {
  260. p.P(`for _, e := range m.`, fieldname, ` {`)
  261. p.In()
  262. p.P(`n+=`, strconv.Itoa(key), `+sov`, p.localName, `(uint64(e))`)
  263. p.Out()
  264. p.P(`}`)
  265. } else if proto3 {
  266. p.P(`if m.`, fieldname, ` != 0 {`)
  267. p.In()
  268. p.P(`n+=`, strconv.Itoa(key), `+sov`, p.localName, `(uint64(m.`, fieldname, `))`)
  269. p.Out()
  270. p.P(`}`)
  271. } else if nullable {
  272. p.P(`n+=`, strconv.Itoa(key), `+sov`, p.localName, `(uint64(*m.`, fieldname, `))`)
  273. } else {
  274. p.P(`n+=`, strconv.Itoa(key), `+sov`, p.localName, `(uint64(m.`, fieldname, `))`)
  275. }
  276. case descriptor.FieldDescriptorProto_TYPE_BOOL:
  277. if packed {
  278. p.P(`n+=`, strconv.Itoa(key), `+sov`, p.localName, `(uint64(len(m.`, fieldname, `)))`, `+len(m.`, fieldname, `)*1`)
  279. } else if repeated {
  280. p.P(`n+=`, strconv.Itoa(key+1), `*len(m.`, fieldname, `)`)
  281. } else if proto3 {
  282. p.P(`if m.`, fieldname, ` {`)
  283. p.In()
  284. p.P(`n+=`, strconv.Itoa(key+1))
  285. p.Out()
  286. p.P(`}`)
  287. } else if nullable {
  288. p.P(`n+=`, strconv.Itoa(key+1))
  289. } else {
  290. p.P(`n+=`, strconv.Itoa(key+1))
  291. }
  292. case descriptor.FieldDescriptorProto_TYPE_STRING:
  293. if repeated {
  294. p.P(`for _, s := range m.`, fieldname, ` { `)
  295. p.In()
  296. p.P(`l = len(s)`)
  297. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  298. p.Out()
  299. p.P(`}`)
  300. } else if proto3 {
  301. p.P(`l=len(m.`, fieldname, `)`)
  302. p.P(`if l > 0 {`)
  303. p.In()
  304. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  305. p.Out()
  306. p.P(`}`)
  307. } else if nullable {
  308. p.P(`l=len(*m.`, fieldname, `)`)
  309. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  310. } else {
  311. p.P(`l=len(m.`, fieldname, `)`)
  312. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  313. }
  314. case descriptor.FieldDescriptorProto_TYPE_GROUP:
  315. panic(fmt.Errorf("size does not support group %v", fieldname))
  316. case descriptor.FieldDescriptorProto_TYPE_MESSAGE:
  317. if p.IsMap(field) {
  318. m := p.GoMapType(nil, field)
  319. _, keywire := p.GoType(nil, m.KeyAliasField)
  320. valuegoTyp, _ := p.GoType(nil, m.ValueField)
  321. valuegoAliasTyp, valuewire := p.GoType(nil, m.ValueAliasField)
  322. _, fieldwire := p.GoType(nil, field)
  323. nullable, valuegoTyp, valuegoAliasTyp = generator.GoMapValueTypes(field, m.ValueField, valuegoTyp, valuegoAliasTyp)
  324. fieldKeySize := keySize(field.GetNumber(), wireToType(fieldwire))
  325. keyKeySize := keySize(1, wireToType(keywire))
  326. valueKeySize := keySize(2, wireToType(valuewire))
  327. p.P(`for k, v := range m.`, fieldname, ` { `)
  328. p.In()
  329. p.P(`_ = k`)
  330. p.P(`_ = v`)
  331. sum := []string{strconv.Itoa(keyKeySize)}
  332. switch m.KeyField.GetType() {
  333. case descriptor.FieldDescriptorProto_TYPE_DOUBLE,
  334. descriptor.FieldDescriptorProto_TYPE_FIXED64,
  335. descriptor.FieldDescriptorProto_TYPE_SFIXED64:
  336. sum = append(sum, `8`)
  337. case descriptor.FieldDescriptorProto_TYPE_FLOAT,
  338. descriptor.FieldDescriptorProto_TYPE_FIXED32,
  339. descriptor.FieldDescriptorProto_TYPE_SFIXED32:
  340. sum = append(sum, `4`)
  341. case descriptor.FieldDescriptorProto_TYPE_INT64,
  342. descriptor.FieldDescriptorProto_TYPE_UINT64,
  343. descriptor.FieldDescriptorProto_TYPE_UINT32,
  344. descriptor.FieldDescriptorProto_TYPE_ENUM,
  345. descriptor.FieldDescriptorProto_TYPE_INT32:
  346. sum = append(sum, `sov`+p.localName+`(uint64(k))`)
  347. case descriptor.FieldDescriptorProto_TYPE_BOOL:
  348. sum = append(sum, `1`)
  349. case descriptor.FieldDescriptorProto_TYPE_STRING,
  350. descriptor.FieldDescriptorProto_TYPE_BYTES:
  351. sum = append(sum, `len(k)+sov`+p.localName+`(uint64(len(k)))`)
  352. case descriptor.FieldDescriptorProto_TYPE_SINT32,
  353. descriptor.FieldDescriptorProto_TYPE_SINT64:
  354. sum = append(sum, `soz`+p.localName+`(uint64(k))`)
  355. }
  356. switch m.ValueField.GetType() {
  357. case descriptor.FieldDescriptorProto_TYPE_DOUBLE,
  358. descriptor.FieldDescriptorProto_TYPE_FIXED64,
  359. descriptor.FieldDescriptorProto_TYPE_SFIXED64:
  360. sum = append(sum, strconv.Itoa(valueKeySize))
  361. sum = append(sum, strconv.Itoa(8))
  362. case descriptor.FieldDescriptorProto_TYPE_FLOAT,
  363. descriptor.FieldDescriptorProto_TYPE_FIXED32,
  364. descriptor.FieldDescriptorProto_TYPE_SFIXED32:
  365. sum = append(sum, strconv.Itoa(valueKeySize))
  366. sum = append(sum, strconv.Itoa(4))
  367. case descriptor.FieldDescriptorProto_TYPE_INT64,
  368. descriptor.FieldDescriptorProto_TYPE_UINT64,
  369. descriptor.FieldDescriptorProto_TYPE_UINT32,
  370. descriptor.FieldDescriptorProto_TYPE_ENUM,
  371. descriptor.FieldDescriptorProto_TYPE_INT32:
  372. sum = append(sum, strconv.Itoa(valueKeySize))
  373. sum = append(sum, `sov`+p.localName+`(uint64(v))`)
  374. case descriptor.FieldDescriptorProto_TYPE_BOOL:
  375. sum = append(sum, strconv.Itoa(valueKeySize))
  376. sum = append(sum, `1`)
  377. case descriptor.FieldDescriptorProto_TYPE_STRING:
  378. sum = append(sum, strconv.Itoa(valueKeySize))
  379. sum = append(sum, `len(v)+sov`+p.localName+`(uint64(len(v)))`)
  380. case descriptor.FieldDescriptorProto_TYPE_BYTES:
  381. if gogoproto.IsCustomType(field) {
  382. p.P(`l = 0`)
  383. if nullable {
  384. p.P(`if v != nil {`)
  385. p.In()
  386. }
  387. p.P(`l = v.`, sizeName, `()`)
  388. p.P(`l += `, strconv.Itoa(valueKeySize), ` + sov`+p.localName+`(uint64(l))`)
  389. if nullable {
  390. p.Out()
  391. p.P(`}`)
  392. }
  393. sum = append(sum, `l`)
  394. } else {
  395. p.P(`l = 0`)
  396. if proto3 {
  397. p.P(`if len(v) > 0 {`)
  398. } else {
  399. p.P(`if v != nil {`)
  400. }
  401. p.In()
  402. p.P(`l = `, strconv.Itoa(valueKeySize), ` + len(v)+sov`+p.localName+`(uint64(len(v)))`)
  403. p.Out()
  404. p.P(`}`)
  405. sum = append(sum, `l`)
  406. }
  407. case descriptor.FieldDescriptorProto_TYPE_SINT32,
  408. descriptor.FieldDescriptorProto_TYPE_SINT64:
  409. sum = append(sum, strconv.Itoa(valueKeySize))
  410. sum = append(sum, `soz`+p.localName+`(uint64(v))`)
  411. case descriptor.FieldDescriptorProto_TYPE_MESSAGE:
  412. stdSizeCall, stdOk := p.std(field, "v")
  413. if nullable {
  414. p.P(`l = 0`)
  415. p.P(`if v != nil {`)
  416. p.In()
  417. if stdOk {
  418. p.P(`l = `, stdSizeCall)
  419. } else if valuegoTyp != valuegoAliasTyp {
  420. p.P(`l = ((`, valuegoTyp, `)(v)).`, sizeName, `()`)
  421. } else {
  422. p.P(`l = v.`, sizeName, `()`)
  423. }
  424. p.P(`l += `, strconv.Itoa(valueKeySize), ` + sov`+p.localName+`(uint64(l))`)
  425. p.Out()
  426. p.P(`}`)
  427. sum = append(sum, `l`)
  428. } else {
  429. if stdOk {
  430. p.P(`l = `, stdSizeCall)
  431. } else if valuegoTyp != valuegoAliasTyp {
  432. p.P(`l = ((*`, valuegoTyp, `)(&v)).`, sizeName, `()`)
  433. } else {
  434. p.P(`l = v.`, sizeName, `()`)
  435. }
  436. sum = append(sum, strconv.Itoa(valueKeySize))
  437. sum = append(sum, `l+sov`+p.localName+`(uint64(l))`)
  438. }
  439. }
  440. p.P(`mapEntrySize := `, strings.Join(sum, "+"))
  441. p.P(`n+=mapEntrySize+`, fieldKeySize, `+sov`, p.localName, `(uint64(mapEntrySize))`)
  442. p.Out()
  443. p.P(`}`)
  444. } else if repeated {
  445. p.P(`for _, e := range m.`, fieldname, ` { `)
  446. p.In()
  447. stdSizeCall, stdOk := p.std(field, "e")
  448. if stdOk {
  449. p.P(`l=`, stdSizeCall)
  450. } else {
  451. p.P(`l=e.`, sizeName, `()`)
  452. }
  453. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  454. p.Out()
  455. p.P(`}`)
  456. } else {
  457. stdSizeCall, stdOk := p.std(field, "m."+fieldname)
  458. if stdOk {
  459. p.P(`l=`, stdSizeCall)
  460. } else {
  461. p.P(`l=m.`, fieldname, `.`, sizeName, `()`)
  462. }
  463. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  464. }
  465. case descriptor.FieldDescriptorProto_TYPE_BYTES:
  466. if !gogoproto.IsCustomType(field) {
  467. if repeated {
  468. p.P(`for _, b := range m.`, fieldname, ` { `)
  469. p.In()
  470. p.P(`l = len(b)`)
  471. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  472. p.Out()
  473. p.P(`}`)
  474. } else if proto3 {
  475. p.P(`l=len(m.`, fieldname, `)`)
  476. p.P(`if l > 0 {`)
  477. p.In()
  478. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  479. p.Out()
  480. p.P(`}`)
  481. } else {
  482. p.P(`l=len(m.`, fieldname, `)`)
  483. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  484. }
  485. } else {
  486. if repeated {
  487. p.P(`for _, e := range m.`, fieldname, ` { `)
  488. p.In()
  489. p.P(`l=e.`, sizeName, `()`)
  490. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  491. p.Out()
  492. p.P(`}`)
  493. } else {
  494. p.P(`l=m.`, fieldname, `.`, sizeName, `()`)
  495. p.P(`n+=`, strconv.Itoa(key), `+l+sov`, p.localName, `(uint64(l))`)
  496. }
  497. }
  498. case descriptor.FieldDescriptorProto_TYPE_SINT32,
  499. descriptor.FieldDescriptorProto_TYPE_SINT64:
  500. if packed {
  501. p.P(`l = 0`)
  502. p.P(`for _, e := range m.`, fieldname, ` {`)
  503. p.In()
  504. p.P(`l+=soz`, p.localName, `(uint64(e))`)
  505. p.Out()
  506. p.P(`}`)
  507. p.P(`n+=`, strconv.Itoa(key), `+sov`, p.localName, `(uint64(l))+l`)
  508. } else if repeated {
  509. p.P(`for _, e := range m.`, fieldname, ` {`)
  510. p.In()
  511. p.P(`n+=`, strconv.Itoa(key), `+soz`, p.localName, `(uint64(e))`)
  512. p.Out()
  513. p.P(`}`)
  514. } else if proto3 {
  515. p.P(`if m.`, fieldname, ` != 0 {`)
  516. p.In()
  517. p.P(`n+=`, strconv.Itoa(key), `+soz`, p.localName, `(uint64(m.`, fieldname, `))`)
  518. p.Out()
  519. p.P(`}`)
  520. } else if nullable {
  521. p.P(`n+=`, strconv.Itoa(key), `+soz`, p.localName, `(uint64(*m.`, fieldname, `))`)
  522. } else {
  523. p.P(`n+=`, strconv.Itoa(key), `+soz`, p.localName, `(uint64(m.`, fieldname, `))`)
  524. }
  525. default:
  526. panic("not implemented")
  527. }
  528. if repeated || doNilCheck {
  529. p.Out()
  530. p.P(`}`)
  531. }
  532. }
  533. func (p *size) Generate(file *generator.FileDescriptor) {
  534. p.PluginImports = generator.NewPluginImports(p.Generator)
  535. p.atleastOne = false
  536. p.localName = generator.FileName(file)
  537. p.typesPkg = p.NewImport("github.com/gogo/protobuf/types")
  538. protoPkg := p.NewImport("github.com/gogo/protobuf/proto")
  539. if !gogoproto.ImportsGoGoProto(file.FileDescriptorProto) {
  540. protoPkg = p.NewImport("github.com/golang/protobuf/proto")
  541. }
  542. for _, message := range file.Messages() {
  543. sizeName := ""
  544. if gogoproto.IsSizer(file.FileDescriptorProto, message.DescriptorProto) && gogoproto.IsProtoSizer(file.FileDescriptorProto, message.DescriptorProto) {
  545. fmt.Fprintf(os.Stderr, "ERROR: message %v cannot support both sizer and protosizer plugins\n", generator.CamelCase(*message.Name))
  546. os.Exit(1)
  547. }
  548. if gogoproto.IsSizer(file.FileDescriptorProto, message.DescriptorProto) {
  549. sizeName = "Size"
  550. } else if gogoproto.IsProtoSizer(file.FileDescriptorProto, message.DescriptorProto) {
  551. sizeName = "ProtoSize"
  552. } else {
  553. continue
  554. }
  555. if message.DescriptorProto.GetOptions().GetMapEntry() {
  556. continue
  557. }
  558. p.atleastOne = true
  559. ccTypeName := generator.CamelCaseSlice(message.TypeName())
  560. p.P(`func (m *`, ccTypeName, `) `, sizeName, `() (n int) {`)
  561. p.In()
  562. p.P(`var l int`)
  563. p.P(`_ = l`)
  564. oneofs := make(map[string]struct{})
  565. for _, field := range message.Field {
  566. oneof := field.OneofIndex != nil
  567. if !oneof {
  568. proto3 := gogoproto.IsProto3(file.FileDescriptorProto)
  569. p.generateField(proto3, file, message, field, sizeName)
  570. } else {
  571. fieldname := p.GetFieldName(message, field)
  572. if _, ok := oneofs[fieldname]; ok {
  573. continue
  574. } else {
  575. oneofs[fieldname] = struct{}{}
  576. }
  577. p.P(`if m.`, fieldname, ` != nil {`)
  578. p.In()
  579. p.P(`n+=m.`, fieldname, `.`, sizeName, `()`)
  580. p.Out()
  581. p.P(`}`)
  582. }
  583. }
  584. if message.DescriptorProto.HasExtension() {
  585. if gogoproto.HasExtensionsMap(file.FileDescriptorProto, message.DescriptorProto) {
  586. p.P(`n += `, protoPkg.Use(), `.SizeOfInternalExtension(m)`)
  587. } else {
  588. p.P(`if m.XXX_extensions != nil {`)
  589. p.In()
  590. p.P(`n+=len(m.XXX_extensions)`)
  591. p.Out()
  592. p.P(`}`)
  593. }
  594. }
  595. if gogoproto.HasUnrecognized(file.FileDescriptorProto, message.DescriptorProto) {
  596. p.P(`if m.XXX_unrecognized != nil {`)
  597. p.In()
  598. p.P(`n+=len(m.XXX_unrecognized)`)
  599. p.Out()
  600. p.P(`}`)
  601. }
  602. p.P(`return n`)
  603. p.Out()
  604. p.P(`}`)
  605. p.P()
  606. //Generate Size methods for oneof fields
  607. m := proto.Clone(message.DescriptorProto).(*descriptor.DescriptorProto)
  608. for _, f := range m.Field {
  609. oneof := f.OneofIndex != nil
  610. if !oneof {
  611. continue
  612. }
  613. ccTypeName := p.OneOfTypeName(message, f)
  614. p.P(`func (m *`, ccTypeName, `) `, sizeName, `() (n int) {`)
  615. p.In()
  616. p.P(`var l int`)
  617. p.P(`_ = l`)
  618. vanity.TurnOffNullableForNativeTypes(f)
  619. p.generateField(false, file, message, f, sizeName)
  620. p.P(`return n`)
  621. p.Out()
  622. p.P(`}`)
  623. }
  624. }
  625. if !p.atleastOne {
  626. return
  627. }
  628. p.sizeVarint()
  629. p.sizeZigZag()
  630. }
  631. func init() {
  632. generator.RegisterPlugin(NewSize())
  633. }