astutil.Applyで抽象構文木を置き換える #golang

December 2, 2017

この記事はQiitaの記事をエクスポートしたものです。内容が古くなっている可能性があります。

はじめに

Goには、標準パッケージ以外にも準標準にあたるgolang.org/x/以下で管理されているパッケージが存在します。ここで管理されているパッケージは、Go本体と同様にGoチームによって管理され、コントリビュートの方法も概ね同じです。

この準標準パッケージには、さまざまなパッケージが存在します。これらは標準パッケージになるほど汎用的だったり、完成度は高くないはないですが、それでも十分に有用です。golang.org/x/tools/go/ast/astutilパッケージ（以降、astutilパッケージと表記）もそのひとつです。

astutilパッケージは静的解析で使用すると便利なパッケージで、抽象構文木（Abstract Syntax Tree, AST）に関連する便利な機能を提供します。静的解析や抽象構文木については、”goパッケージで簡単に静的解析して世界を広げよう“という記事を書いていますので、そちらをご覧ください。

さて、この記事ではastutilパッケージに最近（2017年12月現在）導入された(tools#77811)、抽象構文木のノードを操作するためのApply関数について解説します。

ソースコードの変更と抽象構文木

Goの開発ツールを作成している際に、与えられたソースコードの一部を変更したい場合があります。 ぱっと簡単に思いつく方法は、文字列として置換することですが、複雑な変更の場合には正規表現を駆使して黒魔術のようなコードを書く必要があります。 そういった黒魔術は、聖なる夜に向けて書くアドベントカレンダーには不向きです。また、Goに静的解析を行うためのパッケージであるgoパッケージが標準として存在するため、それを用いない手はありません。 goパッケージを用いることで、ソースコードを構造化された抽象構文木として扱うこともできます。 複雑な変更の場合は、構造化されていた方が扱いやすいですし、クリスマスには黒魔術よりツリーの方が都合が良いでしょう。 そこでここでは、抽象構文木で用いることを考えてみます。

たとえば、x+yという式があった場合に、これを10+20にしたいとしましょう。 x+yは次のような抽象構文木としてパースすることができます。

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

これを10+20に変更したい場合は、”抽象構文木（AST）をいじってフォーマットをかける“という記事でも解説しているように、次のような抽象構文木に変換してあげればよいでしょう。

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

先の記事でも紹介しているように、この場合は単純に抽象構文木のノードを入れ替えることでも実現が可能です。 しかし、特定の条件にマッチするノードをツリー上から探し出し、そのノードを別のノードに入れ替える場合、この方法では難しくなってきます。 難しくなる理由を抽象構文木の探索について解説しながら見ていくことにしましょう。

抽象構文木の探索

抽象構文木をその名の通り、ツリー構造をとっています。 ルートとなるノードが1つあり、そこから子ノードへと枝分かれしていきます。 抽象構文木のノードは、前述した2項演算式を表す*ast.BinaryExprやリテラルを表す*ast.BasicLitなどがあります。 これらの型はastパッケージで定義されており、組み合わせてツリー構造を作ることで抽象構文木を形成します。

さて、ツリー構造のデータを探索する場合、通常は再帰的な関数をルートノードから葉ノードへと、親ノードから子ノードへと適用していくことで探索することが多いでしょう。 astパッケージにも抽象構文木を探索する方法がいくつか用意されており、”抽象構文木（AST）をトラバースする“という記事で解説したast.Inspect関数もそのひとつです。

ast.Inspect関数は引数にルートノードと各ノードに適用する関数を渡すことで、深さ優先でノードを探索し、渡した関数を適用させることができます。 たとえば、ast.Inspect関数は次のように実行することができます。

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引数にルートノード、第2引数に各ノードに適用する関数が渡されていることが分かるでしょう。 このように、非常に簡単に抽象構文木を探索することができます。

また、ast.Inspect関数内で仕様しているast.Walk関数を使うとノードの適用する処理をast.Visitorインタフェースを用いて抽象化することできます。 そして、ast.VisitorインタフェースのVisitメソッドは、次のようにast.Visitorを返すため、ノードごとに適用する処理を変更していくことができます。

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

さて、ast.Inspect関数やast.Walk関数を用いることで、ある程度の処理は行うことができそうです。 特に、抽象構文木の中を探索し、特定の条件にマッチするノードを取得するという処理であれば、これらの関数で十分行うことができるでしょう。

しかし、ここでは抽象構文木のノードを変更する必要があります。 ノードの変更を伴う場合、ast.Inspect関数やast.Walk関数では不十分な場合があります。 たとえば、とあるノードが特定の条件を満たす場合に、親のノードに変更をいれたいとします。 しかし、ast.Nodeインタフェースは次のように定義されているため、親のノードを取得することができません。 また、ast.Nodeインタフェースを実装した、各具象的な型についても子ノードは取ることができても、親のノードは取得することはできません。

type Node interface {
        Pos() token.Pos // position of first character belonging to the node
        End() token.Pos // position of first character immediately after the node
}

そのため、ast.Inspect関数やast.Walk関数を用いて、親ノードの変更を行うためには、自身で親ノードを記録しておく必要がありました。 ast.Inspect関数やast.Walk関数に渡す処理の引数は決まっていますから、メソッドにしてフィールドに状態を保存するなどして、親のノードを記録していく必要があります。 煩雑な場合は、astパッケージの関数では事足りず、自身で再帰的な処理を実装する必要がありました。

この問題は静的解析を行う場合に頻繁に発生する共通のものでした。 そのため、go/ast: provide AST Rewrite functionality (ast.Walk is not good enough)というissueが立てられ、ast.Walk関数に変わる機能が必要という議論が始まりました。

2016年の9月に作られたissueでしたが、1年以上たってようやく実装が完成し、ついにgolang.org/x/tools/go/ast/astutilパッケージにApply関数として入ることになりました。

astutil.Apply関数でノードを操作する

さて、ここではastutil.Apply関数によって抽象構文木を探索し操作する方法について説明しましょう。 astutil.Apply関数は次のようなシグニチャを持ちます。

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

引数は探索を行うルートノードと２つの各ノードに適用するための関数です。 そして戻り値に変更後の抽象構文木を返します。 戻り値にresultという名前で結果を返すようにはなっていますが、astutil.Apply関数は副作用のある関数で抽象構文木がコピーされてノードの変更が行われる訳ではありません。 そのため、ポインタで保持されている抽象構文木のルートノードには当然ながら変更が加えられます。 しかし、ルートノード自体を置き換えた場合などは、新しく作られたルートノードが必要となるので、必ず受け取るようにしておきましょう。

引数にpreとpostという２つの関数を取る理由ですが、preは子ノードに処理を適用する前に、そのノードに適用する関数で、一方、postは子ノードの後に処理が適用される関数です。 ２つある理由は、前述のように破壊的な変更を抽象構文木に加える必要があるため、子ノードの変更を反映してから処理したい場合かそうでないかでpreとpostを使い分ける必要があるためです。

preとpostの型はastutil.ApplyFunc型となっています。 そして、astutil.ApplyFunc型は次のように定義されています。

type ApplyFunc func(*Cursor) bool

これは*astutil.Cursor型の値を引数に取り、bool値を返すような関数です。 詳細は後述しますが、*astutil.Cursorは現在注目しているノードを表す型です。 また、戻り値のbool値がfalseの場合は、astutil.Apply関数はすぐに処理を停止します。

さて、*astutil.Cursorについて見てみましょう。 astutil.Cursor型は構造体ですが、フィールドは公開されていません。 そのため、通常はメソッドのみを使うことを想定して設計されているようです。 *astutil.Cursor型は次の8つのメソッドを持っています。

  • func (c *Cursor) Node() ast.Nodeメソッド
  • func (c *Cursor) Parent() ast.Nodeメソッド
  • func (c *Cursor) Name() stringメソッド
  • func (c *Cursor) Index() intメソッド
  • func (c *Cursor) Replace(n ast.Node)メソッド
  • func (c *Cursor) InsertAfter(n ast.Node)メソッド
  • func (c *Cursor) InsertBefore(n ast.Node)メソッド
  • func (c *Cursor) Delete()メソッド

Nodeメソッド、Indexメソッド、Parentメソッド、Nameメソッドは、現在注目しているノードの情報を取得します。

Nodeメソッドはその名の通り、現在注目しているノードを、Parentメソッドはその親ノードを取得します。

Nameメソッドは、注目しているノードが親ノードにどういう役目で保持されているかを表す文字列を取得します。 たとえば、2項演算式を表す*ast.BinaryExpr型の場合、2つの子ノードを持っており、第1項がX、第2項がYという名前のフィールドで管理されています。 そのため、第1項にあたるノードが注目しているノードである場合に、NameメソッドはXという文字列を返します。

Indexメソッドは、親ノードが注目しているノードをスライスで管理している場合に、何番目のノードかを表す値が入っています。 たとえば、引数はスライスでなお、スライスで管理していない場合は0より小さい値が返ってきます。 スライスでノードが管理されていると、Nameメソッドが返す値だけでは一意にノードを特定できないため、Indexメソッドを使う必要がでてきます。 たとえば、何番目の引数であるか取得したい場合は次のように取得できます。

expr, err := parser.ParseExpr(`func(x, y int){}(10, 20)`)
if err != nil { log.Fatal(err) }
n := astutil.Apply(expr, func(cr *astutil.Cursor) bool {
	if cr.Name() == "Args" {
		fmt.Println(cr.Name(), cr.Index())
	}
	return true
}, nil)

実行すると次のように表示されます。

Args 0
Args 1

さて、*astutil.Cursor型の残りのReplaceメソッド、Deleteメソッド、InsertBeforeメソッド、InsertAfterメソッドは抽象構文木を操作するメソッドです。

Replaceメソッドはその名の通り、ノードを入れ替えるです。 次のように入れ替えたいノードを指定するとすぐにノードを入れ替えてくれます。

expr, err := parser.ParseExpr(`x+y`)
if err != nil { log.Fatal(err) }
n := astutil.Apply(expr, func(cr *astutil.Cursor) bool {
	switch cr.Name() {
	case "X":
		cr.Replace(&ast.BasicLit{
			Kind:  token.INT,
			Value: "10",
		})
	case "Y":
		cr.Replace(&ast.BasicLit{
			Kind:  token.INT,
			Value: "20",
		})
	}
	return true
}, nil)

if err := format.Node(os.Stdout, token.NewFileSet(), n); err != nil {
	log.Fatalln("Error:", err)
}
fmt.Println()

実行すると次のように表示されます。

10 + 20

Deleteメソッドは、注目しているノードを削除します。 次のように削除したいノードに対してDeleteメソッドを呼び出します。

expr, err := parser.ParseExpr(`func(x, y int){}(10, 20)`)
if err != nil { log.Fatal(err) }
n := astutil.Apply(expr, func(cr *astutil.Cursor) bool {
	if cr.Name() == "Args" && cr.Index() == 0 {
		cr.Delete()
	}
	return true
}, nil)

if err := format.Node(os.Stdout, token.NewFileSet(), n); err != nil {
	log.Fatalln("Error:", err)
}
fmt.Println()

実行すると次のように表示されます。

func(x, y int) {
}()

ここで注意したいのは、変更は直ちに行われるという点です。 そのため、引数の0番目を削除したい場合は、次のように1度だけ実行されるようにしたほうが良いでしょう。

expr, err := parser.ParseExpr(`func(x, y int){}(10, 20)`)
if err != nil {	log.Fatal(err) }

var once sync.Once
n := astutil.Apply(expr, func(cr *astutil.Cursor) bool {
	if cr.Name() == "Args" && cr.Index() == 0 {
		once.Do(cr.Delete)
	}
	return true
}, nil)

if err := format.Node(os.Stdout, token.NewFileSet(), n); err != nil {
	log.Fatalln("Error:", err)
}
fmt.Println()

実行すると次のように表示されます。

func(x, y int) {
}(20)

InsertBeforeメソッドとInsertAfterメソッドは、注意のノードがスライスで管理されている場合に、その前後にノードを挿入することができます。 たとえば、引数リストにあたらしい引数を追加したい場合は次のように行います。

expr, err := parser.ParseExpr(`func(x, y int){}(10, 20)`)
if err != nil { log.Fatal(err) }
n := astutil.Apply(expr, func(cr *astutil.Cursor) bool {
	if cr.Name() == "Args" && cr.Index() == 0 {
		cr.InsertBefore(&ast.BasicLit{
			Kind:  token.STRING,
			Value: "hi",
		})
		cr.InsertAfter(&ast.BasicLit{
			Kind:  token.STRING,
			Value: "gopher",
		})
	}
	return true
}, nil)

if err := format.Node(os.Stdout, token.NewFileSet(), n); err != nil {
	log.Fatalln("Error:", err)
}
fmt.Println()

実行すると次のように表示されます。

func(x, y int) {
}(hi, 10, gopher, 20)

このように*astutil.Cursor型のメソッドを使うと簡単に親ノードに関する情報を取得することができたり、抽象構文木の操作することができます。

おわりに

この記事では、最近（2017年12月2日現在）導入されたastutil.Apply関数について説明を行いました。 この関数は静的解析を行ってツールを使っている人たちには待望の機能だったのではないでしょうか。 ノードの操作が即時反映だったり、副作用を伴ったりをするため、まだまだ使い勝手としてはイマイチな点もありますが、ast.Walk関数よりも多くの事ができるため、活躍の場は多そうです。 実際に、@vvakameさんが開発しているgo.mercari.io/datastoreというライブラリでも早速使われているようです。 go.mercari.io/datastoreについては、@vvakameさんの”ぼくが かんがえた さいきょうの でーたすとあ らっぱー“という記事に詳しく書いてありますので、そちらも合わせてどうぞ。