GoのastutilでAST処理をしてみる

前回の記事: [Goの構文解析に入門してみる #Go - Qiita](https://qiita.com/ssc-ynakamura/items/142ce6b729299d82d039)

少し前にGoの構文解析に入門してみたわけですが、前回の記事の最後では[`golang.org/x/tools/go/ast/astutil`](https://golang.org/x/tools/go/ast/astutil)（以下、`astutil`）という便利パッケージがあることに触れました。
今回はこの`astutil`パッケージを使ってみたいと思います。

標準の`ast`での処理を踏まえた上での`astutil`なので、前回の記事もご一読いただけると、より一層理解しやすくなるかと思います。
「AST（抽象構文木）って何？」という方も、前回の記事をご一読ください。

## 本記事の動作環境

本記事は次の環境で動作確認しながら書きました。

- macOS Ventura
- Go 1.22

OSについては特に環境依存はないかと思います。

Goのバージョンについては、大きく離れていなければ（ここ1、2年ぐらいのバージョンであれば）、それほど差異はないかと思います。
（ジェネリクスのような構文的に大きな違いがある場合は、その限りではないです）

## 今回やってみること

今回扱う内容は次の2つです。いずれも`astutil`を使って行います。

- AST（抽象構文木）の走査
- ノードの削除

`astutil`は、標準の`ast.Inspect()`よりも痒いところに手が届く、`astutil.Apply()`という走査関数が特徴となっています。今回はそれを体験してみたいと思います。

## AST（抽象構文木）の走査

今回のコードです。

今回の対象コードはファイルではなく`Greet("hello", "world")`という簡単な関数呼び出しの式にしました。
それ以外の全体の流れは、前回の`ast.Inspect()`を使ったときと大体同じです。

```go
package main

import (
	"fmt"
	"go/parser"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
)

func main() {
	if err := traverse(); err != nil {
		panic(err)
	}
}

func traverse() error {
	// 今回対象とするコード（単純な関数呼び出しの式）
	code := `Greet("hello", "world")`

	f, err := parser.ParseExpr(code)
	if err != nil {
		return err
	}

	// ノードを走査して処理する
	level := 0
	astutil.Apply(f, func(cursor *astutil.Cursor) bool {
		indent := strings.Repeat("  ", level)
		level++

		// ノードの情報を表示
		node := cursor.Node()
		fmt.Printf("%sPRE: %T(%s[%d]): %+v\n", indent, node, cursor.Name(), cursor.Index(), node)
		return true
	}, func(cursor *astutil.Cursor) bool {
		level--
		indent := strings.Repeat("  ", level)

		// ノードの情報を表示
		node := cursor.Node()
		fmt.Printf("%sPOST: %T(%s[%d]): %+v\n", indent, node, cursor.Name(), cursor.Index(), node)
		return true
	})

	return nil
}
```

コードの解説はこのあとしますが、まずはこれを実行すると次のような出力を得られます。

```
PRE: *ast.CallExpr(Node[-1]): &{Fun:Greet Lparen:6 Args:[0xc0000c0020 0xc0000c0040] Ellipsis:0 Rparen:23}
  PRE: *ast.Ident(Fun[-1]): Greet
  POST: *ast.Ident(Fun[-1]): Greet
  PRE: *ast.BasicLit(Args[0]): &{ValuePos:7 Kind:STRING Value:"hello"}
  POST: *ast.BasicLit(Args[0]): &{ValuePos:7 Kind:STRING Value:"hello"}
  PRE: *ast.BasicLit(Args[1]): &{ValuePos:16 Kind:STRING Value:"world"}
  POST: *ast.BasicLit(Args[1]): &{ValuePos:16 Kind:STRING Value:"world"}
POST: *ast.CallExpr(Node[-1]): &{Fun:Greet Lparen:6 Args:[0xc0000c0020 0xc0000c0040] Ellipsis:0 Rparen:23}
```

ノードを順に辿って、ノードの情報が順番に出力されているのですが、それぞれ`PRE`と`POST`で2回出力されているのが見て取れるかと思います。
また、子ノードがある場合は子ノードを辿った後に`POST`が呼ばれているのがわかるかと思います（上記例ではインデントで表現しています）。

### コードの解説

#### [`astutil.Apply()`](https://pkg.go.dev/golang.org/x/tools/go/ast/astutil#Apply)の説明

今回のキモである`astutil.Apply()`は、次のようなシグネチャを持った関数です（引数の`ApplyFunc`についても示します）。

```go
func Apply(root ast.Node, pre, post ApplyFunc) (result ast.Node)

type ApplyFunc func(*Cursor) bool
```

動作としては、`root`に指定したノードを再帰的に走査し、各ノードに対して`pre`と`post`に指定した関数を呼び出します。
戻り値は変更された可能性のあるツリーになります。

##### `pre`のコールバック

`pre`が`nil`でない場合は、対象の子ノードが走査される前に呼び出されます。
`pre`が`false`を返す場合は、子ノードは走査されず、対象ノードに対する`post`も呼ばれません。

##### `post`のコールバック

`post`が`nil`でない場合は、対象ノードの子ノードが走査された後に呼び出されます（`pre`が`false`を返した場合は呼ばれない）。
`post`が`false`を返すと、残りの走査を打ち切って`Apply()`が終了します。

##### コールバックに渡ってくる引数

`pre`と`post`に渡ってくる引数は`ast.Node`ではなく、[`*astutil.Cursor`](https://pkg.go.dev/golang.org/x/tools/go/ast/astutil#Cursor)という、ノードをラップした型になっています。
`ast.Node`を取得できるだけでなく、ノードに対する追加の情報を取得できたり、ノードの削除や置き換えといった操作も行えるものになっています（色々あるので[マニュアル](https://pkg.go.dev/golang.org/x/tools/go/ast/astutil#Cursor)を参照してください）。

##### 戻り値

戻り値は基本的には`root`と同じものが返ってきますが、例えば`root`ノード自身を置き換えるような操作をすると、新しいノードが返ってくることになります。[^1]

[^1]: `append()`で新しいスライスが返ってくる可能性があるのと似ていますね。

#### 今回の例の解説

`astutil.Apply()`の部分を再掲します。

```go
	level := 0
	astutil.Apply(f, func(cursor *astutil.Cursor) bool {
		indent := strings.Repeat("  ", level)
		level++

		// ノードの情報を表示
		node := cursor.Node()
		fmt.Printf("%sPRE: %T(%s[%d]): %+v\n", indent, node, cursor.Name(), cursor.Index(), node)
		return true
	}, func(cursor *astutil.Cursor) bool {
		level--
		indent := strings.Repeat("  ", level)

		// ノードの情報を表示
		node := cursor.Node()
		fmt.Printf("%sPOST: %T(%s[%d]): %+v\n", indent, node, cursor.Name(), cursor.Index(), node)
		return true
	})
 ```

`pre`と`post`に渡す関数は本質的には同じで、ノードの情報を表示しているだけです。
違いは、先頭に`PRE: `や`POST: `をつけるのと、ノードの階層がわかるようにインデントレベルを調整しているだけです。

`cursor.Node()`で注目しているノードの`ast.Node`を得ています。
今回は使っていませんが、`cursor.Parent()`で親ノードを参照することもできます。
目的のノードを見つけたら親に対して何か処理をしたい、といった場合に有用です。

`cursor.Name()`で、このノードが「親から何と言う名前で認識されているか」を得ることができます。

今回の例では面白みがありませんが、例えば`1+2`のような二項演算の場合は、次のような感じになります。
それぞれのオペランドが`ast.BinaryExpr`から見て`X`と`Y`という名前で認識されていることがわかります。

```
PRE: *ast.BinaryExpr(Node[-1]): &{X:0xc000076020 OpPos:2 Op:+ Y:0xc000076040}
  PRE: *ast.BasicLit(X[-1]): &{ValuePos:1 Kind:INT Value:1}
  POST: *ast.BasicLit(X[-1]): &{ValuePos:1 Kind:INT Value:1}
  PRE: *ast.BasicLit(Y[-1]): &{ValuePos:3 Kind:INT Value:2}
  POST: *ast.BasicLit(Y[-1]): &{ValuePos:3 Kind:INT Value:2}
POST: *ast.BinaryExpr(Node[-1]): &{X:0xc000076020 OpPos:2 Op:+ Y:0xc000076040}
```

`cursor.Index()`は、注目ノードが親からスライスで管理されている場合に、そのインデックスを得られます。スライスで管理されていない場合は、0より小さい値が返ってきます。

今回の例では関数の引数の部分がこれに当たります。

- `*ast.BasicLit(Args[0])`
- `*ast.BasicLit(Args[1])`

`Args`という名前のスライスで管理され、それぞれ`0`と`1`というインデックスであることがわかりますね。
また、その他のノードは`-1`になっているので、スライスではないということになります。

## ノードの削除

次はノードを書き換えてみましょう。
次のサンプルコードは`Greet("hello", "world")`の第二引数`"world"`を削除する例です。

```go
package main

import (
	"fmt"
	"go/format"
	"go/parser"
	"go/token"
	"os"

	"golang.org/x/tools/go/ast/astutil"
)

func main() {
	if err := modifyNode(); err != nil {
		panic(err)
	}
}

func modifyNode() error {
	// 今回対象とするコード（単純な関数呼び出しの式）
	code := `Greet("hello", "world")`

	f, err := parser.ParseExpr(code)
	if err != nil {
		return err
	}

	// ノードを走査して、対象のノードがあったら削除する
	modifiedNode := astutil.Apply(f, func(cursor *astutil.Cursor) bool {
		if cursor.Name() == "Args" && cursor.Index() == 1 {
			cursor.Delete()
		}

		return true
	}, nil)

	// 書き換え後のASTをフォーマットして表示
	fset := token.NewFileSet()
	format.Node(os.Stdout, fset, modifiedNode)
	fmt.Println()

	return nil
}
```

実行すると次のような出力を得られます。
第二引数の`"world"`がなくなっていますね。

```
Greet("hello")
```

### コードの解説

今回は`astutil.Apply()`の`pre`で処理をしています。`post`は使っていません。

```go
	modifiedNode := astutil.Apply(f, func(cursor *astutil.Cursor) bool {
		if cursor.Name() == "Args" && cursor.Index() == 1 {
			cursor.Delete()
		}
		return true
	}, nil)
```

注目ノードの名前が`Args`で、インデックスが`1`の場合（つまり引数の2つ目）に、`cursor.Delete()`でノードを削除しています。
また、`astutil.Apply()`の戻り値で変更後のASTを受け取って、後で表示しています。

ここで`cursor.Replace()`を使うと、別のノードに置き換えることもできます。
例えば次のようにすると、`"world"`引数を`"gopher"`引数に置き換えて、`Greet("hello", "gopher")`にすることができます。

```go
cursor.Replace(&ast.BasicLit{
    Kind:  token.STRING,
    Value: `"gopher"`,
})
```

このような処理を標準ライブラリの`ast.Inspect()`で行おうとすると、親ノードを別途覚えておくなどの工夫が必要になり、簡単には実装できませんでした。
`astutil`パッケージを使うとこれらの操作が簡単に行えるのでとてもいいですね。

## （おまけ）import周りの便利機能

`astutil`にはその他の機能としてimport周りの処理を簡単に行える関数もあるので、簡単に紹介しておきます。

- `astutil.Imports()`: import宣言のリストを取得
- `astutil.AddImport()`: import宣言の追加
- `astutil.AddNamedImport()`: 名前付きimport宣言の追加
- `astutil.DeleteImport()`: import宣言の削除
- `astutil.DeleteNamedImport()`: 名前付きimport宣言の削除
- `astutil.RewriteImport()`: import宣言の置き換え
- `astutil.UsesImport()`: import宣言されているパッケージが使われているか判別

コードを修正するツールを書くときに、import周りを簡単に処理できそうですね。

## おわりに

前回に続いて、今度は`astutil`パッケージを使ってAST処理を行ってみました。
標準パッケージだけでは難しそうだった処理も簡単にできて、便利そうなパッケージだと感じました。
何かツールを作るときには是非とも活用したいですね。