【AI Shift Advent Calendar 2023】connect-go から学ぶコード生成

// https://github.com/connectrpc/connect-go/blob/main/internal/gen/connect/ping/v1/pingv1connect/ping.connect.go

type PingServiceClient interface {
    // Ping sends a ping to the server to determine if it's reachable.
    Ping(context.Context, *connect.Request[v1.PingRequest]) (*connect.Response[v1.PingResponse], error)
    // Fail always fails.
    Fail(context.Context, *connect.Request[v1.FailRequest]) (*connect.Response[v1.FailResponse], error)
    // Sum calculates the sum of the numbers sent on the stream.
    Sum(context.Context) *connect.ClientStreamForClient[v1.SumRequest, v1.SumResponse]
    // CountUp returns a stream of the numbers up to the given request.
    CountUp(context.Context, *connect.Request[v1.CountUpRequest]) (*connect.ServerStreamForClient[v1.CountUpResponse], error)
    // CumSum determines the cumulative sum of all the numbers sent on the stream.
    CumSum(context.Context) *connect.BidiStreamForClient[v1.CumSumRequest, v1.CumSumResponse]
}

// ...

type PingServiceHandler interface {
    // Ping sends a ping to the server to determine if it's reachable.
    Ping(context.Context, *connect.Request[v1.PingRequest]) (*connect.Response[v1.PingResponse], error)
    // Fail always fails.
    Fail(context.Context, *connect.Request[v1.FailRequest]) (*connect.Response[v1.FailResponse], error)
    // Sum calculates the sum of the numbers sent on the stream.
    Sum(context.Context, *connect.ClientStream[v1.SumRequest]) (*connect.Response[v1.SumResponse], error)
    // CountUp returns a stream of the numbers up to the given request.
    CountUp(context.Context, *connect.Request[v1.CountUpRequest], *connect.ServerStream[v1.CountUpResponse]) error
    // CumSum determines the cumulative sum of all the numbers sent on the stream.
    CumSum(context.Context, *connect.BidiStream[v1.CumSumRequest, v1.CumSumResponse]) error
}

$ protoc --go_out=gen --connect-go_out=gen path/to/file.proto

version: v1
plugins:
  - name: go
      out: gen
  - name: connect-go
    out: gen

$ buf generate
-->generated
//  gen/path/to/file.pb.go
//  gen/path/to/connectfoov1/file.connect.go

// cmd/protoc-gen-connect-go/main.go
package main

import (
　// ....
    connect "connectrpc.com/connect"
    "google.golang.org/protobuf/compiler/protogen"
    "google.golang.org/protobuf/reflect/protoreflect"
    "google.golang.org/protobuf/types/descriptorpb"
    "google.golang.org/protobuf/types/pluginpb"
)

func main() {
    // ...
    protogen.Options{}.Run(
        func(plugin *protogen.Plugin) error {
            plugin.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
            for _, file := range plugin.Files {
                if file.Generate {
                    generate(plugin, file)
                }
            }
            return nil
        },
    )
}

protobufのcompiler/protogenを使用して必要な構造体を生成
↓
connect-goのgenerate関数に構造体を投げて書き込み
↓
出力

// cmd/protoc-gen-connect-go/main.go
func generate(plugin *protogen.Plugin, file *protogen.File) {
  // ...
    generatedFile := plugin.NewGeneratedFile(
        file.GeneratedFilenamePrefix+generatedFilenameExtension,
        protogen.GoImportPath(path.Join(
            string(file.GoImportPath),
            string(file.GoPackageName),
        )),
    )
    generatedFile.Import(file.GoImportPath)
    generatePreamble(generatedFile, file)
    generateServiceNameConstants(generatedFile, file.Services)
    generateServiceNameVariables(generatedFile, file)
    for _, service := range file.Services {
        generateService(generatedFile, service)
    }
}

// cmd/protoc-gen-connect-go/main.go
func generateServerInterface(g *protogen.GeneratedFile, service *protogen.Service, names names) {
    wrapComments(g, names.Server, " is an implementation of the ", service.Desc.FullName(), " service.")
    if isDeprecatedService(service) {
        g.P("//")
        deprecated(g)
    }
    g.AnnotateSymbol(names.Server, protogen.Annotation{Location: service.Location})
    g.P("type ", names.Server, " interface {")
    for _, method := range service.Methods {
        leadingComments(
            g,
            method.Comments.Leading,
            isDeprecatedMethod(method),
        )
        g.AnnotateSymbol(names.Server+"."+method.GoName, protogen.Annotation{Location: method.Location})
        g.P(serverSignature(g, method))
    }
    g.P("}")
    g.P()
}

// protobuf-go/compiler/protogen/protogen.go
// P prints a line to the generated output. It converts each parameter to a
// string following the same rules as fmt.Print. It never inserts spaces
// between parameters.
func (g *GeneratedFile) P(v ...interface{}) {
    for _, x := range v {
        switch x := x.(type) {
        case GoIdent:
            fmt.Fprint(&g.buf, g.QualifiedGoIdent(x))
        default:
            fmt.Fprint(&g.buf, x)
        }
    }
    fmt.Fprintln(&g.buf)
}

// cmd/protoc-gen-connect-go/main.go
// Raggedy comments in the generated code are driving me insane. This
// word-wrapping function is ruinously inefficient, but it gets the job done.
func wrapComments(g *protogen.GeneratedFile, elems ...any) {
    text := &bytes.Buffer{}
    for _, el := range elems {
        switch el := el.(type) {
        case protogen.GoIdent:
            fmt.Fprint(text, g.QualifiedGoIdent(el))
        default:
            fmt.Fprint(text, el)
        }
    }
    words := strings.Fields(text.String())
    text.Reset()
    var pos int
    for _, word := range words {
        numRunes := utf8.RuneCountInString(word)
        if pos > 0 && pos+numRunes+1 > commentWidth {
            g.P("// ", text.String())
            text.Reset()
            pos = 0
        }
        if pos > 0 {
            text.WriteRune(' ')
            pos++
        }
        text.WriteString(word)
        pos += numRunes
    }
    if text.Len() > 0 {
        g.P("// ", text.String())
    }
}

text := &bytes.Buffer{}
for _, el := range elems {
    switch el := el.(type) {
    case protogen.GoIdent:
        fmt.Fprint(text, g.QualifiedGoIdent(el))
    default:
        fmt.Fprint(text, el)
    }
}
words := strings.Fields(text.String())

func (b *Buffer) Reset() {
    b.buf = b.buf[:0]
    b.off = 0
    b.lastRead = opInvalid
}

var pos int
for _, word := range words {
    numRunes := utf8.RuneCountInString(word)
    if pos > 0 && pos+numRunes+1 > commentWidth {
        g.P("// ", text.String())
        text.Reset()
        pos = 0
    }
    if pos > 0 {
        text.WriteRune(' ')
        pos++
    }
    text.WriteString(word)
    pos += numRunes
}
if text.Len() > 0 {
    g.P("// ", text.String())
}