1// Copyright 2018 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 5package internal_gengo 6 7import ( 8 "fmt" 9 "math" 10 "strings" 11 "unicode/utf8" 12 13 "google.golang.org/protobuf/compiler/protogen" 14 "google.golang.org/protobuf/proto" 15 "google.golang.org/protobuf/reflect/protopath" 16 "google.golang.org/protobuf/reflect/protorange" 17 "google.golang.org/protobuf/reflect/protoreflect" 18 19 "google.golang.org/protobuf/types/descriptorpb" 20) 21 22func genReflectFileDescriptor(gen *protogen.Plugin, g *protogen.GeneratedFile, f *fileInfo) { 23 g.P("var ", f.GoDescriptorIdent, " ", protoreflectPackage.Ident("FileDescriptor")) 24 g.P() 25 26 genFileDescriptor(gen, g, f) 27 if len(f.allEnums) > 0 { 28 g.P("var ", enumTypesVarName(f), " = make([]", protoimplPackage.Ident("EnumInfo"), ",", len(f.allEnums), ")") 29 } 30 if len(f.allMessages) > 0 { 31 g.P("var ", messageTypesVarName(f), " = make([]", protoimplPackage.Ident("MessageInfo"), ",", len(f.allMessages), ")") 32 } 33 34 // Generate a unique list of Go types for all declarations and dependencies, 35 // and the associated index into the type list for all dependencies. 36 var goTypes []string 37 var depIdxs []string 38 seen := map[protoreflect.FullName]int{} 39 genDep := func(name protoreflect.FullName, depSource string) { 40 if depSource != "" { 41 line := fmt.Sprintf("%d, // %d: %s -> %s", seen[name], len(depIdxs), depSource, name) 42 depIdxs = append(depIdxs, line) 43 } 44 } 45 genEnum := func(e *protogen.Enum, depSource string) { 46 if e != nil { 47 name := e.Desc.FullName() 48 if _, ok := seen[name]; !ok { 49 line := fmt.Sprintf("(%s)(0), // %d: %s", g.QualifiedGoIdent(e.GoIdent), len(goTypes), name) 50 goTypes = append(goTypes, line) 51 seen[name] = len(seen) 52 } 53 if depSource != "" { 54 genDep(name, depSource) 55 } 56 } 57 } 58 genMessage := func(m *protogen.Message, depSource string) { 59 if m != nil { 60 name := m.Desc.FullName() 61 if _, ok := seen[name]; !ok { 62 line := fmt.Sprintf("(*%s)(nil), // %d: %s", g.QualifiedGoIdent(m.GoIdent), len(goTypes), name) 63 if m.Desc.IsMapEntry() { 64 // Map entry messages have no associated Go type. 65 line = fmt.Sprintf("nil, // %d: %s", len(goTypes), name) 66 } 67 goTypes = append(goTypes, line) 68 seen[name] = len(seen) 69 } 70 if depSource != "" { 71 genDep(name, depSource) 72 } 73 } 74 } 75 76 // This ordering is significant. 77 // See filetype.TypeBuilder.DependencyIndexes. 78 type offsetEntry struct { 79 start int 80 name string 81 } 82 var depOffsets []offsetEntry 83 for _, enum := range f.allEnums { 84 genEnum(enum.Enum, "") 85 } 86 for _, message := range f.allMessages { 87 genMessage(message.Message, "") 88 } 89 depOffsets = append(depOffsets, offsetEntry{len(depIdxs), "field type_name"}) 90 for _, message := range f.allMessages { 91 for _, field := range message.Fields { 92 if field.Desc.IsWeak() { 93 continue 94 } 95 source := string(field.Desc.FullName()) 96 genEnum(field.Enum, source+":type_name") 97 genMessage(field.Message, source+":type_name") 98 } 99 } 100 depOffsets = append(depOffsets, offsetEntry{len(depIdxs), "extension extendee"}) 101 for _, extension := range f.allExtensions { 102 source := string(extension.Desc.FullName()) 103 genMessage(extension.Extendee, source+":extendee") 104 } 105 depOffsets = append(depOffsets, offsetEntry{len(depIdxs), "extension type_name"}) 106 for _, extension := range f.allExtensions { 107 source := string(extension.Desc.FullName()) 108 genEnum(extension.Enum, source+":type_name") 109 genMessage(extension.Message, source+":type_name") 110 } 111 depOffsets = append(depOffsets, offsetEntry{len(depIdxs), "method input_type"}) 112 for _, service := range f.Services { 113 for _, method := range service.Methods { 114 source := string(method.Desc.FullName()) 115 genMessage(method.Input, source+":input_type") 116 } 117 } 118 depOffsets = append(depOffsets, offsetEntry{len(depIdxs), "method output_type"}) 119 for _, service := range f.Services { 120 for _, method := range service.Methods { 121 source := string(method.Desc.FullName()) 122 genMessage(method.Output, source+":output_type") 123 } 124 } 125 depOffsets = append(depOffsets, offsetEntry{len(depIdxs), ""}) 126 for i := len(depOffsets) - 2; i >= 0; i-- { 127 curr, next := depOffsets[i], depOffsets[i+1] 128 depIdxs = append(depIdxs, fmt.Sprintf("%d, // [%d:%d] is the sub-list for %s", 129 curr.start, curr.start, next.start, curr.name)) 130 } 131 if len(depIdxs) > math.MaxInt32 { 132 panic("too many dependencies") // sanity check 133 } 134 135 g.P("var ", goTypesVarName(f), " = []interface{}{") 136 for _, s := range goTypes { 137 g.P(s) 138 } 139 g.P("}") 140 141 g.P("var ", depIdxsVarName(f), " = []int32{") 142 for _, s := range depIdxs { 143 g.P(s) 144 } 145 g.P("}") 146 147 g.P("func init() { ", initFuncName(f.File), "() }") 148 149 g.P("func ", initFuncName(f.File), "() {") 150 g.P("if ", f.GoDescriptorIdent, " != nil {") 151 g.P("return") 152 g.P("}") 153 154 // Ensure that initialization functions for different files in the same Go 155 // package run in the correct order: Call the init funcs for every .proto file 156 // imported by this one that is in the same Go package. 157 for i, imps := 0, f.Desc.Imports(); i < imps.Len(); i++ { 158 impFile := gen.FilesByPath[imps.Get(i).Path()] 159 if impFile.GoImportPath != f.GoImportPath { 160 continue 161 } 162 g.P(initFuncName(impFile), "()") 163 } 164 165 if len(f.allMessages) > 0 { 166 // Populate MessageInfo.Exporters. 167 g.P("if !", protoimplPackage.Ident("UnsafeEnabled"), " {") 168 for _, message := range f.allMessages { 169 if sf := f.allMessageFieldsByPtr[message]; len(sf.unexported) > 0 { 170 idx := f.allMessagesByPtr[message] 171 typesVar := messageTypesVarName(f) 172 173 g.P(typesVar, "[", idx, "].Exporter = func(v interface{}, i int) interface{} {") 174 g.P("switch v := v.(*", message.GoIdent, "); i {") 175 for i := 0; i < sf.count; i++ { 176 if name := sf.unexported[i]; name != "" { 177 g.P("case ", i, ": return &v.", name) 178 } 179 } 180 g.P("default: return nil") 181 g.P("}") 182 g.P("}") 183 } 184 } 185 g.P("}") 186 187 // Populate MessageInfo.OneofWrappers. 188 for _, message := range f.allMessages { 189 if len(message.Oneofs) > 0 { 190 idx := f.allMessagesByPtr[message] 191 typesVar := messageTypesVarName(f) 192 193 // Associate the wrapper types by directly passing them to the MessageInfo. 194 g.P(typesVar, "[", idx, "].OneofWrappers = []interface{} {") 195 for _, oneof := range message.Oneofs { 196 if !oneof.Desc.IsSynthetic() { 197 for _, field := range oneof.Fields { 198 g.P("(*", field.GoIdent, ")(nil),") 199 } 200 } 201 } 202 g.P("}") 203 } 204 } 205 } 206 207 g.P("type x struct{}") 208 g.P("out := ", protoimplPackage.Ident("TypeBuilder"), "{") 209 g.P("File: ", protoimplPackage.Ident("DescBuilder"), "{") 210 g.P("GoPackagePath: ", reflectPackage.Ident("TypeOf"), "(x{}).PkgPath(),") 211 g.P("RawDescriptor: ", rawDescVarName(f), ",") 212 g.P("NumEnums: ", len(f.allEnums), ",") 213 g.P("NumMessages: ", len(f.allMessages), ",") 214 g.P("NumExtensions: ", len(f.allExtensions), ",") 215 g.P("NumServices: ", len(f.Services), ",") 216 g.P("},") 217 g.P("GoTypes: ", goTypesVarName(f), ",") 218 g.P("DependencyIndexes: ", depIdxsVarName(f), ",") 219 if len(f.allEnums) > 0 { 220 g.P("EnumInfos: ", enumTypesVarName(f), ",") 221 } 222 if len(f.allMessages) > 0 { 223 g.P("MessageInfos: ", messageTypesVarName(f), ",") 224 } 225 if len(f.allExtensions) > 0 { 226 g.P("ExtensionInfos: ", extensionTypesVarName(f), ",") 227 } 228 g.P("}.Build()") 229 g.P(f.GoDescriptorIdent, " = out.File") 230 231 // Set inputs to nil to allow GC to reclaim resources. 232 g.P(rawDescVarName(f), " = nil") 233 g.P(goTypesVarName(f), " = nil") 234 g.P(depIdxsVarName(f), " = nil") 235 g.P("}") 236} 237 238// stripSourceRetentionFieldsFromMessage walks the given message tree recursively 239// and clears any fields with the field option: [retention = RETENTION_SOURCE] 240func stripSourceRetentionFieldsFromMessage(m protoreflect.Message) { 241 protorange.Range(m, func(ppv protopath.Values) error { 242 m2, ok := ppv.Index(-1).Value.Interface().(protoreflect.Message) 243 if !ok { 244 return nil 245 } 246 m2.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool { 247 fdo, ok := fd.Options().(*descriptorpb.FieldOptions) 248 if ok && fdo.GetRetention() == descriptorpb.FieldOptions_RETENTION_SOURCE { 249 m2.Clear(fd) 250 } 251 return true 252 }) 253 return nil 254 }) 255} 256 257func genFileDescriptor(gen *protogen.Plugin, g *protogen.GeneratedFile, f *fileInfo) { 258 descProto := proto.Clone(f.Proto).(*descriptorpb.FileDescriptorProto) 259 descProto.SourceCodeInfo = nil // drop source code information 260 stripSourceRetentionFieldsFromMessage(descProto.ProtoReflect()) 261 b, err := proto.MarshalOptions{AllowPartial: true, Deterministic: true}.Marshal(descProto) 262 if err != nil { 263 gen.Error(err) 264 return 265 } 266 267 g.P("var ", rawDescVarName(f), " = []byte{") 268 for len(b) > 0 { 269 n := 16 270 if n > len(b) { 271 n = len(b) 272 } 273 274 s := "" 275 for _, c := range b[:n] { 276 s += fmt.Sprintf("0x%02x,", c) 277 } 278 g.P(s) 279 280 b = b[n:] 281 } 282 g.P("}") 283 g.P() 284 285 if f.needRawDesc { 286 onceVar := rawDescVarName(f) + "Once" 287 dataVar := rawDescVarName(f) + "Data" 288 g.P("var (") 289 g.P(onceVar, " ", syncPackage.Ident("Once")) 290 g.P(dataVar, " = ", rawDescVarName(f)) 291 g.P(")") 292 g.P() 293 294 g.P("func ", rawDescVarName(f), "GZIP() []byte {") 295 g.P(onceVar, ".Do(func() {") 296 g.P(dataVar, " = ", protoimplPackage.Ident("X"), ".CompressGZIP(", dataVar, ")") 297 g.P("})") 298 g.P("return ", dataVar) 299 g.P("}") 300 g.P() 301 } 302} 303 304func genEnumReflectMethods(g *protogen.GeneratedFile, f *fileInfo, e *enumInfo) { 305 idx := f.allEnumsByPtr[e] 306 typesVar := enumTypesVarName(f) 307 308 // Descriptor method. 309 g.P("func (", e.GoIdent, ") Descriptor() ", protoreflectPackage.Ident("EnumDescriptor"), " {") 310 g.P("return ", typesVar, "[", idx, "].Descriptor()") 311 g.P("}") 312 g.P() 313 314 // Type method. 315 g.P("func (", e.GoIdent, ") Type() ", protoreflectPackage.Ident("EnumType"), " {") 316 g.P("return &", typesVar, "[", idx, "]") 317 g.P("}") 318 g.P() 319 320 // Number method. 321 g.P("func (x ", e.GoIdent, ") Number() ", protoreflectPackage.Ident("EnumNumber"), " {") 322 g.P("return ", protoreflectPackage.Ident("EnumNumber"), "(x)") 323 g.P("}") 324 g.P() 325} 326 327func genMessageReflectMethods(g *protogen.GeneratedFile, f *fileInfo, m *messageInfo) { 328 idx := f.allMessagesByPtr[m] 329 typesVar := messageTypesVarName(f) 330 331 // ProtoReflect method. 332 g.P("func (x *", m.GoIdent, ") ProtoReflect() ", protoreflectPackage.Ident("Message"), " {") 333 g.P("mi := &", typesVar, "[", idx, "]") 334 g.P("if ", protoimplPackage.Ident("UnsafeEnabled"), " && x != nil {") 335 g.P("ms := ", protoimplPackage.Ident("X"), ".MessageStateOf(", protoimplPackage.Ident("Pointer"), "(x))") 336 g.P("if ms.LoadMessageInfo() == nil {") 337 g.P("ms.StoreMessageInfo(mi)") 338 g.P("}") 339 g.P("return ms") 340 g.P("}") 341 g.P("return mi.MessageOf(x)") 342 g.P("}") 343 g.P() 344} 345 346func fileVarName(f *protogen.File, suffix string) string { 347 prefix := f.GoDescriptorIdent.GoName 348 _, n := utf8.DecodeRuneInString(prefix) 349 prefix = strings.ToLower(prefix[:n]) + prefix[n:] 350 return prefix + "_" + suffix 351} 352func rawDescVarName(f *fileInfo) string { 353 return fileVarName(f.File, "rawDesc") 354} 355func goTypesVarName(f *fileInfo) string { 356 return fileVarName(f.File, "goTypes") 357} 358func depIdxsVarName(f *fileInfo) string { 359 return fileVarName(f.File, "depIdxs") 360} 361func enumTypesVarName(f *fileInfo) string { 362 return fileVarName(f.File, "enumTypes") 363} 364func messageTypesVarName(f *fileInfo) string { 365 return fileVarName(f.File, "msgTypes") 366} 367func extensionTypesVarName(f *fileInfo) string { 368 return fileVarName(f.File, "extTypes") 369} 370func initFuncName(f *protogen.File) string { 371 return fileVarName(f, "init") 372} 373