如何利用AST树修改源码

正则匹配

平时在开发工具类程序时，经常会出现需要对源码进行操作。源码操作本质上是对读取进程序内存中的字符进行修改。遇到字符修改时，我们很容易联想到使用正则表达式去匹配替换。

正则方法的好处是面对短小简单的需求，可以快速写出比较短的正则表达式。而正则表达式的缺点也很明显：

1. 即使使用了正则表达式可视化工具后，复杂的正则表达式仍旧有很高的理解成本。
2. 构造复杂正则表达式的开发成本很高。
3. 正则表达式中去做分支判断比较困难。

于是当我们面对复杂且多变的匹配规则时，我们应该要用上源码操作的另一种方法，也是本文的主题——AST树操作。

AST是什么

AST(Abstract Syntax Tree)是一种语法树，其大体结构和数据结构课上学的树结构是一样的。程序的源码经过词法分析和语法分析后即可生成一棵AST树。

let a = b + c;

例如以上这句代码可以分解为一个let声明、一个名为a的变量、一个名为b的变量、一个名为c的变量和一个名为加号的操作符

上图中的body就是代码块，当源代码被转换为AST树形结构时，该树形结构可以用JSON对象来表示。

tip: 这个网址 astexplorer.net/ 可以将代码转换为可视化的AST树。

parser

我们使用AST树操作去修改源码时，一般分为三个步骤：

1. parser: 将代码字符串转换为AST树，即生成一个JSON对象
2. transform: 使用工具对生成的AST树进行增删查改等操作
3. generate: 将修改完毕的AST树重新转换为源码字符串

这里我们使用babel的@babel/parser工具来转换代码。

const parser = require("@babel/parser");
// 将源码转换为AST
const ast = parser.parse(source, {
    sourceType: "module"
});

@babel/parser在不添加plugins的时候只能转换标准的javascript代码，如果有特殊的需要，可以添加相对应的plugins来扩展。

例如添加"jsx"这个plugins后就可以解析jsx代码

const ast = parser.parse(source, {
    sourceType: "module",
    plugins: [
        "jsx"
    ]
});

transform

在对AST树进行增删查改之前，我们首先要通过遍历去找到目标节点。@babel/traverse这个工具库提供了大量API帮助我们快速定位到目标节点。

const traverse = require('@babel/traverse');
traverse.default(ast, {
    enter(path) {
        // 如果变量名为abc
        if (path.isIdentifier({name: "abc"})) {
        }
    }
})

这里@babel/traverse引入了一个path的概念，path是一个对象，它表示两个节点之间的关联，我们可以通过path.node来访问其节点属性、通过path.parentPath访问其父路径等等。

在找到目标节点后我们可以使用@babel/traverse提供的增删查改API来进行修改。

观察上面生成的简易AST树，我们发现AST树节点往往比较复杂，自己构造多层JSON对象会十分繁琐。在添加节点时，可以借助@babel/types来快速生成AST树节点。

const t = require('@babel/types');
t.stringLiteral("this is a test");

上面这句代码可以生成一个内容为this is a test的AST字符串节点。

generate

在对AST树进行修改操作后，我们就可以调用@babel/generator来生成新源码：

const generate = require('@babel/generator');
console.log(generate.default(ast).code)

至此，我们的源码替换过程就完成了。