JavaScriptでASTを使ってコードをインジェクションしてみる

10/22の技術書典3に、東京ラビットハウスでJavaScript AST本を出します。ASTを知らない人でも簡単にソースコードをハックできるようになる入門＆実用的な本です。今回の記事はその本を書いてる過程で生まれた物を記事にしたものです。

dev-injector というツールを作ったのでその過程について説明します。

※注意: Babel系のみ説明します＆Babel7 (現時点ではまだbeta) を使っています。

ユースケース

Twitterでユースケースを募集したところ

「テストフレームワークから関数フックしやすいコーディングとは？」辺りがあります。「exports.hook.method = method」とかやってますけど「もっと適切な方法（設計含む）は何処なら見つけられる？

というのを頂きました。

exportしたい・したくない、あるいは対象のコードが大きい、そもそもユニットテストの無いレガシーコードだった、などとテストにおいては様々なケースがあります。

そういう難しい事情を無視してさっくりテストを書きたいですよね。ASTを使えばテスト対象のコードに手を加えず部分的に改変できるので、楽にテストコードを作れるのです。

ASTではアイデアさえあればコードを削除、改変、追加など任意に行えます。

今回作ったdev-injectorでは、v0.1.2現時点では、変数宣言の初期化を書き換えたり、関数宣言を書き換えたり、クラス定義を書き換えたり、コードの一番最後にコードを追加できます。

ASTをいじるとは

ASTというのは抽象構文木のことで、コンパイラの内部表現でよく使われるものです。JavaScriptの世界ではASTは大きなエコスシステムを構築していて、AST関連で一番有名なツールはBabelでしょうか。Babelのプラグインは基本的には、ASTを加工するものです。

コンパイラの本というとなんか分厚い本が多くて尻込みする人が多いかも知れませんが、字句解析・構文解析と言った難しい処理はparserがやってくれますて、楽にいじれるASTをはきだしてくれるのです。Babel系だとBabylonというparserがあります。BabylonでJavaScript AST事始めという記事に書いているので出来たら目を通しておいてください。

ASTはその名のとおり、ツリー構造です。再帰的に処理するのが一般的ですが、生のASTを触る場合、スコープなどの処理は自前でしなければいけません。そこでtraverserとかwalkerとか呼ばれるタイプのものを使います。Babel系だとbabel-traverseです。

traverserでASTを加工した後は、ASTからソースを生成するgeneratorを使います。babel系だとbabel-generatorが一般的ですが、prettierもこの目的に使えます。

さて、ASTについては、AST explorer を積極的に使いましょう。これはAST操作では手放せないとても便利なサイトです。

実際にやってみる

$ npm i babel-core babylon babel-types -D

$ npm i babel-core@next babylon@next babel-types@next -D

Babel7系ではTypeScriptに対応してたり美味しい点が多そうだし、せっかく新規で覚えるならBabel6よりは7系でしょう。まだベータ版なので@nextが必要ですが、正式版になれば@nextはいりません。

const {transform} = require('babel-core')
const babylon = require('babylon')
const t = require('babel-types')

const visitor = {
    VariableDeclarator: (nodePath) => {
        nodePath.get('init').replaceWith(t.numericLiteral(10))
    }
}

const plugin = {
    visitor,
}

const source = 'const a = 1'

console.log('before:')
console.log(source) // const a = 1

const {code} = transform(source, {plugins: [plugin]})
console.log('\nafter:')
console.log(code)   // const a = 10

visitorはvisitor patternのvisitorです。VariableDeclarator (変数定義)にたどり着いたらそのノードの init (変数初期化)を replaceWithで置き換えています。置き換え先はASTを簡略化したものです。babel-typesを使えば、t.numericLiteralのように簡単にASTを生成できます。

さて、前回の記事BabylonでJavaScript AST事始めでは、Babylonがはき出すASTはNode型であると言いましたが、babel-traverserでは、NodePath型が使われます。NodePathには、そのスコープ内にどの変数があるか？など含めたコンテキスト情報と、Node型がそれぞれ入っていますが直接操作する事はなく参照するだけになるでしょう。

ASTを加工するならnodePath.get('init')のようにアクセスしたいプロパティに対応するNodePathをまず取得します。そしてそれに対してreplaceWith などのメソッドを使います。

pluginは実際にBabelのプラグインです。Babelのプラグインはvisitorだけでもあれば動作します。このプラグインを使って、babel-coreのtransformという関数を叩くと、codeやastが入ったオブジェクトが帰ってきますが、今回はcodeだけを使っています。

ヘルパーを作る

t.numericLiteral()のようなBabel-typesのメソッドを使ってもいいんですが、実はBabylonでparseしたASTを少し加工するだけでも良いのです。

const createNode = (sourceCode) => {
    const {body} = babylon.parse(sourceCode).program
    assert(body.length === 1)
    return body[0]
}

ただし、bodyは配列で複数のステートメントがある場合lengthが2以上に増えます。場合によって、複数のステートメントがある場合、Blockで囲う必要もあるので、ひとまず1つのステートメントにのみ対応しています。

式(Expression)でよければbabylon.parseExpressionで生成できます。場合によってはそれを使う事もあるでしょう。

dev-injectorで実際にやっているコード変換処理

dev-injectorで実際にやっているコード変換をそれぞれ抜粋して説明します。

ソースの最後にコードを追加する

// replacedCode に追加したいコードが文字列で入っているとする。

const visitor = {
  Program: {
    exit: (nodePath) => {
      if (replaceCode) {
        const bodies = nodePath.get('body')
        bodies[bodies.length - 1].insertAfter(createNode(replaceCode))
      }
    }
  },
}

nodePath.get('body')でソースコードのトップレベルのコードに該当するものが配列で取得できます。

こんな感じですね。変数宣言、関数宣言、クラス宣言の3つがbodyに入っています。

ソースの最後は、クラス宣言の後ろという事になります。なので、bodyの長さを取得して最後のやつにinsertAfter()でコードを追加しているのです。

もちろんクラスの先頭であれば、bodies[0].insertBefore()で可能です。

変数宣言の初期化式を書き換える

// replaceCodeに変換したいコードが入っているとする
// targetIdに変数名が入ってるとする

const visitor = {
  VariableDeclarator: (nodePath) => {
    if (t.isIdentifier(nodePath.node.id && nodePath.node.id === targetId)) {
      nodePath.get('init').replaceWith(babylon.parseExpression(replaceCode))
    }
  },
}

JSの変数定義は、VariableDeclarationの中のdeclarations配列に、VariableDeclaratorが複数(基本的には1つ)入っている点に注意が必要です。

今回は初期化式を置き換えるだけなので、VariableDeclarationはいじりません。

あと、初期化式は、式が必要なので、babylon.parseExpressionを使っています。

関数定義を置き換える

// replaceCodeに変換したいコード
// targetIdに関数名

const visitor = {
  FunctionDeclaration: (nodePath) => {
    if (t.isIdentifier(nodePath.node.id) && nodePath.node.id === targetId) {
      nodePath.insertBefore(createNode(replaceCode))
      nodePath.remove()
    }
  },
}

変数宣言の時と似ていますが、違うのは自分自身を書き換えていることです。これには注意が必要で、nodepath.replaceWithだと無限ループになってしまいます。なので、いったんinsertBefore()でこのコードの前に挿入してから自分をremove()しています。

※これが本当に望ましいやり方かは知りません。もし、これより望ましい方法をご存じの方は教えていただけると幸いです

クラス定義を置き換える方法は、ClassDeclaration: に同じコードを書くだけです。

動的にインジェクションする

さて、babelプラグインを作るだけなら、もうこれで作る事ができましたが、動的にインジェクションしないと今回の目的は達成できません。babel-registerを使ってもいいのですが同じ事を自前でやってしまいましょう。

$ npm i source-map-support -D

ソースコードをトランスパイルする時はSourceMapが無いと不便なのでまずこれをインストールします。

const {transformFileSync} = require('babel-core')
const {install} = require('source-map-support')

install({hookRequire: true})

// pluginは既にあるとする

const olds = {}
const ext = '.js'
olds[ext] = require.extensions[ext]
require.extensions[ext] = (m, filename) => {
  if (isNeedInjection(filename)) {
    const {code} = transformFile(filename, {plugins: [plugin])
    m._compile(code, filename)
  } else {
    olds[ext](m, filename)
  }
}

Node.jsのrequireにまつわる隠しAPIを叩いてるようなものですが、こうすることで、require処理に対してフックできます。(今、推奨されるやり方かどうかは不明。一応Node8.6.0でも動作しています)

まとめ

• JavaScript ではASTのサブシステムが充実しているので簡単にソースコードを変換できる
• Babelのプラグインは割と簡単に書ける
• Node.jsのrequireをハックすることで、ソースを動的に変換できる

突っ込みとか、この記事がわかりやすい・わかりにくい、ASTに関して知りたい事とかあったら遠慮無くコメントやTwitterで話しかけて頂ければ幸いです。

10/22の技術書典3に、東京ラビットハウスで今回の記事のようなものをまとめあげた本を出します。良ければサークルチェックなどしていただければ嬉しいです。