Goの抽象構文木（AST）でBoilerplate Codeを自動生成する

この記事は、Merpay Advent Calendar 2022 の19日目の記事です。

　こんにちは。メルペイのPayment PlatformチームでPaymentServiceの開発を担当するエンジニアの @huhu です。主にPaymentServiceに関わるマイクロサービスの開発や運用をしています。

　この記事では、私のチームで行ったGoの抽象構文木（AST）でコードを自動生成する取り組みに関して紹介します。

背景

　PaymentServiceは決済トランザクション管理の基盤サービスとして、エッジレイヤーのサービス（外部サービスも含め）が提供する各種決済手段を利用して、表層のサービス（メルカリ、NFC、コード払いなど）に必要な決済フローを共通APIとして提供しています。我々の提供している多彩な決済手段を表すEntityが存在していて、例えば支払い方法ではお客さまの残高の決済を表すBalancePaymentや、クレカ決済を表すCreditCardPaymentなど、似たような用途のための構造体がたくさん存在しています。

　このような Entity に対する処理は、以前@foghostが紹介したマイクロサービスにおける決済トランザクション管理で、複数のフェーズに細分化して実行するようになっています。Entityは現在どの状態を実行中かを表すフィールド、メソッドが存在しています。

https://engineering.mercari.com/blog/entry/2019-06-07-155849/

　例えばデータが作成開始段階にあることを表すIsCreating()のようなメソッドや、Refund処理のために今現在Paidになっているかどうかを表すIsPaid()のようなメソッドもあります。Entityごとに状態の遷移が異なっているため、それぞれのEntityに対し、このようなBoilerplate Codeを大量に書き、かつメンテナンスする必要があります。
　また、遷移状態が複雑になったり、あるいは単にEntityとメソッドの数が増えると、人間による作業ミス（漏れなど）が非常に発生しやすく、メンテナンス（コードのスタイルやアップデートなど）もしにくくなります。

　そこで、このような問題を解決するために、構造体の特定のメソッドに対し、コードの自動生成やアップデートの方法が必要です。今GoのGenerics(ジェネリクス)はまだ構造体に対応していないため、この部分の抽象化はなかなか難しいです。従ってメタプログラミングでここのコードを生成したいと考えるのは自然な成り行きです。Goではコードを生成するには普通templateのパッケージを利用しますが、文字列でメソッドの修正や生成はあまり賢くはないので、生成したファイルも通常人間による編集が禁止されることになります。今回は人間による編集も併せて行いたいケースがありました。またコードの組み合わせも文字列の組み合わせで表現するのはややこしいため、今回はtemplateを使いませんでした。

今回の取り組み

要件を整理してみます。

• ボイラープレートコード軽減のためのコード生成
• 構造体に特定のメソッドを自動生成したい
• 既に人間が書いたメソッドが存在する場合、設定によってアップデートしないこともできるようにする

つまり、Goのコードをデータとして読んで、必要なメソッドの定義を生成する必要があります。ここは抽象構文木の出番です！

関数の抽象構文木

　要件からまずは構造体の特定のメソッドを見つける必要があります。正規表現などを使うより、抽象構文木を使ってメソッドを見つけましょう。Goではオフィシャルのパッケージgo/astがあります。そしてパッケージの中にはGoのコードの抽象構文木を表す構造体の定義があります。例えば以下は関数定義の抽象構文木の定義です。

type FuncDecl struct {
    Doc  *CommentGroup // associated documentation; or nil
    Recv *FieldList    // receiver (methods); or nil (functions)
    Name *Ident        // function/method name
    Type *FuncType     // function signature: type and value parameters, results, and position of "func" keyword
    Body *BlockStmt    // function body; or nil for external (non-Go) function
}

　この構造体はGoの関数のコードを表すことができます。例えばHello WorldのFuncDeclは以下となります。

func main() {
    fmt.Println("hello world")
}

&ast.FuncDecl{
    Name: &ast.Ident{
        Name: "main",
    },
    Type: &ast.FuncType{
        Params: &ast.FieldList{},
    },
    Body: &ast.BlockStmt{
        List: []ast.Stmt{
            &ast.ExprStmt{
                X: &ast.CallExpr{
                    Fun: &ast.SelectorExpr{
                        X: &ast.Ident{
                            Name: "fmt",
                        },
                        Sel: &ast.Ident{
                            Name: "Println",
                        },
                    },
                    Args: []ast.Expr{
                        &ast.BasicLit{
                            Kind: token.STRING,
                            Value: "\"hello world\"",
                        },
                    },
                },
            },
        },
    },
}

ファイルの抽象構文木の取得

　コードの中からある構造体のメソッドを見つけるには、Goのファイルの抽象構文木を取得する必要があります。go/parserパッケージで簡単にファイルの抽象構文木を取得することができます。

import (
    "go/parser"
    "go/token"
)

fun main() {
    fset := token.NewFileSet()
    node, err := parser.ParseFile(fset, "filename.go", nil, parser.AllErrors)
    if err != nil {
        panic(err)
    }
}

メソッドを見つける

　ファイル全体の抽象構文木からメソッドを見つけるのは難しくありません。名前の通りツリー構造であるため、golang.org/x/tools/go/ast/astutilというパッケージでは再帰的な方式でファイルの抽象構文木の全てノードを​​探索かつ操作することが可能です。ノードのTypeを判断し、FuncDeclであればNameとレシーバーをチェックし、メソッドを見つけることができます。

astutil.Apply(node, nil, func(c *astutil.Cursor) bool {
    n := c.Node()
    switch x := n.(type) {
    case *ast.FuncDecl:
        if x.Name.Name == “IsPaying” && x.Recv != nil {
               …
           }
    }

    return true
})

　メソッドを見つけたら、次はメソッドの内容を更新します。また事前に設定を書き、もしすでに存在したら変更をスキップすることもできます。x.Body.Listはメソッドのボディです、これを操作すると、メソッドの内容を変更することができます。ボディの内容を解析し、細部を変更することもできますが、直接メソッドのボディ全体を差し替えることも可能です。今回はこの方法でメソッドをアップデートします

astutil.Apply(node, nil, func(c *astutil.Cursor) bool {
    n := c.Node()
    switch x := n.(type) {
    case *ast.FuncDecl:
        if x.Name.Name == “IsPaying” && x.Recv != nil {
               c.Replace(f)
           }
    }

    return true
})

変数fは見つけたFuncDeclを置き換える内容になっています。例えば以下のメソッドに置き換えたい場合を考えます。

func (p *Payment) IsPaid() bool {
    return p.done
}

fは以下となります。メソッドの名前をisPayingからisPaidに変更し、Bodyの内容も新しいコンテンツに変更します。

f := &ast.FuncDecl{
    Recv: &ast.FieldList{
        List: []*ast.Field{
            &ast.Field{
                Names: []*ast.Ident{
                    &ast.Ident{
                        Name: "p",
                    },
                },
                Type: &ast.StarExpr{
                    X: &ast.Ident{
                        Name: "Payment",
                    },
                },
            },
        },
    },
    Name: &ast.Ident{
        Name: "IsPaid",
    },
    Type: &ast.FuncType{
        Params: &ast.FieldList{},
        Results: &ast.FieldList{
            List: []*ast.Field{
                &ast.Field{
                    Type: &ast.Ident{
                        Name: "bool",
                    },
                },
            },
        },
    },
    Body: &ast.BlockStmt{
        List: []ast.Stmt{
            &ast.ReturnStmt{
                Results: []ast.Expr{
                    &ast.SelectorExpr{
                        X: &ast.Ident{
                            Name: "p",
                        },
                        Sel: &ast.Ident{
                            Name: "done",
                        },
                    },
                },
            },
        },
    },
}

例えは以下のファィルであれば、

package payment

type Payment struct {
    doing bool
    done bool
}

func (p *Payment) IsPaying() bool {
    return p.doing
}

　処理を実行したら以下のように変更できます

package payment

type Payment struct {
    doing bool
    done bool
}

func (p *Payment) IsPaid() bool {
    return p.done
}

　fの内容はプログラム内で自由に組み立てられるので、うまく操れば、非常に柔軟に生成するプログラムを決定できます。いわゆるメタプログラミングのような効果が実現できます。
　もし元の関数が見つけられない場合は、新たな関数を定義してやる必要があります。これも非常に簡単でトップレベルのノードのdeclarationsに関数を追加すればできます。

node.Decls = append(node.Decls, f)

最後に、元のファイルに出力すれば全ての操作が完了です。

f, err := os.Create(fileName)
if err != nil {
    panic(err)
}

format.Node(f, fset, node)

まとめ

　今回はGoの抽象構文木（AST）でコードを自動生成と題して私のチームでのboilerplateのcodeの管理に関する取り組みを紹介しました。

　Goは構造体のGenericsがまだ存在しない、またOOPのようなコーディングもできません。生成機能によって一定程度コードの再利用効果を達成することができます。かつそれをよりうまく管理するために、抽象構文木を使いました。また実際メソッドだけではなく、構造体やインターフェースの定義、Importの処理ほぼ全てにも使えます。

　今回紹介した方法がもしご参考になれば幸いです。
　明日の記事は @pedroさんです。引き続きお楽しみください。