Goで特定のパターンのリファクタリングをASTを弄って自動化した

これは日記です。技術記事ではないので読みやすくはないです。

仕事のコードで特定のパターンでちょっと泥臭く書き換える必要のあるリファクタリングが必要になっているのだけれど、単純計算で100ファイル、1000箇所以上の書き換えが必要になっている。これまでちまちま手作業でいろんな人が片手間で書き換えをやっていたのだけれど、無限に時間がかかりそうだったのでどうにか出来ないかと考えていた。

先行研究として id:hitode909 のASTを使ってリファクタリングするやつが記憶に残っていたのでもうちょっと簡易版で似たようなことをやってみた。

speakerdeck.com

GoのAST周りはあんまり詳しくなくて、社内の静的解析ツールをちょっと弄ったくらいだったのでまずは勉強した。

motemen.github.io

id:motemen さんの Go のための Go を読んでふむふむという感じで理解した気になってとりかかるも、そんなにすぐに身に付くものでもないので ChatGPT (GPT-4) に相談しながら進める。

本格的なツールを作るつもりじゃなくて、リファクタリングを補助したいだけなので以下のような作戦でいった。

大まかな書き換えだけをやる
書き換えた箇所に FIXME コメントを残して確認出来るようにする
書き換えたことによって使わなくなった変数とかは手動で直す

これくらいのことならそんなにがんばらなくても出来そう。

最初は ast.Inspect しか知らなくて書き換えるの大変だなあと思っていたのだけれど、 astutil.Apply というのが使えると知ったのでこちらで進める。書き換え前の状態と、書き換え後の状態でそれぞれコールバックを渡せるが、今回は書き換え前のみ使う。

astutil.Apply(f, func(c *astutil.Cursor) bool {
    // TODO: ここで対象を見つけて書き換える
}, nil)

コールバックの引数に渡ってくる astutil.Cursor が Replace とか Deleta を持っているのでこれを使って書き換えていく。

書き換え対象を見つけるにはいま着目しているのがなんの型なのかを調べる必要がる。Cursorから取得できるNodeは ast.Node interface の状態なので型をちまちま見ていく。

n := c.Node()
switch x := n.(type) {
    case *ast.AssignStmt:
        // 代入だったらここ
    case *ast.CallExpr:
        // メソッド呼び出しだったらここ
}

ここからさらに詳細な型にキャストしていく必要がある。例えばメソッド呼び出しだった場合は以下のような感じでレシーバを取得できる。

if sel, ok := x.Fun.(*ast.SelectorExpr); ok {
    if ident, ok := sel.X.(*ast.Ident); ok {
        if ident.Name == "hoge" {
            // hoge.(何かのメソッド) という呼び出し
        }
    }
}

型は素晴らしいのだけれど、こういうちょっとしたことをするにはちょっと面倒。

メソッド呼び出しの場合は引数などの情報を持っているのでこれも使う。

for _, arg := range x.Args {
    // もちろん arg は ast.Node 型なのでここでも型を見極める必要がある
}

元の呼び出しを新しいメソッドの呼び出しに変えるにはASTを手動で組み立てて先述の Replace メソッドを使う。

replaced := &ast.CallExpr{
    Fun: &ast.SelectorExpr{
        X: ast.NewIdent("hoge"),
        Sel: ast.NewIdent("Bar"),
    },
    Args: x.Args,
}
c.Replace(replaced)

これで例えば hoge.Piyo みたいな呼び出しだったのが hoge.Bar に書き換えられる。この例では引数はそのままにしているが、実際にはここも弄ることになったのでちょっと泥臭い感じに。

コメントは別途 ast.File のフィールドとして持っているので、適切なポジションを与えつつ append するだけでよい。ただし parse.ParseFile するときの mode 引数に parser.ParseComments を与えておく必要がある。

f.Comments = append(f.Comments, &ast.CommentGroup{
    List: []*ast.Comment{
        {
            Slash: pos,
            Text:  "// FIXME: 自動で置き換えたので確認してください",
        },
    },
})

append するだけで良いと言ったが実はこれでは駄目で、最後に出力するときに正しい位置に配置するにはポジション順にソートされている必要がある。これにしばらくハマってしまった。

sort.Slice(f.Comments, func(i, j int) bool {
    return f.Comments[i].Pos() < f.Comments[j].Pos()
})

最後にASTをGoのコードとして正しくフォーマットしつつ元のファイルに書き戻すには format.Node を使う。第一引数を os.Stdout にすると標準出力に出せるのでdry runとかを用意して確認用に出力すると良さそうです。

file, err := os.Create(filename)
if err != nil {
    panic(err)
}
defer file.Close()
format.Node(file, fset, f)

ここで少し問題があって、各Nodeが持っているポジションは元のコードの位置なので、例えば、引数の数を減らして書き戻したりすると不要な空行が生まれたりします。この辺りは雑にやってしまったのでまだうまい解決方法を見つけていなくて課題になっています。いい方法があったら教えてください。

だいぶ雑ながらも人力で置き換え箇所を見つけてちまちま書き換える必要がなくなってだいぶ楽になりました。引数の数が可変だったり、残すもの、残さないものがあったりと人力でやるにも機械的にはやりにくいところだったのではかどりそうです。

ちなみに

火曜日の空は僕を押しつぶした。

Goで特定のパターンのリファクタリングをASTを弄って自動化した