TypeScriptのgeneratorで再帰的探索をシンプルに実装する

こんにちは、エンジニアの籏野です。

最近、ESLintカスタムルールの開発において、AST（抽象構文木）から特定のノードをすべて抽出してルールを作成したいという要望がありました。
その際に、generatorを活用することで、非常に効率的かつ可読性の高い実装を行うことができたので、その方法を紹介したいと思います。

generatorの基本については既に多くの記事が存在するため、本記事では実際のコード例を通じて、generatorを用いてどのように探索を行うのかを中心に解説します。

最終形

以下のようなコードで、ASTを再帰的に探索し特定のノードを抽出することができます。

const extractor = new ExtractVariableNode();
const allNodeList = Array.from(
	extractor.extract(attr.value),
);

では、ExtractVariableNodeクラスの具体的な実装について紹介していきます。

ExtractVariableNodeクラスの実装

実際のコードは以下のリポジトリにありますので、興味のある方はぜひご覧ください。

https://github.com/taku-hatano/generator_trial

まず、実際に呼び出されるextractメソッドを見ていきます。

	public *extract(
		node: TSESTree.Node,
	): Generator<TSESTree.Identifier, void, unknown> {
		switch (node.type) {
			case AST_NODE_TYPES.Identifier:
				yield node;
				break;
			case AST_NODE_TYPES.JSXExpressionContainer:
				yield* this.jsxExpressionContainer(node);
				break;
			case AST_NODE_TYPES.LogicalExpression:
				yield* this.logicalExpression(node);
				break;
			...
			default:
				console.warn(`Unsupported node type: ${node.type}`);
				break;
		}
	}

上記の通り、extractメソッドの実装は巨大なswitch文になっています。
switch文自体はノードの種類の分だけ巨大になっていきますが、やっていることはノードの種類に応じたハンドラーに値を渡しているだけなので非常にシンプルな実装となっています。
ここでは抽出したいノードの場合だけ、そのノードをyieldで返している点がポイントとなります。
これが再帰的探索の最終目標地点となり、抽出したいノードを返している部分となります。

続いて、各ノードタイプに対するハンドラーの実装を見ていきます。

	private *logicalExpression(node: TSESTree.LogicalExpression) {
		yield* this.extract(node.left);
		yield* this.extract(node.right);
	}

ここで紹介するLogicalExpressionとは、a && bのような論理演算を表すノードです。
ここではノードを左辺と右辺に分割し、それぞれ再度extractメソッドに投げることで再帰的な探索を行うようにしています。
extractメソッドに与えられた値が目的のノードであれば、先に紹介した通りそのノードがyieldされて値が返り、そうでなければ再帰処理を続けてより深く探索を続けることができます。

他のメソッドについても構成は同じで、自身が持つ子ノードの中で必要なものをextractメソッドに投げるという非常にシンプルな形での実装ができました。

これらの処理でyieldされたノードは最終的にArray.fromで配列に集約されて後続の処理で利用することが可能となります。

generatorを使う利点

generatorを使わずに同様の再帰的探索を実装する場合、例えば以下のような配列を使った実装が考えられます。

// generatorを使わない場合の例
private extractNodes(node: TSESTree.Node, result: TSESTree.Identifier[] = []): TSESTree.Identifier[] {
    switch (node.type) {
        case AST_NODE_TYPES.Identifier:
            result.push(node);
            break;
        case AST_NODE_TYPES.LogicalExpression:
            this.extractNodes(node.left, result);
            this.extractNodes(node.right, result);
            break;
        // ...
    }
    return result;
}

この実装と比較すると、generatorを使った場合には以下のような利点があります。

配列管理が不要でシンプル

generatorを使わない場合、各再帰呼び出しで結果を格納する配列を引き回す必要があります。
一方、generatorを使うことで、この配列管理が完全に不要になり、各メソッドは単純にyieldで値を返すだけで済みます。
これにより、コードがより簡潔で理解しやすくなります。

バグを産みづらい

配列を引き回す実装では、配列の参照を誤って変更したり、再帰呼び出し時に配列を正しく渡し忘れるといったバグが発生する可能性があります。
generatorを使った実装では、このような配列に関する操作が一切不要になるため、バグの発生源を大幅に減らすことができます。

遅延評価による効率性

generatorは必要に応じて値を生成する遅延評価の仕組みを持っています。
そのため、すべてのノードを探索する必要がない場合（例：最初の数個だけ取得したい場合）には、不要な探索処理をスキップできる可能性があります。

まとめ

TypeScriptのgeneratorを活用することで、再帰的探索処理を非常にシンプルにかけることがわかりました。
generatorはなんとなく難しいイメージがあったのですが、今回の実装を通して少しは仲良くなれたような気がします。
今後も適切な場面でgeneratorを活用していきたいと思います。

この記事を書いた人

籏野 拓
2018年新卒入社