Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 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 | |
| 5 | package internal_gengo |
| 6 | |
| 7 | import ( |
| 8 | "fmt" |
| 9 | "math" |
Joe Tsai | 7ca7098 | 2019-04-15 13:57:56 -0700 | [diff] [blame] | 10 | "strings" |
| 11 | "unicode/utf8" |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 12 | |
Damien Neil | 5c5b531 | 2019-05-14 12:44:37 -0700 | [diff] [blame] | 13 | "google.golang.org/protobuf/compiler/protogen" |
Damien Neil | e89e624 | 2019-05-13 23:55:40 -0700 | [diff] [blame] | 14 | "google.golang.org/protobuf/proto" |
Damien Neil | e89e624 | 2019-05-13 23:55:40 -0700 | [diff] [blame] | 15 | "google.golang.org/protobuf/reflect/protoreflect" |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 16 | |
Joe Tsai | a95b29f | 2019-05-16 12:47:20 -0700 | [diff] [blame] | 17 | "google.golang.org/protobuf/types/descriptorpb" |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 18 | ) |
| 19 | |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 20 | func genReflectFileDescriptor(gen *protogen.Plugin, g *protogen.GeneratedFile, f *fileInfo) { |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 21 | g.P("var ", f.GoDescriptorIdent, " ", protoreflectPackage.Ident("FileDescriptor")) |
| 22 | g.P() |
| 23 | |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 24 | genFileDescriptor(gen, g, f) |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 25 | if len(f.allEnums) > 0 { |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 26 | g.P("var ", enumTypesVarName(f), " = make([]", prototypePackage.Ident("Enum"), ",", len(f.allEnums), ")") |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 27 | } |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 28 | if len(f.allMessages) > 0 { |
Joe Tsai | 4fe9663 | 2019-05-22 05:12:36 -0400 | [diff] [blame] | 29 | g.P("var ", messageTypesVarName(f), " = make([]", protoimplPackage.Ident("MessageInfo"), ",", len(f.allMessages), ")") |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 30 | } |
| 31 | |
| 32 | // Generate a unique list of Go types for all declarations and dependencies, |
| 33 | // and the associated index into the type list for all dependencies. |
| 34 | var goTypes []string |
| 35 | var depIdxs []string |
| 36 | seen := map[protoreflect.FullName]int{} |
| 37 | genDep := func(name protoreflect.FullName, depSource string) { |
| 38 | if depSource != "" { |
| 39 | line := fmt.Sprintf("%d, // %s -> %s", seen[name], depSource, name) |
| 40 | depIdxs = append(depIdxs, line) |
| 41 | } |
| 42 | } |
| 43 | genEnum := func(e *protogen.Enum, depSource string) { |
| 44 | if e != nil { |
| 45 | name := e.Desc.FullName() |
| 46 | if _, ok := seen[name]; !ok { |
| 47 | line := fmt.Sprintf("(%s)(0), // %d: %s", g.QualifiedGoIdent(e.GoIdent), len(goTypes), name) |
| 48 | goTypes = append(goTypes, line) |
| 49 | seen[name] = len(seen) |
| 50 | } |
| 51 | if depSource != "" { |
| 52 | genDep(name, depSource) |
| 53 | } |
| 54 | } |
| 55 | } |
| 56 | genMessage := func(m *protogen.Message, depSource string) { |
| 57 | if m != nil { |
| 58 | name := m.Desc.FullName() |
| 59 | if _, ok := seen[name]; !ok { |
| 60 | line := fmt.Sprintf("(*%s)(nil), // %d: %s", g.QualifiedGoIdent(m.GoIdent), len(goTypes), name) |
| 61 | if m.Desc.IsMapEntry() { |
| 62 | // Map entry messages have no associated Go type. |
| 63 | line = fmt.Sprintf("nil, // %d: %s", len(goTypes), name) |
| 64 | } |
| 65 | goTypes = append(goTypes, line) |
| 66 | seen[name] = len(seen) |
| 67 | } |
| 68 | if depSource != "" { |
| 69 | genDep(name, depSource) |
| 70 | } |
| 71 | } |
| 72 | } |
| 73 | |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 74 | // This ordering is significant. |
| 75 | // See filetype.TypeBuilder.DependencyIndexes. |
| 76 | var depOffsets []string |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 77 | for _, enum := range f.allEnums { |
| 78 | genEnum(enum, "") |
| 79 | } |
| 80 | for _, message := range f.allMessages { |
| 81 | genMessage(message, "") |
| 82 | } |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 83 | depOffsets = append(depOffsets, fmt.Sprintf("%d, // starting offset of field type_name sub-list", len(depIdxs))) |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 84 | for _, message := range f.allMessages { |
| 85 | for _, field := range message.Fields { |
| 86 | if field.Desc.IsWeak() { |
| 87 | continue |
| 88 | } |
| 89 | source := string(field.Desc.FullName()) |
Joe Tsai | d24bc72 | 2019-04-15 23:39:09 -0700 | [diff] [blame] | 90 | genEnum(field.Enum, source+":type_name") |
| 91 | genMessage(field.Message, source+":type_name") |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 92 | } |
| 93 | } |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 94 | depOffsets = append(depOffsets, fmt.Sprintf("%d, // starting offset of extension extendee sub-list", len(depIdxs))) |
| 95 | for _, extension := range f.allExtensions { |
| 96 | source := string(extension.Desc.FullName()) |
| 97 | genMessage(extension.Extendee, source+":extendee") |
| 98 | } |
| 99 | depOffsets = append(depOffsets, fmt.Sprintf("%d, // starting offset of extension type_name sub-list", len(depIdxs))) |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 100 | for _, extension := range f.allExtensions { |
| 101 | source := string(extension.Desc.FullName()) |
Joe Tsai | d24bc72 | 2019-04-15 23:39:09 -0700 | [diff] [blame] | 102 | genEnum(extension.Enum, source+":type_name") |
| 103 | genMessage(extension.Message, source+":type_name") |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 104 | } |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 105 | depOffsets = append(depOffsets, fmt.Sprintf("%d, // starting offset of method input_type sub-list", len(depIdxs))) |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 106 | for _, service := range f.Services { |
| 107 | for _, method := range service.Methods { |
| 108 | source := string(method.Desc.FullName()) |
Joe Tsai | d24bc72 | 2019-04-15 23:39:09 -0700 | [diff] [blame] | 109 | genMessage(method.Input, source+":input_type") |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 110 | } |
| 111 | } |
| 112 | depOffsets = append(depOffsets, fmt.Sprintf("%d, // starting offset of method output_type sub-list", len(depIdxs))) |
| 113 | for _, service := range f.Services { |
| 114 | for _, method := range service.Methods { |
| 115 | source := string(method.Desc.FullName()) |
Joe Tsai | d24bc72 | 2019-04-15 23:39:09 -0700 | [diff] [blame] | 116 | genMessage(method.Output, source+":output_type") |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 117 | } |
| 118 | } |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 119 | for i := len(depOffsets) - 1; i >= 0; i-- { |
| 120 | depIdxs = append(depIdxs, depOffsets[i]) |
| 121 | } |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 122 | if len(depIdxs) > math.MaxInt32 { |
| 123 | panic("too many dependencies") // sanity check |
| 124 | } |
| 125 | |
| 126 | g.P("var ", goTypesVarName(f), " = []interface{}{") |
| 127 | for _, s := range goTypes { |
| 128 | g.P(s) |
| 129 | } |
| 130 | g.P("}") |
| 131 | |
| 132 | g.P("var ", depIdxsVarName(f), " = []int32{") |
| 133 | for _, s := range depIdxs { |
| 134 | g.P(s) |
| 135 | } |
| 136 | g.P("}") |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 137 | |
Damien Neil | 0fc2245 | 2019-03-08 17:18:11 -0800 | [diff] [blame] | 138 | g.P("func init() { ", initFuncName(f.File), "() }") |
| 139 | |
| 140 | g.P("func ", initFuncName(f.File), "() {") |
| 141 | g.P("if ", f.GoDescriptorIdent, " != nil {") |
| 142 | g.P("return") |
| 143 | g.P("}") |
| 144 | |
| 145 | // Ensure that initialization functions for different files in the same Go |
| 146 | // package run in the correct order: Call the init funcs for every .proto file |
| 147 | // imported by this one that is in the same Go package. |
| 148 | for i, imps := 0, f.Desc.Imports(); i < imps.Len(); i++ { |
| 149 | impFile, _ := gen.FileByName(imps.Get(i).Path()) |
| 150 | if impFile.GoImportPath != f.GoImportPath { |
| 151 | continue |
| 152 | } |
| 153 | g.P(initFuncName(impFile), "()") |
| 154 | } |
| 155 | |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 156 | g.P("out := ", protoimplPackage.Ident("TypeBuilder"), "{") |
| 157 | g.P("File: ", protoimplPackage.Ident("DescBuilder"), "{") |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 158 | g.P("RawDescriptor: ", rawDescVarName(f), ",") |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 159 | g.P("NumEnums: ", len(f.allEnums), ",") |
| 160 | g.P("NumMessages: ", len(f.allMessages), ",") |
| 161 | g.P("NumExtensions: ", len(f.allExtensions), ",") |
| 162 | g.P("NumServices: ", len(f.Services), ",") |
| 163 | g.P("},") |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 164 | g.P("GoTypes: ", goTypesVarName(f), ",") |
| 165 | g.P("DependencyIndexes: ", depIdxsVarName(f), ",") |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 166 | if len(f.allMessages) > 0 { |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 167 | g.P("MessageInfos: ", messageTypesVarName(f), ",") |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 168 | } |
| 169 | if len(f.allExtensions) > 0 { |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 170 | g.P("LegacyExtensions: ", extDescsVarName(f), ",") |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 171 | } |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 172 | g.P("}.Build()") |
| 173 | g.P(f.GoDescriptorIdent, " = out.File") |
| 174 | if len(f.allEnums) > 0 { |
| 175 | g.P(enumTypesVarName(f), " = out.Enums") |
| 176 | } |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 177 | |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 178 | // Set inputs to nil to allow GC to reclaim resources. |
| 179 | g.P(rawDescVarName(f), " = nil") |
| 180 | g.P(goTypesVarName(f), " = nil") |
| 181 | g.P(depIdxsVarName(f), " = nil") |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 182 | g.P("}") |
| 183 | } |
| 184 | |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 185 | func genFileDescriptor(gen *protogen.Plugin, g *protogen.GeneratedFile, f *fileInfo) { |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 186 | descProto := new(descriptorpb.FileDescriptorProto) |
Joe Tsai | 8b0b71e | 2019-07-08 11:33:50 -0700 | [diff] [blame^] | 187 | proto.Merge(descProto, f.Proto) |
| 188 | descProto.SourceCodeInfo = nil // drop source code information |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 189 | |
Joe Tsai | 8b0b71e | 2019-07-08 11:33:50 -0700 | [diff] [blame^] | 190 | b, err := proto.MarshalOptions{AllowPartial: true, Deterministic: true}.Marshal(descProto) |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 191 | if err != nil { |
| 192 | gen.Error(err) |
| 193 | return |
| 194 | } |
| 195 | |
| 196 | g.P("var ", rawDescVarName(f), " = []byte{") |
| 197 | for len(b) > 0 { |
| 198 | n := 16 |
| 199 | if n > len(b) { |
| 200 | n = len(b) |
| 201 | } |
| 202 | |
| 203 | s := "" |
| 204 | for _, c := range b[:n] { |
| 205 | s += fmt.Sprintf("0x%02x,", c) |
| 206 | } |
| 207 | g.P(s) |
| 208 | |
| 209 | b = b[n:] |
| 210 | } |
| 211 | g.P("}") |
| 212 | g.P() |
| 213 | |
Joe Tsai | ab61d41 | 2019-04-16 15:23:29 -0700 | [diff] [blame] | 214 | if generateRawDescMethods { |
| 215 | onceVar := rawDescVarName(f) + "Once" |
| 216 | dataVar := rawDescVarName(f) + "Data" |
| 217 | g.P("var (") |
| 218 | g.P(onceVar, " ", syncPackage.Ident("Once")) |
| 219 | g.P(dataVar, " = ", rawDescVarName(f)) |
| 220 | g.P(")") |
| 221 | g.P() |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 222 | |
Joe Tsai | ab61d41 | 2019-04-16 15:23:29 -0700 | [diff] [blame] | 223 | g.P("func ", rawDescVarName(f), "GZIP() []byte {") |
| 224 | g.P(onceVar, ".Do(func() {") |
| 225 | g.P(dataVar, " = ", protoimplPackage.Ident("X"), ".CompressGZIP(", dataVar, ")") |
| 226 | g.P("})") |
| 227 | g.P("return ", dataVar) |
| 228 | g.P("}") |
| 229 | g.P() |
| 230 | } |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 231 | } |
| 232 | |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 233 | func genReflectEnum(gen *protogen.Plugin, g *protogen.GeneratedFile, f *fileInfo, enum *protogen.Enum) { |
Joe Tsai | 9667c48 | 2018-12-05 15:42:52 -0800 | [diff] [blame] | 234 | idx := f.allEnumsByPtr[enum] |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 235 | typesVar := enumTypesVarName(f) |
Joe Tsai | 61968ce | 2019-04-01 12:59:24 -0700 | [diff] [blame] | 236 | |
Joe Tsai | 0fc49f8 | 2019-05-01 12:29:25 -0700 | [diff] [blame] | 237 | // Descriptor method. |
| 238 | g.P("func (", enum.GoIdent, ") Descriptor() ", protoreflectPackage.Ident("EnumDescriptor"), " {") |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 239 | g.P("return ", typesVar, "[", idx, "].EnumDescriptor") |
Joe Tsai | 0fc49f8 | 2019-05-01 12:29:25 -0700 | [diff] [blame] | 240 | g.P("}") |
| 241 | g.P() |
| 242 | |
Joe Tsai | 61968ce | 2019-04-01 12:59:24 -0700 | [diff] [blame] | 243 | // Number method. |
| 244 | g.P("func (x ", enum.GoIdent, ") Number() ", protoreflectPackage.Ident("EnumNumber"), " {") |
| 245 | g.P("return ", protoreflectPackage.Ident("EnumNumber"), "(x)") |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 246 | g.P("}") |
Joe Tsai | 61968ce | 2019-04-01 12:59:24 -0700 | [diff] [blame] | 247 | g.P() |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 248 | } |
| 249 | |
| 250 | func genReflectMessage(gen *protogen.Plugin, g *protogen.GeneratedFile, f *fileInfo, message *protogen.Message) { |
Joe Tsai | 9667c48 | 2018-12-05 15:42:52 -0800 | [diff] [blame] | 251 | idx := f.allMessagesByPtr[message] |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 252 | typesVar := messageTypesVarName(f) |
Joe Tsai | 61968ce | 2019-04-01 12:59:24 -0700 | [diff] [blame] | 253 | |
| 254 | // ProtoReflect method. |
| 255 | g.P("func (x *", message.GoIdent, ") ProtoReflect() ", protoreflectPackage.Ident("Message"), " {") |
| 256 | g.P("return ", typesVar, "[", idx, "].MessageOf(x)") |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 257 | g.P("}") |
Joe Tsai | 61968ce | 2019-04-01 12:59:24 -0700 | [diff] [blame] | 258 | g.P() |
Damien Neil | 0d3e8cc | 2019-04-01 13:31:55 -0700 | [diff] [blame] | 259 | g.P("func (m *", message.GoIdent, ") XXX_Methods() *", protoifacePackage.Ident("Methods"), " {") |
| 260 | g.P("return ", typesVar, "[", idx, "].Methods()") |
| 261 | g.P("}") |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 262 | } |
| 263 | |
Joe Tsai | 7ca7098 | 2019-04-15 13:57:56 -0700 | [diff] [blame] | 264 | func fileVarName(f *protogen.File, suffix string) string { |
| 265 | prefix := f.GoDescriptorIdent.GoName |
| 266 | _, n := utf8.DecodeRuneInString(prefix) |
| 267 | prefix = strings.ToLower(prefix[:n]) + prefix[n:] |
| 268 | return prefix + "_" + suffix |
| 269 | } |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 270 | func rawDescVarName(f *fileInfo) string { |
Joe Tsai | 7ca7098 | 2019-04-15 13:57:56 -0700 | [diff] [blame] | 271 | return fileVarName(f.File, "rawDesc") |
Joe Tsai | 5d72cc2 | 2019-03-28 01:13:26 -0700 | [diff] [blame] | 272 | } |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 273 | func goTypesVarName(f *fileInfo) string { |
Joe Tsai | 7ca7098 | 2019-04-15 13:57:56 -0700 | [diff] [blame] | 274 | return fileVarName(f.File, "goTypes") |
Damien Neil | 8012b44 | 2019-01-18 09:32:24 -0800 | [diff] [blame] | 275 | } |
| 276 | func depIdxsVarName(f *fileInfo) string { |
Joe Tsai | 7ca7098 | 2019-04-15 13:57:56 -0700 | [diff] [blame] | 277 | return fileVarName(f.File, "depIdxs") |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 278 | } |
| 279 | func enumTypesVarName(f *fileInfo) string { |
Joe Tsai | 7ca7098 | 2019-04-15 13:57:56 -0700 | [diff] [blame] | 280 | return fileVarName(f.File, "enumTypes") |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 281 | } |
| 282 | func messageTypesVarName(f *fileInfo) string { |
Joe Tsai | 7ca7098 | 2019-04-15 13:57:56 -0700 | [diff] [blame] | 283 | return fileVarName(f.File, "msgTypes") |
Joe Tsai | b6405bd | 2018-11-15 14:44:37 -0800 | [diff] [blame] | 284 | } |
Joe Tsai | d888139 | 2019-06-06 13:01:53 -0700 | [diff] [blame] | 285 | func extDescsVarName(f *fileInfo) string { |
Joe Tsai | 7ca7098 | 2019-04-15 13:57:56 -0700 | [diff] [blame] | 286 | return fileVarName(f.File, "extDescs") |
Joe Tsai | afb455e | 2019-03-14 16:08:22 -0700 | [diff] [blame] | 287 | } |
Damien Neil | 0fc2245 | 2019-03-08 17:18:11 -0800 | [diff] [blame] | 288 | func initFuncName(f *protogen.File) string { |
Joe Tsai | 7ca7098 | 2019-04-15 13:57:56 -0700 | [diff] [blame] | 289 | return fileVarName(f, "init") |
Damien Neil | 0fc2245 | 2019-03-08 17:18:11 -0800 | [diff] [blame] | 290 | } |