AST 實戰

歡迎關注個人公衆號睿Talk，獲取我最新的文章：

1、前言

最近忽然對 AST 產生了興趣，深刻了解後發現它的使用場景還真的很多，不少咱們平常開發使用的工具都跟它息息相關，如 Babel、ESLint 和 Prettier 等。本文除了介紹 AST 的一些基本概念外，更偏重實戰，講解如何利用它來對代碼進行修改。

2、基本概念

AST 全稱 Abstract Syntax Tree，也就是抽象語法樹，它是將編程語言轉換成機器語言的橋樑。瀏覽器在解析 JS 的過程當中，會根據 ECMAScript 標準將字符串進行分詞，拆分爲一個個語法單元。而後再遍歷這些語法單元，進行語義分析，構造出 AST。最後再使用 JIT 編譯器的全代碼生成器，將 AST 轉換爲本地可執行的機器碼。以下面一段代碼：

function add(a, b) {
    return a + b;
}

進行分詞後，會獲得這些 token：

對 token 進行分析，最終會獲得這樣一棵 AST（簡化版）：

{
  "type": "Program",
  "body": [
    {
      "type": "FunctionDeclaration",
      "id": {
        "type": "Identifier",
        "name": "add"
      },
      "params": [
        {
          "type": "Identifier",
          "name": "a"
        },
        {
          "type": "Identifier",
          "name": "b"
        }
      ],
      "body": {
        "type": "BlockStatement",
        "body": [
          {
            "type": "ReturnStatement",
            "argument": {
              "type": "BinaryExpression",
              "left": {
                "type": "Identifier",
                "name": "a"
              },
              "operator": "+",
              "right": {
                "type": "Identifier",
                "name": "b"
              }
            }
          }
        ]
      }
    }
  ],
  "sourceType": "module"
}

拿到 AST 後就能夠根據規則轉換爲機器碼了，在此再也不贅述。

3、Babel 工做原理

AST 除了能夠轉換爲機器碼外，還能作不少事情，如 Babel 就能經過分析 AST，將 ES6 的代碼轉換成 ES5。

Babel 的編譯過程分爲 3 個階段：

1. 解析：將代碼字符串解析成抽象語法樹
2. 變換：對抽象語法樹進行變換操做
3. 生成：根據變換後的抽象語法樹生成新的代碼字符串

Babel 實現了一個 JS 版本的解析器Babel parser，它能將 JS 字符串轉換爲 JSON 結構的 AST。爲了方便對這棵樹進行遍歷和變換操做，babel 又提供了traverse工具函數。完成 AST 的修改後，能夠使用generator生成新的代碼。

4、AST 實戰

下面咱們來詳細看看如何對 AST 進行操做。先建好以下的代碼模板：

import parser from "@babel/parser";
import generator from "@babel/generator";
import t from "@babel/types";
import traverser from "@babel/traverse";

const generate = generator.default;
const traverse = traverser.default;

const code = ``;
const ast = parser.parse(code);

// AST 變換

const output = generate(ast, {}, code);

console.log("Input \n", code);
console.log("Output \n", output.code);

• 構造一個 hello world

打開 AST Explorer，將左側代碼清空，再輸入 hello world，能夠看到先後 AST 的樣子：

// 空
{
  "type": "Program",
  "body": [],
  "sourceType": "module"
}

// hello world
{
  "type": "Program",
  "body": [
    {
      "type": "ExpressionStatement",
      "expression": {
        "type": "Literal",
        "value": "hello world",
        "raw": "'hello world'"
      },
      "directive": "hello world"
    }
  ],
  "sourceType": "module"
}

接下來經過代碼構造這個ExpressionStatement:

const code = ``;
const ast = parser.parse(code);

// 生成 literal
const literal = t.stringLiteral('hello world')
// 生成 expressionStatement
const exp = t.expressionStatement(literal)  
// 將表達式放入body中
ast.program.body.push(exp)

const output = generate(ast, {}, code);

能夠看到 AST 的建立過程就是自底向上建立各類節點的過程。這裏咱們藉助 babel 提供的types對象幫咱們建立各類類型的節點。更多類型能夠查閱這裏。

一樣道理，下面咱們來看看如何構造一個賦值語句：

const code = ``;
const ast = parser.parse(code);
 
// 生成 identifier
const id = t.identifier('str')
// 生成 literal
const literal = t.stringLiteral('hello world')
// 生成 variableDeclarator
const declarator = t.variableDeclarator(id, literal)
 // 生成 variableDeclaration
const declaration = t.variableDeclaration('const', [declarator])

// 將表達式放入body中
ast.program.body.push(declaration)

const output = generate(ast, {}, code);

• 獲取 AST 中的節點

下面咱們將對這段代碼進行操做：

export default {
  data() {
    return {
      count: 0
    }
  },
  methods: {
    add() {
      ++this.count
    },
    minus() {
      --this.count
    }
  }
}

假設我想獲取這段代碼中的data方法，能夠直接這麼訪問：

const dataProperty = ast.program.body[0].declaration.properties[0]

也能夠使用 babel 提供的traverse工具方法：

const code = `
export default {
  data() {
    return {
      count: 0
    }
  },
  methods: {
    add() {
      ++this.count
    },
    minus() {
      --this.count
    }
  }
}
`;

const ast = parser.parse(code, {sourceType: 'module'});
 
// const dataProperty = ast.program.body[0].declaration.properties[0]

traverse(ast, {
  ObjectMethod(path) {
    if (path.node.key.name === 'data') {
      path.node.key.name = 'myData';
      // 中止遍歷
      path.stop();
    }
  }
})

const output = generate(ast, {}, code);

traverse方法的第二個參數是一個對象，只要提供與節點類型同名的屬性，就能獲取到全部的這種類型的節點。經過path參數能訪問到節點信息，進而找出須要操做的節點。上面的代碼中，咱們找到方法名爲data的方法後，將其更名爲myData，而後中止遍歷，生成新的代碼。

• 替換 AST 中的節點

能夠使用replaceWith和replaceWithSourceString替換節點，例子以下：

// 將 this.count 改爲 this.data.count

const code = `this.count`;
const ast = parser.parse(code);

traverse(ast, {
  MemberExpression(path) {
    if (
      t.isThisExpression(path.node.object) &&
      t.isIdentifier(path.node.property, {
        name: "count"
      })
    ) {
      // 將 this 替換爲 this.data
      path
        .get("object")
        .replaceWith(
          t.memberExpression(t.thisExpression(), t.identifier("data"))
        );
        
      // 下面的操做跟上一條語句等價，更加直觀方便
      // path.get("object").replaceWithSourceString("this.data");
    }
  }
});

const output = generate(ast, {}, code);

• 插入新的節點

能夠使用pushContainer、insertBefore和insertAfter等方法來插入節點：

// 這個例子示範了 3 種節點插入的方法

const code = `
const obj = {
  count: 0,
  message: 'hello world'
}
`;

const ast = parser.parse(code);

const property = t.objectProperty(
  t.identifier("new"),
  t.stringLiteral("new property")
);

traverse(ast, {
  ObjectExpression(path) {
    path.pushContainer("properties", property);
    
    // path.node.properties.push(property);
  }
});

/* 
traverse(ast, {
  ObjectProperty(path) {
    if (
      t.isIdentifier(path.node.key, {
        name: "message"
      })
    ) {
      path.insertAfter(property);
    }
  }
}); 
*/

const output = generate(ast, {}, code);

• 刪除節點

使用remove方法來刪除節點：

const code = `
const obj = {
  count: 0,
  message: 'hello world'
}
`;

const ast = parser.parse(code);

traverse(ast, {
  ObjectProperty(path) {
    if (
      t.isIdentifier(path.node.key, {
        name: "message"
      })
    ) {
      path.remove();
    }
  }
});

const output = generate(ast, {}, code);

5、總結

本文介紹了 AST 的一些基本概念，講解了如何使用 Babel 提供的 API，對 AST 進行增刪改查的操做。​掌握這項技能，再加上一點想象力，就能製做出實用的代碼分析和轉換工具。