AST fun. Remove a function call from your bundle

What is AST

In its essence the code that we write is a text. As such it can be parsed and analyzed. The abstract syntax tree is still our code but parsed into a static structure. In the world of JavaScript this is one big object.

To illustrate the statement above we will parse a small snippet of code. There are several npm modules that can help us doing that. We will use Esprima. Consider the following app.js file:

var logger = {
  log() {}
}

We can read it as text and pass it to Esprima's parseScript method. Like so:

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

const file = `${__dirname}/app.js`;
const sourceCode = fs.readFileSync(file).toString('utf-8');
const ast = esprima.parseScript(sourceCode);

console.log(JSON.stringify(ast, null, 2));

The resulted ast object contains a static representation of the logger object above:

  {
  "type": "Program",
  "body": [
    {
      "type": "VariableDeclaration",
      "declarations": [
        {
          "type": "VariableDeclarator",
          "id": {
            "type": "Identifier",
            "name": "logger"
          },
          "init": {
            "type": "ObjectExpression",
            "properties": [
              {
                "type": "Property",
                "key": {
                  "type": "Identifier",
                  "name": "log"
                },
                "computed": false,
                "value": {
                  "type": "FunctionExpression",
                  "id": null,
                  "params": [],
                  "body": {
                    "type": "BlockStatement",
                    "body": []
                  },
                  "generator": false,
                  "expression": false,
                  "async": false
                },
                "kind": "init",
                "method": true,
                "shorthand": false
              }
            ]
          }
        }
      ],
      "kind": "var"
    }
  ],
  "sourceType": "script"
}

So, to achieve our goal (removing all logger.log calls), we can use the AST representation of the code. We can traverse it and delete the nodes that we don't need.

Removing AST node

I played with AST before and I know that modifying it is not a simple task. Very often changing or removing a node makes no sense. It really depends where in the tree this node is and what it represents. Thankfully there are some libraries which are here to help. I'm talking about ast-types. It offers bunch of helpful utility functions to operate with the tree.

Let's change our app.js a bit so it becomes close to the real case scenario:

(function (_index) {
  const xyz = 'abc';
  _index.logger.log('SOMETHING');
});

This is roughtly what my final bundle contains. We want to remove those _index.logger.log calls. If we transform that code to AST we will notice that there is an object of type CallExpression:

{
  "type": "CallExpression",
  "callee": {
    "type": "MemberExpression",
    "computed": false,
    "object": {
      "type": "MemberExpression",
      "computed": false,
      "object": {
        "type": "Identifier",
        "name": "_index"
      },
      "property": {
        "type": "Identifier",
        "name": "logger"
      }
    },
    "property": {
      "type": "Identifier",
      "name": "log"
    }
  },
  "arguments": [
    {
      "type": "Literal",
      "value": "SOMETHING",
      "raw": "'SOMETHING'"
    }
  ]
}

That object represents exactly our target.

ast-types offers an easy way to traverse the tree. There is a visit function and we can specify which nodes we want to look at:

const { visit } = require('ast-types');

visit(ast, {
  visitCallExpression(path) {
    console.log(JSON.stringify(path.node, null, 2));
    this.traverse(path);
  },
});

This will print out the same CallExpression object. Then it's a matter of recognizing the logger.log call. We can do it by running an if check on the fields of the node.

const callee = path.node.callee;
if (callee.type === 'MemberExpression') {
  const object = callee.object;
  const property = callee.property;
  if (
    object &&
    object.property &&
    object.property.name === 'logger' &&
    property &&
    property.name === 'log'
  ) {
    // deleting the node
    path.prune();
  }
}

Once found we have to remove the call from the tree. We can do it by using path.prune(). This removes the node at that given path. After this code runs our ast object contains a tree which is free of logger.log calls. We now need to transform it back to JavaScript source code. For this last job we will use another npm module called escodegen. escodegen.generate(ast) returns a string which is our same app.js code but cleaned up. Here is the whole example:

const fs = require('fs');
const esprima = require('esprima');
const { visit } = require('ast-types');
const escodegen = require('escodegen');

const file = `${__dirname}/app.js`; 
const sourceCode = fs.readFileSync(file).toString('utf-8');
const ast = esprima.parseScript(sourceCode);

visit(ast, {
  visitCallExpression(path) {
    const callee = path.node.callee;
    if (callee.type === 'MemberExpression') {
      const object = callee.object;
      const property = callee.property;
      if (
        object &&
        object.property &&
        object.property.name === 'logger' &&
        property &&
        property.name === 'log'
      ) {
        path.prune();
      }
    }
    this.traverse(path);
  },
});

const result = escodegen.generate(ast);
console.log(result);

And the result is:

(function (_index) {
  const xyz = 'abc';
});

---

Sure there are other ways to solve the same problem. However I really like this approach because it doesn't bother my writing. I don't have to add stuff like __DEV__ variables or architect the library in a certain way just so I help myself generating production-ready build. Tools like Esprima and escodegen become handy in such situations.