AST で JavaScript のコードを変換する

ソースコードを AST（抽象構文木）と呼ばれるデータ構造に変換することで、ソースコードの検証や変換をプログラムによって行えるようになる。
例えば ESLint では、ソースコードを AST に変換して、それに対してチェックを行っている。
また、V8 などの JavaScript エンジンも、対象の JavaScript ソースコードを AST に変換してから、後続の処理を行う。

AST を使えるようになると、ソースコードの検証や変換を行うツールを自作できるようになる。
この記事では、JavaScript AST を扱うための方法を見ていく。

この記事で使用しているライブラリのバージョンは以下の通り。

• esprima@4.0.1
• estraverse@5.2.0
• escodegen@2.0.0

動作環境は Node.js のv12.17.0で行っている。

また、この記事に出てくる木構造の画像は、以下のウェブアプリで作成した。

shape-painter.numb86.net

AST（抽象構文木）とは

ソースコードはただの文字列、テキストなので、そのままでは扱いづらい。
そこでまず、文字列を木構造に変換する。この変換のことを構文解析、またはパースと呼び、それによって得られた木構造を構文木と呼ぶ。
そして、コメントやスペース、改行など、プログラムの実行に不要な情報を取り去った構文木が、AST（抽象構文木）である。

ソースコードを AST に変換したら次は、その AST の中身を見て、コーディング規約に違反していないかどうかのチェックや、特定のルールに基づいた AST の書き換えなどを行う。
その作業を自力で行ってもよいが、AST を横断的に走査するためのツールが既にあるので、それを使うことが多い。
ソースコードの変換が目的の場合、最終的に AST をソースコードに戻す必要がある。

ソースコードを変換する場合、以下の流れになる。

1. ソースコードを AST にパースする（Code -> AST）
2. AST を走査し、任意のロジックに基づいて変換する（AST -> AST）
3. AST をソースコードに変換する（AST -> Code）

まず、ソースコードから AST へのパースについて見ていく。

Esprima

JavaScript ソースコードを AST に変換するライブラリは複数あるが、今回はEsprimaを使う。
Esprimaによって得られる AST は、ESTreeという仕様に準拠している。

github.com

parseScript と parseModule

parseScriptもしくはparseModuleに対象のソースコードを渡すことで、AST を得られる。

const esprima = require('esprima');

const ast = esprima.parseScript('x = 1');

console.log(ast);
// Script {
//   type: 'Program',
//   body: [
//     ExpressionStatement {
//       type: 'ExpressionStatement',
//       expression: [AssignmentExpression]
//     }
//   ],
//   sourceType: 'script'
// }

取得できた AST の詳細は後述する。

JavaScript として無効なコードを渡すとエラーになる。

const esprima = require('esprima');

// Error: Line 1: Unexpected end of input
const ast = esprima.parseScript('x =');

parseScriptとparseModuleの違いは、対象のソースコードをモジュールモードとして扱うか否か。
モジュールモードとして扱いたい場合は、parseModuleを使う必要がある。
そのため、以下のコードはエラーになる。import文はモジュールモードでしか使えず、無効なコードになっているため。

const esprima = require('esprima');

// Error: Line 1: Unexpected token
const ast = esprima.parseScript("import {parseScript} from 'esprima'");

parseModuleを使えば問題なく AST を取得できる。

const esprima = require('esprima');

const ast = esprima.parseModule("import {parseScript} from 'esprima'");

各種オプション

parseScriptやparseModuleの第二引数にはオプションを渡せるので、いくつか紹介しておく。

jsx

jsxを有効にすると、JSX もパースできるようになる。

const esprima = require('esprima');

esprima.parseScript('<div>a</div>', {jsx: true}); // ok
esprima.parseScript('<div>a</div>'); // Error: Line 1: Unexpected token <

range

rangeオプションを有効にすると、[開始位置, 終了位置]という形で位置情報を取得できる。

以下のコードは、x = 1は0番目から始まり、5番目が始まる前で終了していることを示している。
そしてx = 1に含まれる1は、4番目から始まって5番目が始まる前で終了している。

const esprima = require('esprima');

const ast = esprima.parseScript('x = 1', {range: true});

console.log(ast.range); // [ 0, 5 ]

const {right} = ast.body[0].expression;
console.log(right.value); // 1
console.log(right.range); // [ 4, 5 ]

loc

行番号や列番号を取得するには、locオプションを使用する。

以下のコードは、ソースコード全体は1行目で始まり、3行目の5列目が始まる前で終了していることを示している。

const esprima = require('esprima');

const sourceCode = `x = 1
2
y = 3`;

const ast = esprima.parseScript(sourceCode, {loc: true});

console.log(ast.loc); // { start: { line: 1, column: 0 }, end: { line: 3, column: 5 } }

comment

AST に変換する際、コメントは無視される。だがcommentオプションを有効にすると、コメントに関する情報を得られる。

const esprima = require('esprima');

const sourceCode = `// start code
const a = 1;
/*
comment area
*/
`;

const ast = esprima.parseScript(sourceCode, {comment: true});

console.log(ast.comments);
// [
//   { type: 'Line', value: ' start code' },
//   { type: 'Block', value: '\ncomment area\n' }
// ]

typeはLineとBlockの 2 種類で、それぞれ一行コメントと複数行コメントを表している。

AST はどのような構造になっているのか

次は、Esprima で取得できた AST、つまり ESTree の中身を見ていく。

ESTree は、ノードと呼ばれる要素が入れ子になって、木構造を構成している。
そして各ノードはオブジェクトとして表現され、typeプロパティを持っている。

ルート要素は必ずProgramになる。

const esprima = require('esprima');

const ast = esprima.parseScript('x = 1');

console.log(ast.type); // Program

そしてProgramのbodyに、子要素が入っている。

const esprima = require('esprima');

console.log(esprima.parseScript('x = 1').body.length); // 1
console.log(esprima.parseScript('x = 1; y = 2').body.length); // 2

x = 1だと式がひとつしかないので、bodyの要素はひとつになる。
x = 1; y = 2だと式が二つなので、bodyの要素も二つになる。

x = 1について確認してみると、ExpressionStatementというノードで、expressionプロパティを持っていることが分かる。

const esprima = require('esprima');

const ast = esprima.parseScript('x = 1');

const childElement = ast.body[0];

console.log(childElement);
// ExpressionStatement {
//   type: 'ExpressionStatement',
//   expression: AssignmentExpression {
//     type: 'AssignmentExpression',
//     operator: '=',
//     left: Identifier { type: 'Identifier', name: 'x' },
//     right: Literal { type: 'Literal', value: 1, raw: '1' }
//   }
// }

expressionにはAssignmentExpressionというノードが入っており、このノードはleftやrightというプロパティを持っている。
そしてleftとrightに、さらにノードが入っている。

図にすると、このような構成になっている。

ただのテキストであるソースコードをこのようなデータ構造に変換することで、プログラムで扱いやすくなる。

Estraverse

次は、この木構造を辿って、検証や変換を行う。

Estraverse というライブラリを使うと、ESTree のノードを簡単に探索していける。

github.com

Esprima で作った ESTree を、traverseメソッドの第一引数に渡す。
そして第二引数にオブジェクトを渡し、そこでenterメソッドを定義する。
enterメソッドの引数には各ノードが渡されるので、そこで任意の処理を行うことができる。

const esprima = require('esprima');
const estraverse = require('estraverse');

const ast = esprima.parseScript('x = 1');

estraverse.traverse(ast, {
  enter(node) {
    console.log(node.type);
  },
});
// Program
// ExpressionStatement
// AssignmentExpression
// Identifier
// Literal

木構造を、上から順番に辿っているのが分かる。
先程の画像に、処理する順番を追加したのが以下。

探索は、深さ優先探索で行う。
例えば、esprima.parseScript('foo(1) === bar(2)')は、以下の構造の AST を返す。

そしてこれをtraverseメソッドに渡すと、以下の順番で探索していく。

BinaryExpressionまで探索したあと、まずleftのCallExpressionを探索する。
それが終わってから、rightのCallExpressionを探索していく。

const esprima = require('esprima');
const estraverse = require('estraverse');

const ast = esprima.parseScript('foo(1) === bar(2)');

estraverse.traverse(ast, {
  enter(node) {
    if (node.type === 'Identifier') {
      console.log(node.type, node.name);
    } else if (node.type === 'Literal') {
      console.log(node.type, node.value);
    } else {
      console.log(node.type);
    }
  },
});
// Program
// ExpressionStatement
// BinaryExpression
// CallExpression
// Identifier foo
// Literal 1
// CallExpression
// Identifier bar
// Literal 2

enter と leave の違い

enterメソッドではなくleaveメソッドを定義することもできる。もちろん、両方のメソッドを定義してもよい。

estraverse.traverse(ast, {
  enter(node) {/* ... */},
  leave(node) {/* ... */},
});

両者とも出来ることは同じだが、呼ばれるタイミングが異なる。

深さ優先探索で探索していくことは既に述べたが、enterは、新しいノードに辿り着いたタイミングで実行される。
leaveはそれとは逆で、そのノードへ戻ってきたタイミングで、実行される。

const esprima = require('esprima');
const estraverse = require('estraverse');

const ast = esprima.parseScript('foo(1) === bar(2)');

estraverse.traverse(ast, {
  enter(node) {
    if (node.type === 'Identifier') {
      console.log('enter', node.type, node.name);
    } else if (node.type === 'Literal') {
      console.log('enter', node.type, node.value);
    } else {
      console.log('enter', node.type);
    }
  },
  leave(node) {
    if (node.type === 'Identifier') {
      console.log('leave', node.type, node.name);
    } else if (node.type === 'Literal') {
      console.log('leave', node.type, node.value);
    } else {
      console.log('leave', node.type);
    }
  },
});
// enter Program
// enter ExpressionStatement
// enter BinaryExpression
// enter CallExpression
// enter Identifier foo
// leave Identifier foo
// enter Literal 1
// leave Literal 1
// leave CallExpression
// enter CallExpression
// enter Identifier bar
// leave Identifier bar
// enter Literal 2
// leave Literal 2
// leave CallExpression
// leave BinaryExpression
// leave ExpressionStatement
// leave Program

見やすくするために CHALK を使って色を付けたのが以下。

これでも分かりづらいと思うので、木構造に、実行される順番を書いた。
各ノードの左側に書かれた数字はenterによる実行、右側はleaveによる実行。
これで、処理の流れがイメージできると思う。

skip と break

enterやleaveメソッドのなかでは、this.skipやthis.breakを呼び出すことができる。

まずはthis.skipについて。
this.skipを実行すると、その子要素への探索をスキップする。

const esprima = require('esprima');
const estraverse = require('estraverse');

const ast = esprima.parseScript('foo(1) === bar(2)');

estraverse.traverse(ast, {
  enter(node) {
    if (node.type === 'CallExpression' && node.callee.name === 'foo') {
      this.skip();
    }
    console.log('enter', node.type);
  },
  leave(node) {
    console.log('leave', node.type);
  },
});
// enter Program
// enter ExpressionStatement
// enter BinaryExpression
// enter CallExpression
// leave CallExpression
// enter CallExpression
// enter Identifier
// leave Identifier
// enter Literal
// leave Literal
// leave CallExpression
// leave BinaryExpression
// leave ExpressionStatement
// leave Program

この例の場合、左辺のCallExpressionの子要素への探索がスキップされる。

赤いノードのenterでthis.skipが呼び出されたため、黒いノードの処理がスキップされる。

this.breakを実行すると、現在実行中のenterやleaveの処理が終わり次第、traverseメソッドによる探索を終了する。

const esprima = require('esprima');
const estraverse = require('estraverse');

const ast = esprima.parseScript('foo(1) === bar(2)');

estraverse.traverse(ast, {
  enter(node) {
    console.log('enter', node.type);
  },
  leave(node) {
    if (node.type === 'CallExpression') {
      this.break();
    }
    console.log('leave', node.type);
  },
});
// enter Program
// enter ExpressionStatement
// enter BinaryExpression
// enter CallExpression
// enter Identifier
// leave Identifier
// enter Literal
// leave Literal
// leave CallExpression

この例では、左辺のCallExpressionのleaveを最後に、traverseメソッドが終了する。
実行中のenterメソッドやleaveメソッドは最後まで実行されるので、leave CallExpressionは表示される。

赤いノードのleaveでthis.breakが呼び出されたため、黒いノードへの探索は行われず、緑のノードのleaveも実行されない。

replace で AST を書き換える

渡された AST を書き換えたい場合は、traverseメソッドではなくreplaceメソッドを使う。
基本的な使い方はtraverseと一緒で、enterやleaveのなかで値を返すと、その値を元のノードと置き換える。

Literalの値を9に変えたい場合、以下のように書く。

const esprima = require('esprima');
const estraverse = require('estraverse');

const ast = esprima.parseScript('x = 1');

// Literal { type: 'Literal', value: 1, raw: '1' }
console.log(ast.body[0].expression.right);

const convertedAst = estraverse.replace(ast, {
  enter(node) {
    if (node.type === 'Literal') {
      return {
        ...node,
        value: 9,
        raw: '9',
      };
    }
    return node;
  },
});

// { type: 'Literal', value: 9, raw: '9' }
console.log(convertedAst.body[0].expression.right);

// { type: 'Literal', value: 9, raw: '9' }
console.log(ast.body[0].expression.right);

replaceは破壊的メソッドなので、convertedAstだけでなく、元の AST であるastも書き換えられているので注意する。
valueやrawが変わっている他、Literalのインスタンスではなくただのオブジェクトになっている。

対象のノードを削除したい場合は、enterやleaveのなかでthis.removeを呼び出す。

以下のコードでは、x = 1;を削除している。

const esprima = require('esprima');
const estraverse = require('estraverse');

const ast = esprima.parseScript('x = 1; y = 2;');

console.log(ast.body.length); // 2
console.log(ast.body[0].expression.left.name); // x
console.log(ast.body[1].expression.left.name); // y

estraverse.replace(ast, {
  enter(node) {
    if (node.type === 'ExpressionStatement') {
      if (node.expression.left.name === 'x') {
        this.remove();
      }
    }
  },
});

console.log(ast.body.length); // 1
console.log(ast.body[0].expression.left.name); // y
console.log(ast.body[1]); // undefined

元は文が二つあるのでast.body.lengthは2だったが、x = 1;を削除してy = 2;のみが残ったため、ast.body.lengthは1になった。

Escodegen

ソースコードの検証なら Esprima と Estraverse だけで行えるが、ソースコードの変換を行いたい場合、Estraverse で書き換えた AST をソースコードに戻す必要がある。

Escodegen を使うと、AST をソースコードに変換できる。

github.com

generateメソッドに AST を渡すと、ソースコードを返す。

先程のreplaceした結果を、ソースコードに復元してみる。

const esprima = require('esprima');
const estraverse = require('estraverse');
const escodegen = require('escodegen');

const ast = esprima.parseScript('x = 1');

const convertedAst = estraverse.replace(ast, {
  enter(node) {
    if (node.type === 'Literal') {
      return {
        ...node,
        value: 9,
        raw: '9',
      };
    }
    return node;
  },
});

console.log(escodegen.generate(convertedAst)); // x = 9;

AST に変換した時点で空行やコメントなどは取り除かれているので、それらを含まない状態のソースコードが生成される。

const esprima = require('esprima');
const escodegen = require('escodegen');

const str = `const x = 1;

const   y =   2 ;

// comment
const z=3;

`;

const ast = esprima.parseScript(str);

console.log(escodegen.generate(ast));
// const x = 1;
// const y = 2;
// const z = 3;