Rewrite using protogen

README.md

@@ -10,7 +10,7 @@ protoc-gen-marshal-zap is a protoc plugin to generate code that implements zapco
 - protoc compiler
 - protoc-gen-go
 
-See https://github.com/golang/protobuf#installation
+See https://github.com/protocolbuffers/protobuf-go
 
 ## Installation
 
@@ -47,23 +47,18 @@ simple.pb.marshal-zap.go  # auto-generated by protoc-gen-marshal-zap
 simple.proto              # original proto file
 ```
 
-`simple.pb.marshal-zap.go` is generated as follows
+`simple_marshal_zap.pb.go` is generated as follows
 
 ```go
 // Code generated by protoc-gen-marshal-zap. DO NOT EDIT.
-// source: simple.pb.go
+// source: simple.proto
 
 package simple
 
 import (
         "go.uber.org/zap/zapcore"
-        "strconv"
 )
 
-// Reference imports to suppress errors if they are not otherwise used.
-var _ = zapcore.NewNopCore
-var _ = strconv.FormatInt
-
 func (x *SimpleMessage) MarshalLogObject(enc zapcore.ObjectEncoder) error {
         if x == nil {
                 return nil

go.mod

@@ -1,21 +1,17 @@
 module github.com/kei2100/protoc-gen-marshal-zap
 
-go 1.17
+go 1.18
 
 require (
-	github.com/lyft/protoc-gen-star v0.6.0
-	github.com/stretchr/testify v1.7.2
+	github.com/stretchr/testify v1.8.0
 	go.uber.org/zap v1.21.0
 	google.golang.org/protobuf v1.28.0
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/spf13/afero v1.8.2 // indirect
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.8.0 // indirect
-	golang.org/x/text v0.3.7 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

plugin/protoc-gen-marshal-zap/main.go

@@ -1,13 +1,212 @@
 package main
 
 import (
-	pgs "github.com/lyft/protoc-gen-star"
-	pgsgo "github.com/lyft/protoc-gen-star/lang/go"
+	"fmt"
+
+	pbzap "github.com/kei2100/protoc-gen-marshal-zap"
+	"google.golang.org/protobuf/compiler/protogen"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protoreflect"
+	"google.golang.org/protobuf/types/descriptorpb"
+)
+
+const (
+	zapcorePkg = protogen.GoImportPath("go.uber.org/zap/zapcore")
+	fmtPkg     = protogen.GoImportPath("fmt")
 )
 
+func generateListField(g *protogen.GeneratedFile, f *protogen.Field) {
+	fname := f.Desc.Name()
+	g.P(fname, "ArrMarshaller := func(enc ", g.QualifiedGoIdent(zapcorePkg.Ident("ArrayEncoder")), ") error {")
+	g.P("for _, v := range x.", f.GoName, " {")
+	switch f.Desc.Kind() {
+	case protoreflect.BoolKind:
+		g.P("enc.AppendBool(v)")
+	case protoreflect.BytesKind:
+		g.P("enc.AppendByteString(v)")
+	case protoreflect.DoubleKind:
+		g.P("enc.AppendFloat64(v)")
+	case protoreflect.EnumKind:
+		g.P("enc.AppendString(v.String())")
+	case protoreflect.Fixed32Kind, protoreflect.Uint32Kind:
+		g.P("enc.AppendUint32(v)")
+	case protoreflect.Fixed64Kind, protoreflect.Uint64Kind:
+		g.P("enc.AppendUint64(v)")
+	case protoreflect.FloatKind:
+		g.P("enc.AppendFloat32(v)")
+	case protoreflect.Int32Kind, protoreflect.Sfixed32Kind, protoreflect.Sint32Kind:
+		g.P("enc.AppendInt32(v)")
+	case protoreflect.Int64Kind, protoreflect.Sfixed64Kind, protoreflect.Sint64Kind:
+		g.P("enc.AppendInt64(v)")
+	case protoreflect.GroupKind:
+		g.P("enc.AppendReflected(v)")
+	case protoreflect.MessageKind:
+		g.P("if obj, ok := interface{}(v).(", g.QualifiedGoIdent(zapcorePkg.Ident("ObjectMarshaler")), "); ok {")
+		g.P("enc.AppendObject(obj)")
+		g.P("} else {")
+		g.P("enc.AppendReflected(v)")
+		g.P("}")
+	case protoreflect.StringKind:
+		g.P("enc.AppendString(v)")
+	default:
+		g.P("enc.AppendReflected(v)")
+	}
+	g.P("}")
+	g.P("return nil")
+	g.P("}")
+	g.P("enc.AddArray(\"", fname, "\",", g.QualifiedGoIdent(zapcorePkg.Ident("ArrayMarshalerFunc")), "(", fname, "ArrMarshaller))")
+	g.P()
+}
+
+func generateMapField(g *protogen.GeneratedFile, f *protogen.Field) {
+	fname := f.Desc.Name()
+	g.P("enc.AddObject(\"", fname, "\", ", g.QualifiedGoIdent(zapcorePkg.Ident("ObjectMarshalerFunc")), "(func(enc ", g.QualifiedGoIdent(zapcorePkg.Ident("ObjectEncoder")), ") error {")
+	g.P("for k, v := range x.", f.GoName, " {")
+	switch f.Desc.MapValue().Kind() {
+	case protoreflect.BoolKind:
+		g.P("enc.AddBool(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	case protoreflect.BytesKind:
+		g.P("enc.AddBinary(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	case protoreflect.DoubleKind:
+		g.P("enc.AddFloat64(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	case protoreflect.EnumKind:
+		g.P("enc.AddString(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v.String())")
+	case protoreflect.Fixed32Kind, protoreflect.Uint32Kind:
+		g.P("enc.AddUint32(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	case protoreflect.Fixed64Kind, protoreflect.Uint64Kind:
+		g.P("enc.AddUint64(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	case protoreflect.FloatKind:
+		g.P("enc.AddFloat32(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	case protoreflect.Int32Kind, protoreflect.Sfixed32Kind, protoreflect.Sint32Kind:
+		g.P("enc.AddInt32(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	case protoreflect.Int64Kind, protoreflect.Sfixed64Kind, protoreflect.Sint64Kind:
+		g.P("enc.AddInt64(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	case protoreflect.GroupKind:
+		g.P("enc.AddReflected(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	case protoreflect.MessageKind:
+		g.P("if obj, ok := interface{}(v).(", g.QualifiedGoIdent(zapcorePkg.Ident("ObjectMarshaler")), "); ok {")
+		g.P("enc.AddObject(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), obj)")
+		g.P("} else {")
+		g.P("enc.AddReflected(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+		g.P("}")
+	case protoreflect.StringKind:
+		g.P("enc.AddString(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	default:
+		g.P("enc.Addlected(", g.QualifiedGoIdent(fmtPkg.Ident("Sprintf")), "(\"%v\", k), v)")
+	}
+	g.P("}")
+	g.P("return nil")
+	g.P("}))")
+	g.P()
+}
+
+func generatePrimitiveField(g *protogen.GeneratedFile, f *protogen.Field) {
+	fname := f.Desc.Name()
+	var gname string
+	if f.Oneof != nil {
+		gname = fmt.Sprintf("Get%s()", f.GoName)
+	} else {
+		gname = f.GoName
+	}
+	switch f.Desc.Kind() {
+	case protoreflect.BoolKind:
+		g.P("enc.AddBool(\"", fname, "\", x.", gname, ")")
+	case protoreflect.BytesKind:
+		g.P("enc.AddBinary(\"", fname, "\", x.", gname, ")")
+	case protoreflect.DoubleKind:
+		g.P("enc.AddFloat64(\"", fname, "\", x.", gname, ")")
+	case protoreflect.EnumKind:
+		g.P("enc.AddString(\"", fname, "\", x.", gname, ".String())")
+	case protoreflect.Fixed32Kind, protoreflect.Uint32Kind:
+		g.P("enc.AddUint32(\"", fname, "\", x.", gname, ")")
+	case protoreflect.Fixed64Kind, protoreflect.Uint64Kind:
+		g.P("enc.AddUint64(\"", fname, "\", x.", gname, ")")
+	case protoreflect.FloatKind:
+		g.P("enc.AddFloat32(\"", fname, "\", x.", gname, ")")
+	case protoreflect.Int32Kind, protoreflect.Sfixed32Kind, protoreflect.Sint32Kind:
+		g.P("enc.AddInt32(\"", fname, "\", x.", gname, ")")
+	case protoreflect.Int64Kind, protoreflect.Sfixed64Kind, protoreflect.Sint64Kind:
+		g.P("enc.AddInt64(\"", fname, "\", x.", gname, ")")
+	case protoreflect.GroupKind:
+		g.P("enc.AddReflected(\"", fname, "\", x.", gname, ")")
+	case protoreflect.MessageKind:
+		g.P("if obj, ok := interface{}(x.", gname, ").(", g.QualifiedGoIdent(zapcorePkg.Ident("ObjectMarshaler")), "); ok {")
+		g.P("enc.AddObject(\"", fname, "\", obj)")
+		g.P("} else {")
+		g.P("enc.AddReflected(\"", fname, "\", x.", gname, ")")
+		g.P("}")
+	case protoreflect.StringKind:
+		g.P("enc.AddString(\"", fname, "\", x.", gname, ")")
+	default:
+		g.P("enc.AddReflected(\"", fname, "\", x.", gname, ")")
+	}
+	g.P()
+}
+
+func isMasked(opts *descriptorpb.FieldOptions) bool {
+	return proto.GetExtension(opts, pbzap.E_Mask).(bool)
+}
+
+func generateMessage(g *protogen.GeneratedFile, m *protogen.Message) {
+	ident := g.QualifiedGoIdent(m.GoIdent)
+	g.P("func (x *", ident, ") MarshalLogObject(enc ", g.QualifiedGoIdent(zapcorePkg.Ident("ObjectEncoder")), ") error {")
+	g.P("if x == nil {")
+	g.P("return nil")
+	g.P("}")
+	g.P()
+	for _, f := range m.Fields {
+		if isMasked(f.Desc.Options().(*descriptorpb.FieldOptions)) {
+			g.P("enc.AddString(\"", f.Desc.Name(), "\", \"[MASKED]\")")
+			g.P()
+		} else if f.Desc.IsList() {
+			generateListField(g, f)
+		} else if f.Desc.IsMap() {
+			generateMapField(g, f)
+		} else {
+			generatePrimitiveField(g, f)
+		}
+	}
+	g.P("return nil")
+	g.P("}")
+	g.P()
+	for _, submsg := range m.Messages {
+		if submsg.Desc.IsMapEntry() {
+			continue
+		}
+		generateMessage(g, submsg)
+	}
+}
+
+func generateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
+	if len(file.Messages) == 0 {
+		return nil
+	}
+
+	filename := fmt.Sprintf("%s_marshall_zap.pb.go", file.GeneratedFilenamePrefix)
+	g := gen.NewGeneratedFile(filename, file.GoImportPath)
+	g.P("// Code generated by protoc-gen-marshall-zap. DO NOT EDIT.")
+	g.P("//")
+	g.P("// source: ", file.Desc.Path())
+	g.P()
+	g.P("package ", file.GoPackageName)
+	g.P()
+
+	for _, m := range file.Messages {
+		generateMessage(g, m)
+	}
+
+	return g
+}
+
 func main() {
-	pgs.Init(pgs.DebugEnv("DEBUG")).
-		RegisterModule(&module{}).
-		RegisterPostProcessor(pgsgo.GoFmt()).
-		Render()
+	protogen.Options{}.Run(func(plugin *protogen.Plugin) error {
+		for _, file := range plugin.FilesByPath {
+			if !file.Generate {
+				continue
+			}
+
+			generateFile(plugin, file)
+		}
+		return nil
+	})
 }