plugin: move generator to separate package

Signed-off-by: Stephen J Day <stephen.day@docker.com>

plugin/generator.go

@@ -0,0 +1,125 @@
+package plugin
+
+import (
+	"strings"
+
+	"github.com/gogo/protobuf/protoc-gen-gogo/descriptor"
+	"github.com/gogo/protobuf/protoc-gen-gogo/generator"
+)
+
+type ttrpcGenerator struct {
+	*generator.Generator
+	generator.PluginImports
+
+	typeurlPkg generator.Single
+	ttrpcPkg   generator.Single
+	contextPkg generator.Single
+}
+
+func init() {
+	generator.RegisterPlugin(new(ttrpcGenerator))
+}
+
+func (p *ttrpcGenerator) Name() string {
+	return "ttrpc"
+}
+
+func (p *ttrpcGenerator) Init(g *generator.Generator) {
+	p.Generator = g
+}
+
+func (p *ttrpcGenerator) Generate(file *generator.FileDescriptor) {
+	p.PluginImports = generator.NewPluginImports(p.Generator)
+	p.contextPkg = p.NewImport("context")
+	p.typeurlPkg = p.NewImport("github.com/containerd/typeurl")
+	p.ttrpcPkg = p.NewImport("github.com/stevvooe/ttrpc")
+
+	for _, service := range file.GetService() {
+		serviceName := service.GetName()
+		if pkg := file.GetPackage(); pkg != "" {
+			serviceName = pkg + "." + serviceName
+		}
+
+		p.genService(serviceName, service)
+	}
+}
+
+func (p *ttrpcGenerator) genService(fullName string, service *descriptor.ServiceDescriptorProto) {
+	serviceName := service.GetName() + "Service"
+	p.P()
+	p.P("type ", serviceName, " interface{")
+	p.In()
+	for _, method := range service.Method {
+		p.P(method.GetName(),
+			"(ctx ", p.contextPkg.Use(), ".Context, ",
+			"req *", p.typeName(method.GetInputType()), ") ",
+			"(*", p.typeName(method.GetOutputType()), ", error)")
+
+	}
+	p.Out()
+	p.P("}")
+
+	p.P()
+	// registration method
+	p.P("func Register", serviceName, "(srv *", p.ttrpcPkg.Use(), ".Server, svc ", serviceName, ") error {")
+	p.In()
+	p.P(`return srv.Register("`, fullName, `", map[string]`, p.ttrpcPkg.Use(), ".Handler{")
+	p.In()
+	for _, method := range service.Method {
+		p.P(`"`, method.GetName(), `": `, p.ttrpcPkg.Use(), `.HandlerFunc(func(ctx context.Context, req interface{}) (interface{}, error) {`)
+		p.In()
+		p.P("return svc.", method.GetName(), "(ctx, req.(*", p.typeName(method.GetInputType()), "))")
+		p.Out()
+		p.P("}),")
+	}
+	p.Out()
+	p.P("})")
+	p.Out()
+	p.P("}")
+
+	clientType := service.GetName() + "Client"
+	clientStructType := strings.ToLower(clientType[:1]) + clientType[1:]
+	p.P()
+	p.P("type ", clientStructType, " struct{")
+	p.In()
+	p.P("client *", p.ttrpcPkg.Use(), ".Client")
+	p.Out()
+	p.P("}")
+	p.P()
+	p.P("func New", clientType, "(client *", p.ttrpcPkg.Use(), ".Client)", serviceName, "{")
+	p.In()
+	p.P("return &", clientStructType, "{")
+	p.In()
+	p.P("client: client,")
+	p.Out()
+	p.P("}")
+	p.Out()
+	p.P("}")
+	p.P()
+	for _, method := range service.Method {
+		p.P()
+		p.P("func (c *", clientStructType, ") ", method.GetName(),
+			"(ctx ", p.contextPkg.Use(), ".Context, ",
+			"req *", p.typeName(method.GetInputType()), ") ",
+			"(*", p.typeName(method.GetOutputType()), ", error) {")
+		p.In()
+		p.P("resp, err := c.client.Call(ctx, ", `"`+fullName+`", `, `"`+method.GetName()+`"`, ", req)")
+		p.P("if err != nil {")
+		p.In()
+		p.P("return nil, err")
+		p.Out()
+		p.P("}")
+		p.P("return resp.(*", p.typeName(method.GetOutputType()), "), nil")
+		p.Out()
+		p.P("}")
+	}
+}
+
+func (p *ttrpcGenerator) objectNamed(name string) generator.Object {
+	p.Generator.RecordTypeUse(name)
+	return p.Generator.ObjectNamed(name)
+}
+
+func (p *ttrpcGenerator) typeName(str string) string {
+	return p.Generator.TypeName(p.objectNamed(str))
+}