func (p *plugin) Generate(file *generator.FileDescriptor) { p.atleastOne = false p.PluginImports = generator.NewPluginImports(p.Generator) p.varGen = NewVarGen() proto3 := gogoproto.IsProto3(file.FileDescriptorProto) 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.HasPopulate(file.FileDescriptorProto, message.DescriptorProto) { continue } if message.DescriptorProto.GetOptions().GetMapEntry() { continue } p.atleastOne = true ccTypeName := generator.CamelCaseSlice(message.TypeName()) p.P(`func NewPopulated`, ccTypeName, `(r randy`, p.localName, `, easy bool) *`, ccTypeName, ` {`) p.In() p.P(`this := &`, ccTypeName, `{}`) if gogoproto.IsUnion(message.File(), message.DescriptorProto) && len(message.Field) > 0 { loopLevels := make([]int, len(message.Field)) maxLoopLevel := 0 for i, field := range message.Field { loopLevels[i] = p.loops(field, message) if loopLevels[i] > maxLoopLevel { maxLoopLevel = loopLevels[i] } } ran := 0 for i := range loopLevels { ran += int(math.Pow10(maxLoopLevel - loopLevels[i])) } p.P(`fieldNum := r.Intn(`, fmt.Sprintf("%d", ran), `)`) p.P(`switch fieldNum {`) k := 0 for i, field := range message.Field { is := []string{} ran := int(math.Pow10(maxLoopLevel - loopLevels[i])) for j := 0; j < ran; j++ { is = append(is, fmt.Sprintf("%d", j+k)) } k += ran p.P(`case `, strings.Join(is, ","), `:`) p.In() p.GenerateField(file, message, field) p.Out() } p.P(`}`) } else { var maxFieldNumber int32 oneofs := make(map[string]struct{}) for _, field := range message.Field { if field.GetNumber() > maxFieldNumber { maxFieldNumber = field.GetNumber() } oneof := field.OneofIndex != nil if !oneof { if field.IsRequired() || (!gogoproto.IsNullable(field) && !field.IsRepeated()) || (proto3 && !field.IsMessage()) { p.GenerateField(file, message, field) } else { p.P(`if r.Intn(10) != 0 {`) p.In() p.GenerateField(file, message, field) p.Out() p.P(`}`) } } else { fieldname := p.GetFieldName(message, field) if _, ok := oneofs[fieldname]; ok { continue } else { oneofs[fieldname] = struct{}{} } fieldNumbers := []int32{} for _, f := range message.Field { fname := p.GetFieldName(message, f) if fname == fieldname { fieldNumbers = append(fieldNumbers, f.GetNumber()) } } p.P(`oneofNumber_`, fieldname, ` := `, fmt.Sprintf("%#v", fieldNumbers), `[r.Intn(`, strconv.Itoa(len(fieldNumbers)), `)]`) p.P(`switch oneofNumber_`, fieldname, ` {`) for _, f := range message.Field { fname := p.GetFieldName(message, f) if fname != fieldname { continue } p.P(`case `, strconv.Itoa(int(f.GetNumber())), `:`) p.In() ccTypeName := p.OneOfTypeName(message, f) p.P(`this.`, fname, ` = NewPopulated`, ccTypeName, `(r, easy)`) p.Out() } p.P(`}`) } } if message.DescriptorProto.HasExtension() { p.P(`if !easy && r.Intn(10) != 0 {`) p.In() p.P(`l := r.Intn(5)`) p.P(`for i := 0; i < l; i++ {`) p.In() if len(message.DescriptorProto.GetExtensionRange()) > 1 { p.P(`eIndex := r.Intn(`, strconv.Itoa(len(message.DescriptorProto.GetExtensionRange())), `)`) p.P(`fieldNumber := 0`) p.P(`switch eIndex {`) for i, e := range message.DescriptorProto.GetExtensionRange() { p.P(`case `, strconv.Itoa(i), `:`) p.In() p.P(`fieldNumber = r.Intn(`, strconv.Itoa(int(e.GetEnd()-e.GetStart())), `) + `, strconv.Itoa(int(e.GetStart()))) p.Out() if e.GetEnd() > maxFieldNumber { maxFieldNumber = e.GetEnd() } } p.P(`}`) } else { e := message.DescriptorProto.GetExtensionRange()[0] p.P(`fieldNumber := r.Intn(`, strconv.Itoa(int(e.GetEnd()-e.GetStart())), `) + `, strconv.Itoa(int(e.GetStart()))) if e.GetEnd() > maxFieldNumber { maxFieldNumber = e.GetEnd() } } p.P(`wire := r.Intn(4)`) p.P(`if wire == 3 { wire = 5 }`) p.P(`data := randField`, p.localName, `(nil, r, fieldNumber, wire)`) p.P(protoPkg.Use(), `.SetRawExtension(this, int32(fieldNumber), data)`) p.Out() p.P(`}`) p.Out() p.P(`}`) } if maxFieldNumber < (1 << 10) { p.P(`if !easy && r.Intn(10) != 0 {`) p.In() if gogoproto.HasUnrecognized(file.FileDescriptorProto, message.DescriptorProto) { p.P(`this.XXX_unrecognized = randUnrecognized`, p.localName, `(r, `, strconv.Itoa(int(maxFieldNumber+1)), `)`) } p.Out() p.P(`}`) } } p.P(`return this`) p.Out() p.P(`}`) p.P(``) //Generate NewPopulated functions 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 NewPopulated`, ccTypeName, `(r randy`, p.localName, `, easy bool) *`, ccTypeName, ` {`) p.In() p.P(`this := &`, ccTypeName, `{}`) vanity.TurnOffNullableForNativeTypesWithoutDefaultsOnly(f) p.GenerateField(file, message, f) p.P(`return this`) p.Out() p.P(`}`) } } if !p.atleastOne { return } p.P(`type randy`, p.localName, ` interface {`) p.In() p.P(`Float32() float32`) p.P(`Float64() float64`) p.P(`Int63() int64`) p.P(`Int31() int32`) p.P(`Uint32() uint32`) p.P(`Intn(n int) int`) p.Out() p.P(`}`) p.P(`func randUTF8Rune`, p.localName, `(r randy`, p.localName, `) rune {`) p.In() p.P(`ru := r.Intn(62)`) p.P(`if ru < 10 {`) p.In() p.P(`return rune(ru+48)`) p.Out() p.P(`} else if ru < 36 {`) p.In() p.P(`return rune(ru+55)`) p.Out() p.P(`}`) p.P(`return rune(ru+61)`) p.Out() p.P(`}`) p.P(`func randString`, p.localName, `(r randy`, p.localName, `) string {`) p.In() p.P(p.varGen.Next(), ` := r.Intn(100)`) p.P(`tmps := make([]rune, `, p.varGen.Current(), `)`) p.P(`for i := 0; i < `, p.varGen.Current(), `; i++ {`) p.In() p.P(`tmps[i] = randUTF8Rune`, p.localName, `(r)`) p.Out() p.P(`}`) p.P(`return string(tmps)`) p.Out() p.P(`}`) p.P(`func randUnrecognized`, p.localName, `(r randy`, p.localName, `, maxFieldNumber int) (data []byte) {`) p.In() p.P(`l := r.Intn(5)`) p.P(`for i := 0; i < l; i++ {`) p.In() p.P(`wire := r.Intn(4)`) p.P(`if wire == 3 { wire = 5 }`) p.P(`fieldNumber := maxFieldNumber + r.Intn(100)`) p.P(`data = randField`, p.localName, `(data, r, fieldNumber, wire)`) p.Out() p.P(`}`) p.P(`return data`) p.Out() p.P(`}`) p.P(`func randField`, p.localName, `(data []byte, r randy`, p.localName, `, fieldNumber int, wire int) []byte {`) p.In() p.P(`key := uint32(fieldNumber)<<3 | uint32(wire)`) p.P(`switch wire {`) p.P(`case 0:`) p.In() p.P(`data = encodeVarintPopulate`, p.localName, `(data, uint64(key))`) p.P(p.varGen.Next(), ` := r.Int63()`) p.P(`if r.Intn(2) == 0 {`) p.In() p.P(p.varGen.Current(), ` *= -1`) p.Out() p.P(`}`) p.P(`data = encodeVarintPopulate`, p.localName, `(data, uint64(`, p.varGen.Current(), `))`) p.Out() p.P(`case 1:`) p.In() p.P(`data = encodeVarintPopulate`, p.localName, `(data, uint64(key))`) p.P(`data = append(data, byte(r.Intn(256)), byte(r.Intn(256)), byte(r.Intn(256)), byte(r.Intn(256)), byte(r.Intn(256)), byte(r.Intn(256)), byte(r.Intn(256)), byte(r.Intn(256)))`) p.Out() p.P(`case 2:`) p.In() p.P(`data = encodeVarintPopulate`, p.localName, `(data, uint64(key))`) p.P(`ll := r.Intn(100)`) p.P(`data = encodeVarintPopulate`, p.localName, `(data, uint64(ll))`) p.P(`for j := 0; j < ll; j++ {`) p.In() p.P(`data = append(data, byte(r.Intn(256)))`) p.Out() p.P(`}`) p.Out() p.P(`default:`) p.In() p.P(`data = encodeVarintPopulate`, p.localName, `(data, uint64(key))`) p.P(`data = append(data, byte(r.Intn(256)), byte(r.Intn(256)), byte(r.Intn(256)), byte(r.Intn(256)))`) p.Out() p.P(`}`) p.P(`return data`) p.Out() p.P(`}`) p.P(`func encodeVarintPopulate`, p.localName, `(data []byte, v uint64) []byte {`) p.In() p.P(`for v >= 1<<7 {`) p.In() p.P(`data = append(data, uint8(uint64(v)&0x7f|0x80))`) p.P(`v >>= 7`) p.Out() p.P(`}`) p.P(`data = append(data, uint8(v))`) p.P(`return data`) p.Out() p.P(`}`) }
func (p *plugin) generateMessage(file *generator.FileDescriptor, message *generator.Descriptor, verbose bool) { ccTypeName := generator.CamelCaseSlice(message.TypeName()) if verbose { p.P(`func (this *`, ccTypeName, `) VerboseEqual(that interface{}) error {`) } else { p.P(`func (this *`, ccTypeName, `) Equal(that interface{}) bool {`) } p.In() p.generateMsgNullAndTypeCheck(ccTypeName, verbose) oneofs := make(map[string]struct{}) for _, field := range message.Field { oneof := field.OneofIndex != nil if oneof { fieldname := p.GetFieldName(message, field) if _, ok := oneofs[fieldname]; ok { continue } else { oneofs[fieldname] = struct{}{} } p.P(`if that1.`, fieldname, ` == nil {`) p.In() p.P(`if this.`, fieldname, ` != nil {`) p.In() if verbose { p.P(`return `, p.fmtPkg.Use(), `.Errorf("this.`, fieldname, ` != nil && that1.`, fieldname, ` == nil")`) } else { p.P(`return false`) } p.Out() p.P(`}`) p.Out() p.P(`} else if this.`, fieldname, ` == nil {`) p.In() if verbose { p.P(`return `, p.fmtPkg.Use(), `.Errorf("this.`, fieldname, ` == nil && that1.`, fieldname, ` != nil")`) } else { p.P(`return false`) } p.Out() if verbose { p.P(`} else if err := this.`, fieldname, `.VerboseEqual(that1.`, fieldname, `); err != nil {`) } else { p.P(`} else if !this.`, fieldname, `.Equal(that1.`, fieldname, `) {`) } p.In() if verbose { p.P(`return err`) } else { p.P(`return false`) } p.Out() p.P(`}`) } else { p.generateField(file, message, field, verbose) } } if message.DescriptorProto.HasExtension() { fieldname := "XXX_extensions" if gogoproto.HasExtensionsMap(file.FileDescriptorProto, message.DescriptorProto) { p.P(`for k, v := range this.`, fieldname, ` {`) p.In() p.P(`if v2, ok := that1.`, fieldname, `[k]; ok {`) p.In() p.P(`if !v.Equal(&v2) {`) p.In() if verbose { p.P(`return `, p.fmtPkg.Use(), `.Errorf("`, fieldname, ` this[%v](%v) Not Equal that[%v](%v)", k, this.`, fieldname, `[k], k, that1.`, fieldname, `[k])`) } else { p.P(`return false`) } p.Out() p.P(`}`) p.Out() p.P(`} else {`) p.In() if verbose { p.P(`return `, p.fmtPkg.Use(), `.Errorf("`, fieldname, `[%v] Not In that", k)`) } else { p.P(`return false`) } p.Out() p.P(`}`) p.Out() p.P(`}`) p.P(`for k, _ := range that1.`, fieldname, ` {`) p.In() p.P(`if _, ok := this.`, fieldname, `[k]; !ok {`) p.In() if verbose { p.P(`return `, p.fmtPkg.Use(), `.Errorf("`, fieldname, `[%v] Not In this", k)`) } else { p.P(`return false`) } p.Out() p.P(`}`) p.Out() p.P(`}`) } else { p.P(`if !`, p.bytesPkg.Use(), `.Equal(this.`, fieldname, `, that1.`, fieldname, `) {`) p.In() if verbose { p.P(`return `, p.fmtPkg.Use(), `.Errorf("`, fieldname, ` this(%v) Not Equal that(%v)", this.`, fieldname, `, that1.`, fieldname, `)`) } else { p.P(`return false`) } p.Out() p.P(`}`) } } if gogoproto.HasUnrecognized(file.FileDescriptorProto, message.DescriptorProto) { fieldname := "XXX_unrecognized" p.P(`if !`, p.bytesPkg.Use(), `.Equal(this.`, fieldname, `, that1.`, fieldname, `) {`) p.In() if verbose { p.P(`return `, p.fmtPkg.Use(), `.Errorf("`, fieldname, ` this(%v) Not Equal that(%v)", this.`, fieldname, `, that1.`, fieldname, `)`) } else { p.P(`return false`) } p.Out() p.P(`}`) } if verbose { p.P(`return nil`) } else { p.P(`return true`) } p.Out() p.P(`}`) //Generate Equal methods for oneof fields m := proto.Clone(message.DescriptorProto).(*descriptor.DescriptorProto) for _, field := range m.Field { oneof := field.OneofIndex != nil if !oneof { continue } ccTypeName := p.OneOfTypeName(message, field) if verbose { p.P(`func (this *`, ccTypeName, `) VerboseEqual(that interface{}) error {`) } else { p.P(`func (this *`, ccTypeName, `) Equal(that interface{}) bool {`) } p.In() p.generateMsgNullAndTypeCheck(ccTypeName, verbose) vanity.TurnOffNullableForNativeTypesWithoutDefaultsOnly(field) p.generateField(file, message, field, verbose) if verbose { p.P(`return nil`) } else { p.P(`return true`) } p.Out() p.P(`}`) } }
func (p *marshalto) Generate(file *generator.FileDescriptor) { numGen := NewNumGen() p.PluginImports = generator.NewPluginImports(p.Generator) p.atleastOne = false p.localName = generator.FileName(file) p.mathPkg = p.NewImport("math") p.sortKeysPkg = p.NewImport("github.com/gogo/protobuf/sortkeys") p.protoPkg = p.NewImport("github.com/gogo/protobuf/proto") if !gogoproto.ImportsGoGoProto(file.FileDescriptorProto) { p.protoPkg = p.NewImport("github.com/golang/protobuf/proto") } p.unsafePkg = p.NewImport("unsafe") p.errorsPkg = p.NewImport("errors") for _, message := range file.Messages() { if message.DescriptorProto.GetOptions().GetMapEntry() { continue } ccTypeName := generator.CamelCaseSlice(message.TypeName()) if p.unsafe { if !gogoproto.IsUnsafeMarshaler(file.FileDescriptorProto, message.DescriptorProto) { continue } if gogoproto.IsMarshaler(file.FileDescriptorProto, message.DescriptorProto) { panic(fmt.Sprintf("unsafe_marshaler and marshalto enabled for %v", ccTypeName)) } } if !p.unsafe { if !gogoproto.IsMarshaler(file.FileDescriptorProto, message.DescriptorProto) { continue } if gogoproto.IsUnsafeMarshaler(file.FileDescriptorProto, message.DescriptorProto) { panic(fmt.Sprintf("unsafe_marshaler and marshalto enabled for %v", ccTypeName)) } } p.atleastOne = true p.P(`func (m *`, ccTypeName, `) Marshal() (data []byte, err error) {`) p.In() p.P(`size := m.Size()`) p.P(`data = make([]byte, size)`) p.P(`n, err := m.MarshalTo(data)`) p.P(`if err != nil {`) p.In() p.P(`return nil, err`) p.Out() p.P(`}`) p.P(`return data[:n], nil`) p.Out() p.P(`}`) p.P(``) p.P(`func (m *`, ccTypeName, `) MarshalTo(data []byte) (int, error) {`) p.In() p.P(`var i int`) p.P(`_ = i`) p.P(`var l int`) p.P(`_ = l`) fields := orderFields(message.GetField()) sort.Sort(fields) oneofs := make(map[string]struct{}) for _, field := range message.Field { oneof := field.OneofIndex != nil if !oneof { proto3 := gogoproto.IsProto3(file.FileDescriptorProto) p.generateField(proto3, numGen, file, message, field) } else { fieldname := p.GetFieldName(message, field) if _, ok := oneofs[fieldname]; !ok { oneofs[fieldname] = struct{}{} p.P(`if m.`, fieldname, ` != nil {`) p.In() p.P(`nn`, numGen.Next(), `, err := m.`, fieldname, `.MarshalTo(data[i:])`) p.P(`if err != nil {`) p.In() p.P(`return 0, err`) p.Out() p.P(`}`) p.P(`i+=nn`, numGen.Current()) p.Out() p.P(`}`) } } } if message.DescriptorProto.HasExtension() { if gogoproto.HasExtensionsMap(file.FileDescriptorProto, message.DescriptorProto) { p.P(`if len(m.XXX_extensions) > 0 {`) p.In() p.P(`n, err := `, p.protoPkg.Use(), `.EncodeExtensionMap(m.XXX_extensions, data[i:])`) p.P(`if err != nil {`) p.In() p.P(`return 0, err`) p.Out() p.P(`}`) p.P(`i+=n`) p.Out() p.P(`}`) } else { p.P(`if m.XXX_extensions != nil {`) p.In() p.P(`i+=copy(data[i:], 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(`i+=copy(data[i:], m.XXX_unrecognized)`) p.Out() p.P(`}`) } p.P(`return i, nil`) p.Out() p.P(`}`) p.P() //Generate MarshalTo methods for oneof fields m := proto.Clone(message.DescriptorProto).(*descriptor.DescriptorProto) for _, field := range m.Field { oneof := field.OneofIndex != nil if !oneof { continue } ccTypeName := p.OneOfTypeName(message, field) p.P(`func (m *`, ccTypeName, `) MarshalTo(data []byte) (int, error) {`) p.In() p.P(`i := 0`) vanity.TurnOffNullableForNativeTypesWithoutDefaultsOnly(field) p.generateField(false, numGen, file, message, field) p.P(`return i, nil`) p.Out() p.P(`}`) } } if p.atleastOne { p.P(`func encodeFixed64`, p.localName, `(data []byte, offset int, v uint64) int {`) p.In() p.P(`data[offset] = uint8(v)`) p.P(`data[offset+1] = uint8(v >> 8)`) p.P(`data[offset+2] = uint8(v >> 16)`) p.P(`data[offset+3] = uint8(v >> 24)`) p.P(`data[offset+4] = uint8(v >> 32)`) p.P(`data[offset+5] = uint8(v >> 40)`) p.P(`data[offset+6] = uint8(v >> 48)`) p.P(`data[offset+7] = uint8(v >> 56)`) p.P(`return offset+8`) p.Out() p.P(`}`) p.P(`func encodeFixed32`, p.localName, `(data []byte, offset int, v uint32) int {`) p.In() p.P(`data[offset] = uint8(v)`) p.P(`data[offset+1] = uint8(v >> 8)`) p.P(`data[offset+2] = uint8(v >> 16)`) p.P(`data[offset+3] = uint8(v >> 24)`) p.P(`return offset+4`) p.Out() p.P(`}`) p.P(`func encodeVarint`, p.localName, `(data []byte, offset int, v uint64) int {`) p.In() p.P(`for v >= 1<<7 {`) p.In() p.P(`data[offset] = uint8(v&0x7f|0x80)`) p.P(`v >>= 7`) p.P(`offset++`) p.Out() p.P(`}`) p.P(`data[offset] = uint8(v)`) p.P(`return offset+1`) p.Out() p.P(`}`) } }
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() }