AST の基礎を理解：TypeScript でコードを解析し操作する🔰 - 弁護士ドットコム株式会社 Creators’ blog

この記事は、弁護士ドットコム Advent Calendar 2024 の 16 日目の記事です。
クラウドサイン事業本部のフロントエンドエンジニアの辻@t0daaayです。
最近カンファレンスや勉強会で AST 関連の話を見聞きすることが多い中、実際に AST を触ったことがなかったため学びました。

本記事では、イチから TypeScript の AST の概要の理解をし、実際の AST の操作やそれを簡単に行うためのライブラリを使った内容をまとめています。

AST とは

AST（Abstract Syntax Tree : 抽象構文木）はソースコードをパース（解析）し、木構造で表現したものです。
AST 化することで、リンター、トランスパイラ、コンパイラなどのツールによるコード構造の理解、エラーチェック、コード変換などを効率的に実行可能になります。

TypeScript AST Viewer

ブラウザで AST を確認するためのツールとして TypeScript AST Viewer があるので、これを使って AST を確認してみます。

TypeScript AST Viewer

（1）ソースコード

ここにパースしたいソースコードを入力します。
この例では console.log("helloWorld"); のシンプルなコードをパースしています。

（2）AST

ソースコードがパースされて AST 化されたものが表示されます。

SourceFile
├── ExpressionStatement
│   └── CallExpression
│       ├── PropertyAccessExpression
│       │   ├── Identifier
│       │   └── Identifier
│       └── StringLiteral
└── EndOfFileToken

ExpressionStatement（式文）
CallExpression（関数呼び出し）
PropertyAccessExpression（プロパティアクセス）
Identifier（識別子）
StringLiteral（文字列リテラル）

と SourceFile 内に要素が細かく分解されて階層的に表現されていることが分かります。この各要素のことをノードと言います。
TypeScript の AST では SourceFile がトップレベルの親ノードとなり、 EndOfFileToken が末尾のノードとしてファイルの終端を表します。

（3）プロパティ

選択したノードの詳細情報が表示されます。

（4）ファクトリーコード

選択したノードの AST を生成するためのコードであるファクトリーコードが表示されます。
console.log("helloWorld"); のファクトリーコードの内容は以下のとおりです。

[
  factory.createExpressionStatement(
    factory.createCallExpression(
      factory.createPropertyAccessExpression(
        factory.createIdentifier("console"),
        factory.createIdentifier("log")
      ),
      undefined,
      [factory.createStringLiteral("helloWorld")]
    )
  ),
];

AST からソースコードを出力する

前段の TypeScript AST Viewer 上に生成されたファクトリーコードで AST を作成して、それをソースコードに変換してみます。
TypeScript パッケージに含まれる Compiler API を使用することで、 AST の作成やソースコードへの変換できます。

generateCode.ts を作成して、以下のコードを記述します。

1. モジュールのインポート

必要なモジュールをインポートします。 fs はファイルの書き出しに使います。

import * as ts from "typescript";
import * as fs from "fs";

2. AST を作成

TypeScript AST Viewer のファクトリーコードを statements（AST ノード）として宣言します。

const factory = ts.factory;

const statements = [
  // TypeScript AST Viewer のファクトリーコードを設定
  factory.createExpressionStatement(
    factory.createCallExpression(
      factory.createPropertyAccessExpression(
        factory.createIdentifier("console"),
        factory.createIdentifier("log")
      ),
      undefined,
      [factory.createStringLiteral("helloWorld")]
    )
  ),
];

3. ノードをソースファイルに配置

AST ノードを配置した、 sourceFile を作成します。
第二引数には末尾のノードを表す EndOfFileToken、第三引数には SourceFile に特別なフラグの設定が不要なことを表す ts.NodeFlags.None を設定します。

const sourceFile = factory.createSourceFile(
  statements,
  factory.createToken(ts.SyntaxKind.EndOfFileToken),
  ts.NodeFlags.None
);

4. プリンターで AST からコードを生成

プリンターを使って、AST からソースコード文字列を生成します。

const printer = ts.createPrinter();
const code = printer.printFile(sourceFile);

5. ファイルに書き込み

最後に、生成したコード文字列をファイルに書き出します。

fs.writeFileSync("./output.ts", code, { encoding: "utf8" });

コード全文

import * as ts from "typescript";
import * as fs from "fs";

const factory = ts.factory;

const statements = [
  factory.createExpressionStatement(
    factory.createCallExpression(
      factory.createPropertyAccessExpression(
        factory.createIdentifier("console"),
        factory.createIdentifier("log")
      ),
      undefined,
      [factory.createStringLiteral("helloWorld")]
    )
  ),
];

const sourceFile = factory.createSourceFile(
  statements,
  factory.createToken(ts.SyntaxKind.EndOfFileToken),
  ts.NodeFlags.None
);

const printer = ts.createPrinter();
const code = printer.printFile(sourceFile);

fs.writeFileSync("./output.ts", code, { encoding: "utf8" });

これを実行することで console.log("helloWorld"); が書かれた output.ts を出力できます。

AST でソースコードを変更する

ここでは、前段で生成したファイル output.ts を読み込み、AST を操作することでソースコードを変換する transformCode.ts のファイルを作成します。
AST を操作することにより前段で出力した output.ts の "helloWorld" の文字列を、"helloTypeScript" に変更します。

1. モジュールのインポート

import * as ts from "typescript";
import * as fs from "fs";

2. コードの読み込みと AST 化

output.ts ファイルに書かれたコードを読み込み、そのコードから AST を生成します。
ここでは、ファイル名、出力先、スクリプトターゲットなどの設定しています。

const sourceText = fs.readFileSync("./output.ts", "utf8");

const sourceFile = ts.createSourceFile(
  "output.ts",
  sourceText,
  ts.ScriptTarget.ESNext,
  true,
  ts.ScriptKind.TS
);

3. トランスフォーマーの作成

AST を変換するための関数（トランスフォーマー）を作成し、ノードが文字列リテラルで、かつテキストが oldName に一致する場合、newName に置き換えるように実装しています。
関数 visitNode で再起的にノードを辿って（トラバース）条件に一致したノードを変換しています。

const createRenameTransformer = (
  oldName: string,
  newName: string
): ts.TransformerFactory<ts.SourceFile> => {
  return (context) => (sourceFile) => {
    const visitNode = (node: ts.Node): ts.Node =>
      ts.isStringLiteral(node) && node.text === oldName
        ? ts.factory.createStringLiteral(newName)
        : ts.visitEachChild(node, visitNode, context);

    return ts.visitEachChild(sourceFile, visitNode, context);
  };
};

4. AST の変換実行

ts.transform によって AST 変換します。

const result = ts.transform(sourceFile, [
  createRenameTransformer("helloWorld", "helloTypeScript"),
]);

5. プリンターで AST からコードを生成

const printer = ts.createPrinter();
const transformedCode = printer.printFile(result.transformed[0]);

6. ファイルに書き込み

fs.writeFileSync("./output.ts", transformedCode, { encoding: "utf8" });

コード全文

import * as ts from "typescript";
import * as fs from "fs";

const sourceText = fs.readFileSync("./output.ts", "utf8");
const sourceFile = ts.createSourceFile(
  "output.ts",
  sourceText,
  ts.ScriptTarget.ESNext,
  true,
  ts.ScriptKind.TS
);

const createRenameTransformer = (
  oldName: string,
  newName: string
): ts.TransformerFactory<ts.SourceFile> => {
  return (context) => (sourceFile) => {
    const visitNode = (node: ts.Node): ts.Node =>
      ts.isStringLiteral(node) && node.text === oldName
        ? ts.factory.createStringLiteral(newName)
        : ts.visitEachChild(node, visitNode, context);

    return ts.visitEachChild(sourceFile, visitNode, context);
  };
};

const result = ts.transform(sourceFile, [
  createRenameTransformer("helloWorld", "helloTypeScript"),
]);

const printer = ts.createPrinter();
const code = printer.printFile(result.transformed[0]);

fs.writeFileSync("./output.ts", code, { encoding: "utf8" });

上記手順を実行すると、output.ts 内の console.log("helloWorld");が console.log("helloTypeScript"); に変換されます。
このように AST でソースコードを変更することで、プログラム的にリファクタリングが可能になります。
ただし、コードを見て分かるとおり AST のコードは低レベルな記述が多く、複雑かつ記述量が膨大になり書くことが難しいです。
次に説明する ts-morph のような TypeScript Compiler API をラップした高レベルなライブラリを使うことで、直感的かつ簡潔なコードで AST の操作ができます。

ts-morph を使った AST 操作

dsherret/ts-morph

前述したとおり、 ts-morph は TypeScript Compiler API を使いやすい形に抽象化したライブラリです。
AST ノードの生成・変更・出力といった処理を直感的かつ簡潔なコードで実現できます。
詳細な説明は省きますが、前段までのコードを ts-morph を使用すると以下のように記述できます。

generateCode.ts

import { Project } from "ts-morph";

const project = new Project();

const sourceFile = project.createSourceFile("output.ts", "", {
  overwrite: true,
});

sourceFile.addStatements(`console.log("helloWorld");`);

project.saveSync();

transformCode.ts

import { Project, SyntaxKind } from "ts-morph";

function createRenameTransformer(oldName: string, newName: string) {
  return (sourceFile) => {
    sourceFile
      .getDescendantsOfKind(SyntaxKind.StringLiteral)
      .filter((strNode) => strNode.getLiteralValue() === oldName)
      .forEach((strNode) => strNode.setLiteralValue(newName));
  };
}

const project = new Project();
const sourceFile = project.addSourceFileAtPath("output.ts");

const renameTransformer = createRenameTransformer(
  "helloWorld",
  "helloTypeScript"
);
renameTransformer(sourceFile);

project.saveSync();

ts-morph の詳細については、公式ドキュメントが用意されているので気になる人はチェックしてみてください。

まとめ

AST はソースコードを構文的な要素に分解・構造化して表現する手法であり、リンターやトランスパイラ、コンパイラなどで活用されている
TypeScript Compiler API を使うと、TypeScript で書かれたコードから AST を生成し、ソースコードへ出力したり、コード中の識別子・構造を変更することが可能
ただし、TypeScript Compiler API は低レベルな記述が多く、大量かつ複雑な記述が必要になる
ts-morph のような高レベルなライブラリを使うことで、より直感的・簡潔に AST を操作・変更・生成できる
AST を使うことで、大規模なコードリファクタリングや独自ルールによるコード変換が容易になり、開発効率とコード品質の向上につながる

AST とは

TypeScript AST Viewer

（1） ソースコード

（2）AST

（3） プロパティ

（4） ファクトリーコード

AST からソースコードを出力する

1. モジュールのインポート

2. AST を作成

3. ノードをソースファイルに配置

4. プリンターで AST からコードを生成

5. ファイルに書き込み

コード全文

AST でソースコードを変更する

1. モジュールのインポート

2. コードの読み込みと AST 化

3. トランスフォーマーの作成

4. AST の変換実行

5. プリンターで AST からコードを生成

6. ファイルに書き込み

コード全文

ts-morph を使った AST 操作

まとめ

（1）ソースコード

（3）プロパティ

（4）ファクトリーコード