goパッケージを使って複数ファイルを1つにまとめる gma を作った

これは Go2 advent calendar 16日目の記事です。

モチベーション

最近競技プログラミングをやり始めたが、web上のエディタで書いたコードをそのままsubmitする形式が多いように思う。つまりシングルファイルにまとめる必要が出てくるが、いくつか問題を解いていると似たような処理が多くなりutilファイルが欲しくなってきたので作った。

作ったもの

go パッケージを使って複数のファイルを1つにまとめる gma というツールを作った。

何ができるか

単機能なのでREADMEに書いている以上のことはないが、現状以下のように複数のファイルがあったとき、良い感じにシングルファイルにまとめてくれる。

元ファイルたち:

$ tree example/
example/
├── main.go
├── util
│   └── util.go
└── util.go

package main

import (
        "fmt"

        "github.com/takashabe/go-main-aggregator/example/util"
)

func main() {
        fmt.Println(util.Foo())
        Foo()
}

example/util.go:

package main

func Foo() {}

example/util/util.go:

package util

import "fmt"

func Foo() string {
        return fmt.Sprintf("util")
}

結合する:

$ gma -main main.go -depends util.go,util2.go -depends util/util.go
package main

import (
        "fmt"
)

func main() {
        fmt.Println(_util_Foo())
        Foo()
}
func Foo() {
}
func _util_Foo() string {
        return fmt.Sprintf("util")
}

"github.com/takashabe/go-main-aggregator/example/util" への依存が無くなり、単体で実行可能になっていることが分かると思う。

実装について

今回のように他パッケージのファイルを扱おうとするとシングルファイルにしたときimportの問題が出てくるので、ASTをゴニョゴニョするのが良さそうというのがわかる。 goでASTを扱うことについては日本語でも素晴らしい資料があるのでそれらを参照するのが良いと思う。特に参考にさせてもらったのは以下。

• https://qiita.com/tenntenn/items/a312d2c5381e36cf4cd3
• https://motemen.github.io/go-for-go-book/

ここでは gma を実装する上で特にハマった複数ファイルの結合と、外部パッケージの外部関数呼び出しをローカル呼び出しに変換する部分について、実際のコードから抜き出して紹介したい。

複数ファイルの結合

インポートや変数、型、関数定義といった宣言は全て ast.Decl インタフェースを実装している。そのためファイルごとのASTを得たら、それらから ast.Decl を抽出して新しい *ast.File とすれば良い。

以下ではimportを別に解決したかったのでそれだけ個別に除外しているが、雰囲気は伝わると思う。

func mergeFiles(files []*ast.File) (*ast.File, error) {
...
  decls := []ast.Decl{}
  for _, file := range files {
    for _, d := range file.Decls {
    g, ok := d.(*ast.GenDecl)
    if ok && g.Tok == token.IMPORT {
      continue
    }
    decls = append(decls, d)
    }
  }

  file := &ast.File{
    Package: files[0].Package,
    Name:    files[0].Name,
    Decls:   decls,
  }
...

上記コードでは触れていないが、複数パッケージを対象にしたときに同名の宣言があるとコンフリクトするので、実際には予めパッケージごとにユニークになるように名前を変換している。

外部関数呼び出しの変換

関数呼び出しをしているのは *ast.CallExpr で、更にその中で外部パッケージの関数呼び出しをしているのは callExpr.Fun が *ast.SelectorExpr のものになる。

以下では条件に合致するnodeを探索している。ASTをいじるときは大体こんな感じで必要なnodeをwalkなりして探してゴニョゴニョというコードが多くなると思う。

func(c *astutil.Cursor) bool {
  n := c.Node()

  callExpr, ok := n.(*ast.CallExpr)
  if !ok {
    return true
  }
  selector, ok := callExpr.Fun.(*ast.SelectorExpr)
  if !ok {
    return true
  }
  x, ok := selector.X.(*ast.Ident)
  if !ok {
    return true
  }
...

外部関数呼び出しを行っているnodeが特定できたら、あとは実際にそれが変換する必要があるかどうかを判定して、変換の必要があれば astutil パッケージの Cursor.Replace でnodeの変換を行っている。(この用途なら普通にast.Walkしてnodeのフィールドを上書きしても良かったかもしれない)

実際には事前に結合したファイル側の関数を全て抜き出して、変換対象リストを作ってmainファイル側で呼び出される関数ごとに突き合わせを行っている。泥臭い感じのコードになっているがもし興味があればリポジトリを見てもらえればと思う。

  cn := &ast.CallExpr{
    Fun:  repNode.Name,
    Args: callExpr.Args,
  }
  c.Replace(cn)
  return true
}

astutil パッケージは golang.org/x/tools 配下にある準公式っぽいやつで、使い方はテストを見ると何となく分かると思う。 https://github.com/golang/tools/blob/master/go/ast/astutil/rewrite_test.go

まとめ

go パッケージを使って複数のファイルを1つにまとめる gma というツールを作った。またその過程で得たASTの扱いなどについて紹介した。

まともにASTを触ったのは初めてだったが、goではASTにアクセスするためのインタフェースが揃っているのでとても楽だった。

今回のようにAST経由で変換して何かしたいといった場合、要素ごとにそれを表すASTノードが何であるかを把握出来るとあとは整合性を保って変換していくだけなので、各ノードの関係性が分かるとそれなりに動くものが作れそうな気がする。

またgoでは公式ツール内で go パッケージを使っているものも多く、特に cmd/gofmt や golang/x/tools/cmd は非常に参考にさせてもらった。