前回の記事: Goの構文解析に入門してみる #Go - Qiita

少し前にGoの構文解析に入門してみたわけですが、前回の記事の最後では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()を使ったときと大体同じです。

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()`の説明

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

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という、ノードをラップした型になっています。
ast.Nodeを取得できるだけでなく、ノードに対する追加の情報を取得できたり、ノードの削除や置き換えといった操作も行えるものになっています（色々あるのでマニュアルを参照してください）。

戻り値

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

今回の例の解説

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

	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"を削除する例です。

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は使っていません。

	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")にすることができます。

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処理を行ってみました。
標準パッケージだけでは難しそうだった処理も簡単にできて、便利そうなパッケージだと感じました。
何かツールを作るときには是非とも活用したいですね。

append()で新しいスライスが返ってくる可能性があるのと似ていますね。 ↩