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

はじめに

Goには、gofmtというコマンドがあることはご存知かと思います。
そして、そのフォーマットがGoコミュニティの標準となっています。
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構造体は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としておいて構いません。なお、値が0のtoken.Posはtoken.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

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

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