init.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  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 internal_gengo
  5. import (
  6. "unicode"
  7. "unicode/utf8"
  8. "google.golang.org/protobuf/compiler/protogen"
  9. "google.golang.org/protobuf/encoding/protowire"
  10. "google.golang.org/protobuf/types/descriptorpb"
  11. )
  12. type fileInfo struct {
  13. *protogen.File
  14. allEnums []*enumInfo
  15. allMessages []*messageInfo
  16. allExtensions []*extensionInfo
  17. allEnumsByPtr map[*enumInfo]int // value is index into allEnums
  18. allMessagesByPtr map[*messageInfo]int // value is index into allMessages
  19. allMessageFieldsByPtr map[*messageInfo]*structFields
  20. // needRawDesc specifies whether the generator should emit logic to provide
  21. // the legacy raw descriptor in GZIP'd form.
  22. // This is updated by enum and message generation logic as necessary,
  23. // and checked at the end of file generation.
  24. needRawDesc bool
  25. }
  26. type structFields struct {
  27. count int
  28. unexported map[int]string
  29. }
  30. func (sf *structFields) append(name string) {
  31. if r, _ := utf8.DecodeRuneInString(name); !unicode.IsUpper(r) {
  32. if sf.unexported == nil {
  33. sf.unexported = make(map[int]string)
  34. }
  35. sf.unexported[sf.count] = name
  36. }
  37. sf.count++
  38. }
  39. func newFileInfo(file *protogen.File) *fileInfo {
  40. f := &fileInfo{File: file}
  41. // Collect all enums, messages, and extensions in "flattened ordering".
  42. // See filetype.TypeBuilder.
  43. var walkMessages func([]*protogen.Message, func(*protogen.Message))
  44. walkMessages = func(messages []*protogen.Message, f func(*protogen.Message)) {
  45. for _, m := range messages {
  46. f(m)
  47. walkMessages(m.Messages, f)
  48. }
  49. }
  50. initEnumInfos := func(enums []*protogen.Enum) {
  51. for _, enum := range enums {
  52. f.allEnums = append(f.allEnums, newEnumInfo(f, enum))
  53. }
  54. }
  55. initMessageInfos := func(messages []*protogen.Message) {
  56. for _, message := range messages {
  57. f.allMessages = append(f.allMessages, newMessageInfo(f, message))
  58. }
  59. }
  60. initExtensionInfos := func(extensions []*protogen.Extension) {
  61. for _, extension := range extensions {
  62. f.allExtensions = append(f.allExtensions, newExtensionInfo(f, extension))
  63. }
  64. }
  65. initEnumInfos(f.Enums)
  66. initMessageInfos(f.Messages)
  67. initExtensionInfos(f.Extensions)
  68. walkMessages(f.Messages, func(m *protogen.Message) {
  69. initEnumInfos(m.Enums)
  70. initMessageInfos(m.Messages)
  71. initExtensionInfos(m.Extensions)
  72. })
  73. // Derive a reverse mapping of enum and message pointers to their index
  74. // in allEnums and allMessages.
  75. if len(f.allEnums) > 0 {
  76. f.allEnumsByPtr = make(map[*enumInfo]int)
  77. for i, e := range f.allEnums {
  78. f.allEnumsByPtr[e] = i
  79. }
  80. }
  81. if len(f.allMessages) > 0 {
  82. f.allMessagesByPtr = make(map[*messageInfo]int)
  83. f.allMessageFieldsByPtr = make(map[*messageInfo]*structFields)
  84. for i, m := range f.allMessages {
  85. f.allMessagesByPtr[m] = i
  86. f.allMessageFieldsByPtr[m] = new(structFields)
  87. }
  88. }
  89. return f
  90. }
  91. type enumInfo struct {
  92. *protogen.Enum
  93. genJSONMethod bool
  94. genRawDescMethod bool
  95. }
  96. func newEnumInfo(f *fileInfo, enum *protogen.Enum) *enumInfo {
  97. e := &enumInfo{Enum: enum}
  98. e.genJSONMethod = true
  99. e.genRawDescMethod = true
  100. return e
  101. }
  102. type messageInfo struct {
  103. *protogen.Message
  104. genRawDescMethod bool
  105. genExtRangeMethod bool
  106. isTracked bool
  107. hasWeak bool
  108. }
  109. func newMessageInfo(f *fileInfo, message *protogen.Message) *messageInfo {
  110. m := &messageInfo{Message: message}
  111. m.genRawDescMethod = true
  112. m.genExtRangeMethod = true
  113. m.isTracked = isTrackedMessage(m)
  114. for _, field := range m.Fields {
  115. m.hasWeak = m.hasWeak || field.Desc.IsWeak()
  116. }
  117. return m
  118. }
  119. // isTrackedMessage reports whether field tracking is enabled on the message.
  120. func isTrackedMessage(m *messageInfo) (tracked bool) {
  121. const trackFieldUse_fieldNumber = 37383685
  122. // Decode the option from unknown fields to avoid a dependency on the
  123. // annotation proto from protoc-gen-go.
  124. b := m.Desc.Options().(*descriptorpb.MessageOptions).ProtoReflect().GetUnknown()
  125. for len(b) > 0 {
  126. num, typ, n := protowire.ConsumeTag(b)
  127. b = b[n:]
  128. if num == trackFieldUse_fieldNumber && typ == protowire.VarintType {
  129. v, _ := protowire.ConsumeVarint(b)
  130. tracked = protowire.DecodeBool(v)
  131. }
  132. m := protowire.ConsumeFieldValue(num, typ, b)
  133. b = b[m:]
  134. }
  135. return tracked
  136. }
  137. type extensionInfo struct {
  138. *protogen.Extension
  139. }
  140. func newExtensionInfo(f *fileInfo, extension *protogen.Extension) *extensionInfo {
  141. x := &extensionInfo{Extension: extension}
  142. return x
  143. }