通过 ast_walker 来操作 AST

2011-09-29 21:34

通过 ast_walker 来操作 AST

by lifesinger

at 2011-09-29 13:34:46

original http://lifesinger.wordpress.com/2011/09/29/ast-walker/

通过 UglifyJS 可以得到 js 代码的 AST:

var fs = require('fs');
var jsp = require('uglify-js').parser;

var code = fs.readFileSync('test.js', 'utf-8');
var ast = jsp.parse(code);

得到的 ast 是以数组形式表示的抽象语法树。

spm 项目里,要提取 js 模块里的依赖信息。最裸的做法是通过正则来实现:

function parseStatic(inputFile, charset) {
  var ast = getAst(inputFile, charset);
  var deps;

  var pattern = /,stat,call,name,define,(?:(?:[^,]+,){2})?array(,.*?,)(?:function|name),/;
  var match = ast.toString().match(pattern);

  if (match && match[1]) {
    deps = [];
    var t = match[1].match(/,string,[^,]+/g);

    if (t && t.length) {
      deps = t.map(function(s) {
        // s: ,string,xxx
        return s.slice(8);
      });
    }
  }

  return deps;
}

详见:Dependences.js

正则的实现方式在绝大部分情况下不会有问题,但代码的可读性和扩展性不好。

查看 UglifyJS 的源码,发现内部提供了 ast_walker 方法。我们可以简单封装成 Ast.js

var uglifyjs = require('uglify-js');
var pro = uglifyjs.uglify;
var Ast = exports;

Ast.walk = function(ast, type, walker) {
  var w = pro.ast_walker();

  var walkers = {};
  walkers[type] = function() {
    walker(this);
  };

  ast = w.with_walkers(walkers, function() {
    return w.walk(ast);
  });

  return ast;
};

这样,获取依赖信息的代码可以调整为:

function parseStatic(inputFile, charset) {
  var ast = getAst(inputFile, charset);
  var deps, times = 0;

  Ast.walk(ast, 'stat', function(stat) {
    if (stat.toString().indexOf('stat,call,name,define,') !== 0) {
      return stat;
    }

    if (++times > 1) {
      throw 'Found multiple "define" in one module file. It is NOT allowed.';
    }

    // stat[1]:
    //     [ 'call',
    //       [ 'name', 'define' ],
    //       [ [Object], [Object], [Object ] ] ]
    var argsAst = stat[1][2];
    var depsAst;

    // argsAst:
    //   [ [ 'string', 'a' ],
    //     [ 'array', [ [Object], [Object] ] ],
    //     [ 'function', null, [], [] ] ]
    argsAst.some(function(item) {
      // NOTICE: the deps MUST be literal, it can NOT be a reference.
      if (item[0] === 'array') {
        depsAst = item[1];
        return true;
      }
    });

    if (!depsAst) {
      return stat;
    }

    // depsAst:
    //   [ [ 'string', 'b' ], [ 'string', 'c' ] ]
    deps = [];
    depsAst.forEach(function(item) {
      if (item[0] === 'string') {
        deps.push(item[1]);
      }
    });

    return stat;
  });

  return deps;
}

详见:Dependences.js. 功能更可靠,可读性也更好了。

还可以通过 UglifyJS 隐藏的第三参数来让 AST 带上 type/value/start/end 等信息:

var ast = jsp.parse(code, false, true);

通过 ast_walker 来操作 AST, 可以完成很多意想不到的工作。比如获取某个闭包里的所有局部变量、删除某个特定函数、或者给所有 constructor 添加上调试信息等。老赵的 Jscex 也是基于这一原理,有兴趣的可以深入挖掘。