Transform AST in UglifyJS

Question

I am using UglifyJS to parse, minify and transform JS codes. One of my goal is to transform the AST by inserting new variable definitions. For example:

var x;
var y;
x = 1;
y = x;
x = 3;

I want to insert a new variable definition "var _x" into a random position, such as before the statement "y = x". The transformed code should be like this:

var x;
var y;
x = 1;
var _x = x;
y = _x;
_x = 3;

I have tried the TreeTransformer in UglifyJS. Updating the symbol reference to be the new one (x -> _x) is no problem. But I am unclear about obtaining the right position for the insertion with TreeTransformer. Can anyone share some insights? Some code examples would be better!

Answer 1

Here's how I solved this.

var _ = require('lodash');
var path = require('path');
var Promise = require('bluebird');
var fs = Promise.promisifyAll(require('fs'));
var UglifyJS = require('uglify-js');

fs.readFileAsync(path.join(__dirname, 'source.js'), 'utf8').then(function (file) {
  var topLevel = UglifyJS.parse(file);

  var transformed = topLevel.transform(new UglifyJS.TreeTransformer(function (node, descend) {
    if (node instanceof UglifyJS.AST_Toplevel) {
      node = node.clone();

      // Add new variable declaration
      node.body.unshift(new UglifyJS.AST_Var({
        definitions: [
          new UglifyJS.AST_VarDef({
            name: new UglifyJS.AST_SymbolVar({
              name: '_x'
            })
          })
        ]
      }));

      // Replace existing assignment statement with new statements
      var index = _(node.body).findIndex(function (node) {
        return node instanceof UglifyJS.AST_SimpleStatement &&
          node.body instanceof UglifyJS.AST_Assign &&
          node.body.left instanceof UglifyJS.AST_SymbolRef &&
          node.body.left.name === 'y';
      });
      var assignmentStatement = node.body[index].clone();
      node.body.splice(
        index,
        1,
        new UglifyJS.AST_SimpleStatement({
          body: new UglifyJS.AST_Assign({
            left: new UglifyJS.AST_SymbolRef({name: '_x'}),
            operator: '=',
            right: new UglifyJS.AST_SymbolRef({name: 'x'})
          })
        }),
        new UglifyJS.AST_SimpleStatement({
          body: new UglifyJS.AST_Assign({
            left: assignmentStatement.body.left.clone(),
            operator: assignmentStatement.body.operator,
            right: new UglifyJS.AST_SymbolRef({name: '_x'})
          })
        })
      );

      descend(node, this);
      return node;
    }

    node = node.clone();
    descend(node, this);
    return node;
  }));

  var result = transformed.print_to_string({beautify: true});
  return fs.writeFileAsync(path.join(__dirname, 'output.js'), result);
});

Here is the output:

var _x;

var x;

var y;

x = 1;

_x = x;

y = _x;

x = 3;

Uglify hoists variable declarations in its AST format, so I follow the same rule.