func (p *test) Generate(imports generator.PluginImports, file *generator.FileDescriptor) bool { used := false randPkg := imports.NewImport("math/rand") timePkg := imports.NewImport("time") testingPkg := imports.NewImport("testing") protoPkg := imports.NewImport("github.com/gogo/protobuf/proto") if !gogoproto.ImportsGoGoProto(file.FileDescriptorProto) { protoPkg = imports.NewImport("github.com/golang/protobuf/proto") } for _, message := range file.Messages() { ccTypeName := generator.CamelCaseSlice(message.TypeName()) if !gogoproto.IsSizer(file.FileDescriptorProto, message.DescriptorProto) { continue } if message.DescriptorProto.GetOptions().GetMapEntry() { continue } if gogoproto.HasTestGen(file.FileDescriptorProto, message.DescriptorProto) { used = true p.P(`func Test`, ccTypeName, `Size(t *`, testingPkg.Use(), `.T) {`) p.In() p.P(`seed := `, timePkg.Use(), `.Now().UnixNano()`) p.P(`popr := `, randPkg.Use(), `.New(`, randPkg.Use(), `.NewSource(seed))`) p.P(`p := NewPopulated`, ccTypeName, `(popr, true)`) p.P(`size2 := `, protoPkg.Use(), `.Size(p)`) p.P(`data, err := `, protoPkg.Use(), `.Marshal(p)`) p.P(`if err != nil {`) p.In() p.P(`t.Fatalf("seed = %d, err = %v", seed, err)`) p.Out() p.P(`}`) p.P(`size := p.Size()`) p.P(`if len(data) != size {`) p.In() p.P(`t.Errorf("seed = %d, size %v != marshalled size %v", seed, size, len(data))`) p.Out() p.P(`}`) p.P(`if size2 != size {`) p.In() p.P(`t.Errorf("seed = %d, size %v != before marshal proto.Size %v", seed, size, size2)`) p.Out() p.P(`}`) p.P(`size3 := `, protoPkg.Use(), `.Size(p)`) p.P(`if size3 != size {`) p.In() p.P(`t.Errorf("seed = %d, size %v != after marshal proto.Size %v", seed, size, size3)`) p.Out() p.P(`}`) p.Out() p.P(`}`) p.P() } if gogoproto.HasBenchGen(file.FileDescriptorProto, message.DescriptorProto) { used = true p.P(`func Benchmark`, ccTypeName, `Size(b *`, testingPkg.Use(), `.B) {`) p.In() p.P(`popr := `, randPkg.Use(), `.New(`, randPkg.Use(), `.NewSource(616))`) p.P(`total := 0`) p.P(`pops := make([]*`, ccTypeName, `, 1000)`) p.P(`for i := 0; i < 1000; i++ {`) p.In() p.P(`pops[i] = NewPopulated`, ccTypeName, `(popr, false)`) p.Out() p.P(`}`) p.P(`b.ResetTimer()`) p.P(`for i := 0; i < b.N; i++ {`) p.In() p.P(`total += pops[i%1000].Size()`) p.Out() p.P(`}`) p.P(`b.SetBytes(int64(total / b.N))`) p.Out() p.P(`}`) p.P() } } return used }
func (p *size) Generate(file *generator.FileDescriptor) { p.PluginImports = generator.NewPluginImports(p.Generator) p.atleastOne = false p.localName = generator.FileName(file) protoPkg := p.NewImport("github.com/gogo/protobuf/proto") if !gogoproto.ImportsGoGoProto(file.FileDescriptorProto) { protoPkg = p.NewImport("github.com/golang/protobuf/proto") } for _, message := range file.Messages() { if !gogoproto.IsSizer(file.FileDescriptorProto, message.DescriptorProto) { continue } if message.DescriptorProto.GetOptions().GetMapEntry() { continue } p.atleastOne = true ccTypeName := generator.CamelCaseSlice(message.TypeName()) p.P(`func (m *`, ccTypeName, `) Size() (n int) {`) p.In() p.P(`var l int`) p.P(`_ = l`) oneofs := make(map[string]struct{}) for _, field := range message.Field { oneof := field.OneofIndex != nil if !oneof { proto3 := gogoproto.IsProto3(file.FileDescriptorProto) p.generateField(proto3, file, message, field) } else { fieldname := p.GetFieldName(message, field) if _, ok := oneofs[fieldname]; ok { continue } else { oneofs[fieldname] = struct{}{} } p.P(`if m.`, fieldname, ` != nil {`) p.In() p.P(`n+=m.`, fieldname, `.Size()`) p.Out() p.P(`}`) } } if message.DescriptorProto.HasExtension() { p.P(`if m.XXX_extensions != nil {`) p.In() if gogoproto.HasExtensionsMap(file.FileDescriptorProto, message.DescriptorProto) { p.P(`n += `, protoPkg.Use(), `.SizeOfExtensionMap(m.XXX_extensions)`) } else { p.P(`n+=len(m.XXX_extensions)`) } p.Out() p.P(`}`) } if gogoproto.HasUnrecognized(file.FileDescriptorProto, message.DescriptorProto) { p.P(`if m.XXX_unrecognized != nil {`) p.In() p.P(`n+=len(m.XXX_unrecognized)`) p.Out() p.P(`}`) } p.P(`return n`) p.Out() p.P(`}`) p.P() //Generate Size methods for oneof fields m := proto.Clone(message.DescriptorProto).(*descriptor.DescriptorProto) for _, f := range m.Field { oneof := f.OneofIndex != nil if !oneof { continue } ccTypeName := p.OneOfTypeName(message, f) p.P(`func (m *`, ccTypeName, `) Size() (n int) {`) p.In() p.P(`var l int`) p.P(`_ = l`) vanity.TurnOffNullableForNativeTypesWithoutDefaultsOnly(f) p.generateField(false, file, message, f) p.P(`return n`) p.Out() p.P(`}`) } } if !p.atleastOne { return } p.sizeVarint() p.sizeZigZag() }