抽象構文木（AST）をトラバースする #golang

## はじめに

筆者が以前書いた「[簡単な式の評価機を作ってみる](http://qiita.com/tenntenn/items/a312d2c5381e36cf4cd3)」や「[【実践goパッケージ】文字列から複素数型の値をパースする](http://qiita.com/tenntenn/items/02e9a211d791901312a3)」などの記事では、文字列で表された式から抽象構文木(AST)を取得し、ASTのノードの種類ごとに再帰的に処理する関数を書きました。

この記事では、[`go/ast`](https://golang.org/pkg/go/ast/)パッケージに用意されている、[`ast.Inspect`](https://golang.org/pkg/go/ast/#Inspect)関数や[`ast.Walk`](https://golang.org/pkg/go/ast/#Walk)関数を使って、ASTをトラバースする方法について解説します。

なお、ASTを取得する方法については、「[ASTを取得する方法を調べる](http://qiita.com/tenntenn/items/13340f2845316532b55a)」という記事を書いていますので、そちらを参考にしてください。

この記事で扱っているGoのバージョンは1.7.4です。

## `ast.Inspect`関数

`ast.Inspect`は、ASTを深さ優先でトラバースする関数で、以下のようなシグニチャを持ちます。

```
func Inspect(node Node, f func(Node) bool)
```

第1引数の`node`は、トラバースするASTのルートノードを表し、第2引数の`f`はトラバースに使用される関数です。
関数`f`は、引数に[`ast.Node`](https://golang.org/pkg/go/ast/#Node)、つまりAST上の任意のノードを取り、戻り値に`bool`型の値を返します。戻り値が`true`の場合は、子ノードに対して再帰的に関数`f`を適用していきます。なお、すべての`nil`ではない子ノードに対して、関数`f`を適用した後に、`f(nil)`が実行されます。

たとえば、以下のように`1+1`のASTに対して`ast.Inspect`関数を呼び出してみましょう（[The Go Playgroundで実行する](https://play.golang.org/p/sb0c2gNNBe)）。

```go
expr, err := parser.ParseExpr(`1+1`)
if err != nil {
	log.Fatalln("Error:", err)
}

ast.Inspect(expr, func(n ast.Node) bool {
	fmt.Printf("%[1]T %[1]v\n", n)
	return true
})
```

`1+1`のASTは以下のようになっており、[`ast.BinaryExpr`](https://golang.org/pkg/go/ast/#BinaryExpr)の子ノードとして2つの[`ast.BasicLit`](https://golang.org/pkg/go/ast/#BasicLit)が保持されています。

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

そのため、上記のコードを実行すると以下のような結果が表示されます。

```
*ast.BinaryExpr &{0x1050e150 2 + 0x1050e160}
*ast.BasicLit &{1 INT 1}
<nil> <nil>
*ast.BasicLit &{3 INT 1}
<nil> <nil>
<nil> <nil>
```

`ast.Inspect`関数は、まずルールノードの`*ast.BinaryExpr`型の値に対して、関数`f`を適用します。そしてその後、子ノードである2つの`*ast.BasicLit`型の値に対して、関数`f`を適用していきます。ここで関数`f`の適用は、深さ優先で行われるため、`*ast.BasicLit`型の子ノードについても関数`f`が適用されようとします。しかし、`*ast.BasicLit`型のノードには、子ノードが存在しないため、最後に`f(nil)`が実行され、`*ast.BinaryExpr`型のノードの処理へ戻ります。`*ast.BinaryExpr`型のノードのすべての子ノードについて関数`f`を適用したので、最後に`f(nil)`を実行し、処理が終了します。

少し分かりづらいかもしれないので、深さによって先頭に適度にスペースを入れる処理を入れてみます（[The Go Playgroundで実行する](https://play.golang.org/p/EcICnnbJ8w)）。

```go
expr, err := parser.ParseExpr(`1+1`)
if err != nil {
	log.Fatalln("Error:", err)
}

var i int
ast.Inspect(expr, func(n ast.Node) bool {
	fmt.Printf("%s%[2]T %[2]v\n", strings.Repeat("  ", i), n)
	if n != nil {
		i++
	} else {
		i--
	}
	return true
})
```

ここで`i`は、そのノードの深さを表しています。
関数`f`の引数である`n`が`nil`になる場合は、すべての子要素の処理が終わっているということなので、`i`の値を引いておき、それ以外は加算しています。

改めて実行すると、以下のような結果が得られます。

```
*ast.BinaryExpr &{0x1050e150 2 + 0x1050e160}
  *ast.BasicLit &{1 INT 1}
    <nil> <nil>
  *ast.BasicLit &{3 INT 1}
    <nil> <nil>
  <nil> <nil>
```

さきほどの実行結果より幾分分かりやすくなったんじゃないかと思います。

`ast.Inspect`関数は、ASTの中から特定のノードを探し出して処理する用途に用いられる事が多く、`ast.Walk`関数よりも簡易的に用いられます。実際、`ast.Inpsect`関数の中では、以下のように[`ast.Walk`関数が用いられています](https://github.com/golang/go/blob/release-branch.go1.7/src/go/ast/walk.go#L370-L386)。

```go
type inspector func(Node) bool

func (f inspector) Visit(node Node) Visitor {
	if f(node) {
		return f
	}
	return nil
}

// Inspect traverses an AST in depth-first order: It starts by calling
// f(node); node must not be nil. If f returns true, Inspect invokes f
// recursively for each of the non-nil children of node, followed by a
// call of f(nil).
//
func Inspect(node Node, f func(Node) bool) {
	Walk(inspector(f), node)
}
```

それでは次に`ast.Walk`関数について解説します。

## `ast.Walk`関数

`ast.Inspect`内で`ast.Walk`関数が使われていることから分かるように、子ノードに対する処理の仕方や深さ優先で探索する点は同じです。
一方で、`ast.Walk`関数は、以下のようなシグニチャをとっており、`ast.Inspect`関数より柔軟です。

```go
func Walk(v Visitor, node Node)
```

第2引数の`node`に関しては、`ast.Inspect`関数の第1引数の`node`と同じでトラバースするASTのルートノードを表します。
第1引数の[`ast.Visitor`](https://golang.org/pkg/go/ast/#Visitor)型はインタフェースであり、以下のように定義されています。

```go
type Visitor interface {
        Visit(node Node) (w Visitor)
}
```

`Visit`メソッドは、ノードに対する処理を記述したメソッドで、引数に`ast.Node`型の値を取り、戻り値に`ast.Visitor`を返します。`ast.Walk`は各子ノードに対して、この`Visit`メソッドを適用していき、戻り値の`w`が`nil`になった場合にそれより下の子ノードに対しての処理をやめます。

また、各子ノードについては、`Visit`メソッドが返す`ast.Visitor`型の値が持つ`Visit`メソッドを適用します。こうすることで、以下のようにノードの種類によって処理を分けることができます（[The Go Playgroundで実行する](https://play.golang.org/p/YwugB-Rxhc)）。

```go
package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"log"
)

type visitorFunc func(node ast.Node) (w ast.Visitor)

func (f visitorFunc) Visit(node ast.Node) (w ast.Visitor) {
	return f(node)
}

func main() {
	expr, err := parser.ParseExpr(`1+1`)
	if err != nil {
		log.Fatalln("Error:", err)
	}

	var v1, v2 visitorFunc

	v1 = visitorFunc(func(node ast.Node) (w ast.Visitor) {
		fmt.Printf("%T\n", node)
		switch node.(type) {
		case *ast.BinaryExpr:
			return v2
		}
		return nil
	})

	v2 = visitorFunc(func(node ast.Node) (w ast.Visitor) {
		if node == nil {
			return nil
		}
		lit := node.(*ast.BasicLit)
		fmt.Printf("BasicLit: %s\n", lit.Value)
		return nil
	})

	ast.Walk(v1, expr)
}
```

このように、`v2`は`*ast.BasicLit`の子ノードだけを処理することができます。なお、すべての子ノードの処理が終わった後に、`node`が`nil`で`Visit`メソッドが呼ばれるという点については注意しなければいけません。

なお、標準パッケージでは、途中で`Visitor`を切り替えるような処理は見受けられませんでした。そのため、ノードの種類によって切り替えるということはあまりしないのかもしれません。

## おわりに

この記事では、2種類の`AST`をトラバースする方法について解説しました。
`ast.Inspect`関数は簡易的にASTをトラバースでき、`ast.Walk`関数についてはより柔軟性のあるトラバースができます。
ぜひ、ASTを処理する際は、自前の再帰呼び出しだけではなく、これらの関数も使ってみてください。