Posted at

抽象構文木(AST)をいじってフォーマットをかける #golang

More than 1 year has passed since last update.


はじめに

Goには、gofmtというコマンドがあることはご存知かと思います。

そして、そのフォーマットがGoコミュニティの標準となっています。

go/formatパッケージでは、gofmtと同じスタイルのフォーマットでソースコードを整形することができます。

また、go/formatパッケージを使うと、ASTについても整形することがきます。

この記事では、ASTのノードをいじって、整形する方法について説明します。

なお、この記事を書いた2016年12月現在のGoのバージョンは1.7.4が最新です。


ASTをいじる

ASTをいじると書きましたが、具体的に何をするんでしょうか?

ここでは、AST上のノードを別のノードに入れ替えたり、ノードのフィールドを差し替えることを指しています。

たとえば、x+yという式があった場合、このxyをそれぞれ1020に置き換えて、10+20のような式に変えたいとします。

x+yは以下のように構成されるASTになるはずなので、これの*ast.Identの部分を入れ替えてやればよいでしょう。

*ast.BinaryExpr (+)

├── *ast.Ident (x)
└── *ast.Ident (y)

ast.BinaryExpr構造体は2項演算式を表し、オペランドはXフィールドとYフィールドにast.Expr型の値として保持されます。

つまり、XフィールドとYフィールドを差し替えて、以下のような構成のASTにしてやればよいでしょう。

*ast.BinaryExpr (+)

├── *ast.BasicLit (10)
└── *ast.BasicLit (20)

ast.BasicLit構造体は、以下のようなフィールドを持つ構造体です。

type BasicLit struct {

ValuePos token.Pos // literal position
Kind token.Token // token.INT, token.FLOAT, token.IMAG, token.CHAR, or token.STRING
Value string // literal string; e.g. 42, 0x7f, 3.14, 1e-9, 2.4i, 'a', '\x7f', "foo" or `\m\n\o`
}

そのため、Kindフィールドをtoken.INTにして、Valueフィールドを"10""20"にしたものをオペランドとして差し替えてやればOKです。

この時、ValuePosフィールドについては、ゼロ値、つまり0としておいて構いません。なお、値が0token.Postoken.NoPosという名前の定数として定義されています。

上記の処理をコードに落とすと以下のようになります。

fset := token.NewFileSet()

expr, err := parser.ParseExprFrom(fset, "sample.go", `x+y`, 0)
if err != nil {
log.Fatalln("Error:", err)
}

fmt.Println("==== BEFORE ====")
ast.Print(fset, expr)

binaryExpr := expr.(*ast.BinaryExpr)

binaryExpr.X = &ast.BasicLit{
Kind: token.INT,
Value: "10",
}

binaryExpr.Y = &ast.BasicLit{
Kind: token.INT,
Value: "20",
}

fmt.Println("==== AFTER ====")
ast.Print(fset, expr)

そして、出力結果は以下のようになります。

==== BEFORE ====

0 *ast.BinaryExpr {
1 . X: *ast.Ident {
2 . . NamePos: sample.go:1:1
3 . . Name: "x"
4 . . Obj: *ast.Object {
5 . . . Kind: bad
6 . . . Name: ""
7 . . }
8 . }
9 . OpPos: sample.go:1:2
10 . Op: +
11 . Y: *ast.Ident {
12 . . NamePos: sample.go:1:3
13 . . Name: "y"
14 . . Obj: *(obj @ 4)
15 . }
16 }
==== AFTER ====
0 *ast.BinaryExpr {
1 . X: *ast.BasicLit {
2 . . ValuePos: -
3 . . Kind: INT
4 . . Value: "10"
5 . }
6 . OpPos: sample.go:1:2
7 . Op: +
8 . Y: *ast.BasicLit {
9 . . ValuePos: -
10 . . Kind: INT
11 . . Value: "20"
12 . }
13 }

ここではast.Print関数を使ってASTの構造を出力しています。

うまくオペランドが入れ替わっていることが分かるでしょう。

さて、次に変更したASTから文字列として式を得たいと思います。


フォーマットをかける

go/formatパッケージには、以下の2つの関数が用意されています。

func Node(dst io.Writer, fset *token.FileSet, node interface{}) error

func Source(src []byte) ([]byte, error)

format.Node関数は、ASTのノードに対してフォーマットをかけ、format.Source関数は、文字列で表されたコードに対してフォーマットを掛けます。

ここでは、ASTのノードを対象としているのでformat.Node関数を用います。

format.Node関数の第1引数は、出力先のio.Writeです。文字列としてコードが欲しい場合は、bytes.Buffer型などを渡すと良いでしょう。

第2引数のfsetは、パースする際に渡したtoken.FileSet型の値です。詳しくは、「ASTを取得する方法を調べる」という記事に書いてますので、そちらを御覧ください。

第3引数は、整形するASTのノードです。

さて、前置きが長くなりましたが、ast.Node関数を呼び出して、ソースコードを整形してみましょう。

先程オペランドを置き換えた、2項演算式のASTを整形してみます。

if err := format.Node(os.Stdout, fset, expr); err != nil {

log.Fatalln("Error:", err)
}

出力結果は以下のようになります。

10 + 20

うまく、xy1020に入れ変わってますね。

また、ついでに+の間にスペースが入っているので整形されていることが分かります。

最後にすべてのコードを載せておきます。The Go Playgroundでも実行可能です。

package main

import (
"fmt"
"go/ast"
"go/format"
"go/parser"
"go/token"
"log"
"os"
)

func main() {
fset := token.NewFileSet()
expr, err := parser.ParseExprFrom(fset, "sample.go", `x+y`, 0)
if err != nil {
log.Fatalln("Error:", err)
}

fmt.Println("==== BEFORE ====")
ast.Print(fset, expr)

binaryExpr := expr.(*ast.BinaryExpr)

binaryExpr.X = &ast.BasicLit{
Kind: token.INT,
Value: "10",
}

binaryExpr.Y = &ast.BasicLit{
Kind: token.INT,
Value: "20",
}

fmt.Println("==== AFTER ====")
ast.Print(fset, expr)

if err := format.Node(os.Stdout, fset, expr); err != nil {
log.Fatalln("Error:", err)
}
}


おわりに

この記事では、ASTをいじる方法とその後go/formatパッケージをつかってソースコードを整形する方法について説明しました。

ソースコード中の宣言の順番をアルファベット順に並べかえたりするツールなどを作る際にきっと使えると思いますので、ぜひやってみてください。