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

## はじめに

Goには、`gofmt`というコマンドがあることはご存知かと思います。
そして、そのフォーマットがGoコミュニティの標準となっています。
[`go/format`](https://golang.org/pkg/go/format/)パッケージでは、`gofmt`と同じスタイルのフォーマットでソースコードを整形することができます。
また、`go/format`パッケージを使うと、ASTについても整形することがきます。
この記事では、ASTのノードをいじって、整形する方法について説明します。

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

## ASTをいじる

ASTをいじると書きましたが、具体的に何をするんでしょうか？
ここでは、AST上のノードを別のノードに入れ替えたり、ノードのフィールドを差し替えることを指しています。

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

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

```
*ast.BinaryExpr (+)
├── *ast.Ident (x)
└── *ast.Ident (y)
```

[`ast.BinaryExpr`](https://golang.org/pkg/go/ast/#BinaryExpr)構造体は2項演算式を表し、オペランドは`X`フィールドと`Y`フィールドに[`ast.Expr`](https://golang.org/pkg/go/ast/#Expr)型の値として保持されます。

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

```
*ast.BinaryExpr (+)
├── *ast.BasicLit (10)
└── *ast.BasicLit (20)
```

[`ast.BasicLit`](https://golang.org/pkg/go/ast/#BasicLit)構造体は、以下のようなフィールドを持つ構造体です。

```go
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`としておいて構いません。なお、値が`0`の`token.Pos`は[`token.NoPos`](https://golang.org/pkg/go/token/#NoPos)という名前の定数として定義されています。

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

```go
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](https://golang.org/pkg/go/ast/#Print)関数を使ってASTの構造を出力しています。
うまくオペランドが入れ替わっていることが分かるでしょう。

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

## フォーマットをかける

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

```go
func Node(dst io.Writer, fset *token.FileSet, node interface{}) error
func Source(src []byte) ([]byte, error)
```

[`format.Node`](https://golang.org/pkg/go/format/#Node)関数は、ASTのノードに対してフォーマットをかけ、[`format.Source`](https://golang.org/pkg/go/format/#Source)関数は、文字列で表されたコードに対してフォーマットを掛けます。

ここでは、ASTのノードを対象としているので`format.Node`関数を用います。
`format.Node`関数の第1引数は、出力先の`io.Write`です。文字列としてコードが欲しい場合は、`bytes.Buffer`型などを渡すと良いでしょう。
第2引数の`fset`は、パースする際に渡した[`token.FileSet`](https://golang.org/pkg/go/token/#FileSet)型の値です。詳しくは、「[ASTを取得する方法を調べる](http://qiita.com/tenntenn/items/13340f2845316532b55a)」という記事に書いてますので、そちらを御覧ください。
第3引数は、整形するASTのノードです。

さて、前置きが長くなりましたが、`ast.Node`関数を呼び出して、ソースコードを整形してみましょう。
先程オペランドを置き換えた、2項演算式のASTを整形してみます。

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

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

```
10 + 20
```

うまく、`x`と`y`が`10`と`20`に入れ変わってますね。
また、ついでに`+`の間にスペースが入っているので整形されていることが分かります。

最後にすべてのコードを載せておきます。[The Go Playground](https://play.golang.org/p/Ay2t63ymwZ)でも実行可能です。

```go
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`パッケージをつかってソースコードを整形する方法について説明しました。
ソースコード中の宣言の順番をアルファベット順に並べかえたりするツールなどを作る際にきっと使えると思いますので、ぜひやってみてください。