Fix and refactor ast-collector.

lib/abstract-collector.js

@@ -38,7 +38,7 @@ function AbstractCollector(config) {
 
   p.render = render;
   function render(buffer) {
-    var modules = this._modules;
+    var modules = this.getModules();
 
     buffer.push(this.renderRuntime());
 

lib/ast-collector.js

@@ -6,7 +6,13 @@ var util = require('util'),
     processor = uglify.uglify,
     parser = uglify.parser;
 
-var RUNTIME = parser.parse(abstractCollector.getRuntimeSrcCode('modulr.sync.js'));
+var RUNTIME = getRuntimeAst('modulr.sync.js');
+
+exports.getRuntimeAst = getRuntimeAst;
+function getRuntimeAst(filename) {
+  var src = abstractCollector.getRuntimeSrcCode(filename);
+  return toAstBody(parser.parse(src));
+}
 
 exports.createAstCollector = createAstCollector;
 exports.create = createAstCollector;
@@ -30,7 +36,7 @@ util.inherits(AstCollector, SuperClass);
   p.encloseModule = encloseModule;
   function encloseModule(m) {
     var ast = this.getModuleAst(m);
-    return ["function", null, ["require", "exports", "module"], this.getSubtree(ast)];
+    return ["function", null, ["require", "exports", "module"], toAstBody(ast)];
   }
 
   p.escapeModule = escapeModule;
@@ -63,14 +69,10 @@ util.inherits(AstCollector, SuperClass);
   function getModuleAst(m) {
     var ast;
     if (m.duplicateOf) {
-      ast = ["toplevel",
-        [
-          ["stat",
-            ["assign", true,
-              ["dot", ["name", "module"], "exports"],
-              ["call", ["name", "require"], [["string", this.getModuleId(m.duplicateOf)]]]
-            ]
-          ]
+      ast = ["stat",
+        ["assign", true,
+          ["dot", ["name", "module"], "exports"],
+          ["call", ["name", "require"], [["string", this.getModuleId(m.duplicateOf)]]]
         ]
       ];
     } else {
@@ -79,15 +81,75 @@ util.inherits(AstCollector, SuperClass);
     return ast;
   }
 
-  p.getSubtree = getSubtree;
-  function getSubtree(ast) {
-    return ast[1];
-  }
-
   p.toString = toString;
   function toString() {
-    var ast = parser.parse('');
-    this.render(this.getSubtree(ast));
-    return this.generateCode(ast);
+    var buffer = createAstBuffer();
+    this.render(buffer);
+    return this.generateCode(buffer.ast);
   }
+
 })(AstCollector.prototype);
+
+
+exports.toAstBody = toAstBody;
+function toAstBody(ast) {
+  // A body of an AST is a plain array. It can be used
+  // both as the body of a function and the body of a
+  // program.
+  if (isAstBody(ast)) {
+    // We're all good. This is already a body.
+    return ast;
+  }
+
+  if (isAstRoot(ast)) {
+    // We're dealing with a root. Get it's body.
+    return ast[1];
+  }
+  // Lower level than a body. Convert it to a body by
+  // wrapping it into an array.
+  return [ast];
+}
+
+exports.toAstRoot = toAstRoot;
+function toAstRoot(ast) {
+  if (isAstRoot(ast)) {
+    // Already an AST root.
+    return ast;
+  }
+  if (isAstBody(ast)) {
+    // This is a body, wrap it up in a root node.
+    return ['toplevel', ast];
+  }
+  // Lower level than a body. More wrapping up.
+  return ['toplevel', [ast]];
+}
+
+exports.isAstRoot = isAstRoot;
+function isAstRoot(ast) {
+  return ast[0] === 'toplevel';
+}
+
+exports.isAstBody = isAstBody;
+function isAstBody(ast) {
+  // AST bodies are the only anonymous nodes.
+  return typeof ast[0] === 'object';
+}
+
+exports.createAstBuffer = createAstBuffer;
+function createAstBuffer() {
+  return new AstBuffer();
+}
+
+exports.AstBuffer = AstBuffer;
+function AstBuffer() {
+  this.ast = parser.parse('');
+  this.body = toAstBody(this.ast);
+}
+
+(function(p) {
+  p.push = push;
+  function push(items) {
+    var body = this.body;
+    body.push.apply(body, toAstBody(items));
+  }
+})(AstBuffer.prototype);

lib/resolved-ast-collector.js

@@ -1,14 +1,12 @@
 var util = require('util'),
-    abstractCollector = require('./abstract-collector'),
-    SuperClass = require('./ast-collector').AstCollector,
+    astCollector = require('./ast-collector'),
+    SuperClass = astCollector.AstCollector,
     _super = SuperClass.prototype,
     uglify = require('uglify-js'),
-    processor = uglify.uglify,
-    parser = uglify.parser,
-    walker = processor.ast_walker(),
+    walker = uglify.uglify.ast_walker(),
     identifier = require('module-grapher/lib/identifier');
 
-var RUNTIME = parser.parse(abstractCollector.getRuntimeSrcCode('modulr.sync.resolved.js'));
+var RUNTIME = astCollector.getRuntimeAst('modulr.sync.resolved.js');
 
 exports.createResolvedAstCollector = createResolvedAstCollector;
 exports.create = createResolvedAstCollector;