better whitespace detection in generator

lib/6to5/generator.js

@@ -14,14 +14,17 @@ function CodeGenerator(code, ast, opts) {
   opts = opts || {};
 
   this.tabWidth = 2;
+  this.comments = ast.comments || [];
+  this.tokens   = _.sortBy(_.compact([].concat(ast.tokens, this.comments)), "start");
   this.code     = code;
   this.opts     = opts;
   this.ast      = ast;
   this.buf      = "";
 
-  this.line    = 1;
+  this.usedLineWhitespace = [];
-  this.column  = 0;
+  this._indent            = 0;
-  this._indent = 0;
+  this.column             = 0;
+  this.line               = 1;
 
   if (opts.sourceMap) {
     this.map = new sourceMap.SourceMapGenerator({
@@ -34,7 +37,7 @@ function CodeGenerator(code, ast, opts) {
   }
 }
 
-CodeGenerator.prototype.mark = function (node, locType) {
+CodeGenerator.prototype.mark = function (node, type) {
   var loc = node.loc;
   if (!loc) return; // no location info
 
@@ -47,37 +50,76 @@ CodeGenerator.prototype.mark = function (node, locType) {
       line: this.line,
       column: this.column
     },
-    original: loc[locType]
+    original: loc[type]
   });
 };
 
-CodeGenerator.prototype.newline = function (removeLast) {
+CodeGenerator.prototype.newline = function (i, removeLast) {
   if (!this.buf) return;
-  if (removeLast) this.removeLastNewline();
+
+  if (_.isBoolean(i)) {
+    removeLast = i;
+    i = null;
+  }
+
+  if (_.isNumber(i)) {
+    var self = this;
+    _.times(i, function () {
+      self.newline(null, removeLast);
+    });
+    return;
+  }
+
+  if (removeLast && this.isLast("\n")) this.removeLast("\n");
+
+  this.removeLast(" ");
   this.buf = this.buf.replace(/\n(\s+)$/, "\n");
   this._push("\n");
 };
 
-CodeGenerator.prototype.removeLastNewline = function () {
+CodeGenerator.prototype.removeLast = function (cha) {
-  if (this.isLast("\n")) {
+  if (!this.isLast(cha)) return;
-    this.buf = this.buf.slice(0, -1);
+
+  this.buf = this.buf.slice(0, -1);
+
+  if (cha === "\n") {
     this.line--;
+  } else {
+    this.column--;
   }
 };
 
 CodeGenerator.prototype.semicolon = function () {
-  this._push(";");
+  this.push(";");
 };
 
 CodeGenerator.prototype.ensureSemicolon = function () {
   if (!this.isLast(";")) this.semicolon();
 };
 
+CodeGenerator.prototype.rightBrace = function () {
+  this.newline(true);
+  this.push("}");
+};
+
 CodeGenerator.prototype.keyword = function (name) {
   this.push(name);
   this.push(" ");
 };
 
+CodeGenerator.prototype.space = function () {
+  if (this.buf && !this.isLast([" ", "\n"])) {
+    this.push(" ");
+  }
+};
+
+CodeGenerator.prototype.printAndIndentOnComments = function (print, node) {
+  var indent = !!node.leadingComments;
+  if (indent) this.indent();
+  print(node);
+  if (indent) this.dedent();
+};
+
 CodeGenerator.prototype.printBlock = function (print, node) {
   if (t.isEmptyStatement(node)) {
     this.semicolon();
@@ -120,11 +162,17 @@ CodeGenerator.prototype._push = function (str) {
 CodeGenerator.prototype.isLast = function (cha, trimRight) {
   var buf = this.buf;
   if (trimRight) buf = buf.trimRight();
-  return _.last(buf) === cha;
+
+  var chars = [].concat(cha);
+  return _.contains(chars, _.last(buf));
 };
 
 CodeGenerator.prototype.getIndent = function () {
-  return util.repeat(this._indent * this.tabWidth);
+  return util.repeat(this.indentSize());
+};
+
+CodeGenerator.prototype.indentSize = function () {
+  return this._indent * this.tabWidth;
 };
 
 CodeGenerator.prototype.indent = function () {
@@ -141,6 +189,7 @@ CodeGenerator.prototype.generate = function () {
   this.print(ast);
 
   this.buf = this.buf.trimRight();
+  this.buf = this.buf.replace(/\{\n\n/g, "{\n");
 
   var map = this.map;
   if (map) map = map.toJSON();
@@ -160,125 +209,106 @@ CodeGenerator.prototype.buildPrint = function (parent) {
   };
 
   print.sequence = function (nodes, opts) {
-    return self.printSequence(print, nodes, opts);
+    opts = opts || {};
+    opts.statement = true;
+    return self.printJoin(print, nodes, opts);
   };
 
   return print;
 };
 
-CodeGenerator.prototype.printSequence = function (print, nodes, opts) {
+CodeGenerator.prototype.needsNewlineBefore = function (node) {
-  var tokens   = this.ast.tokens;
+  var startToken;
+  var endToken;
+  var tokens = this.tokens;
 
-  var self = this;
+  _.each(tokens, function (token, i) {
-  opts = opts || {};
+    // this is the token this node starts with
+    if (node.start === token.start) {
+      startToken = tokens[i - 1];
+      endToken = token;
+      return false;
+    }
+  });
 
-  var needsNewlineBefore = function (node) {
+  return this.hasWhitespaceBetween(startToken, endToken);
-    var startToken;
+};
-    var endToken;
 
-    _.each(tokens, function (token, i) {
+CodeGenerator.prototype.needsNewlineAfter = function (node) {
-      // this is the token this node starts with
+  var startToken;
-      if (node.start === token.start) {
+  var endToken;
-        startToken = tokens[i - 1];
+  var tokens = this.tokens;
-        endToken = token;
-        return false;
-      }
-    });
 
-    return self.hasWhitespaceBetween(startToken, endToken);
+  _.each(tokens, function (token, i) {
-  };
+    // this is the token this node ends with
+    if (node.end === token.end) {
+      startToken = token;
+      endToken = tokens[i + 1];
+      return false;
+    }
+  });
 
-  var needsNewlineAfter = function (node) {
+  return this.hasWhitespaceBetween(startToken, endToken);
-    var startToken;
-    var endToken;
-
-    _.each(tokens, function (token, i) {
-      // this is the token this node ends with
-      if (node.end === token.end) {
-        startToken = token;
-        endToken = tokens[i + 1];
-        return false;
-      }
-    });
-
-    return self.hasWhitespaceBetween(startToken, endToken);
-  };
-
-  opts.print = function (node) {
-    var needs = function (fn) {
-      if (!self.opts.whitespace) return;
-
-      if (node.start != null) {
-        // user node
-        if (!fn(node)) return;
-      } else {
-        // generated node
-        if (!t.needsWhitespace(node)) return;
-        if (self.isLast("{", true)) return;
-      }
-
-      self.newline();
-    };
-
-    print(node, node && {
-      before: function () {
-        needs(needsNewlineBefore);
-      },
-
-      after: function () {
-        needs(needsNewlineAfter);
-      }
-    });
-  };
-
-  return this.printJoin(print, nodes, "\n", opts);
 };
 
 CodeGenerator.prototype.hasWhitespaceBetween = function (startToken, endToken) {
   if (!endToken) return false;
 
-  var comments = this.ast.comments;
+  var start = startToken ? startToken.loc.end.line : 1;
-
+  var end   = endToken.loc.start.line;
-  var start = startToken ? startToken.end : 0;
+  var sep   = this.code.slice(start, end);
-  var end   = endToken.start;
-
-  var sep = this.code.slice(start, end);
 
   var lines = 0;
-  if (start > 0) lines--; // take off the current line
 
-  // remove comments
+  for (var line = start; line < end; line++) {
-  _.each(comments, function (comment) {
+    if (!_.contains(this.usedLineWhitespace, line)) {
-    // this comment is after the last token or befor ethe first
+      this.usedLineWhitespace.push(line);
-    if (comment.end > end || comment.start < start) return;
+      lines++;
+    }
+  }
 
-    var length = comment.end - comment.start;
+  return lines;
-
-    // compute the relative positions of the comment within the sliced node
-    // string
-    var startRelative = comment.start - start;
-    var endRelative = comment.end - start;
-
-    // remove the line this comment ends with
-    if (comment.type === "Line") lines--;
-
-    // remove comment
-    sep = sep.slice(0, startRelative) + util.repeat(length) + sep.slice(endRelative);
-  });
-
-  // check if there was a newline between the two nodes
-  lines += _.size(sep.match(/\n/g));
-  return lines > 0;
 };
 
 CodeGenerator.prototype.print = function (node, parent, opts) {
   if (!node) return "";
 
+  var self = this;
   opts = opts || {};
 
+  var newline = function (leading, fn) {
+    var ignoreDuplicates = false;
+
+    if (!opts.statement) {
+      if (t.isUserWhitespacable(node)) {
+        ignoreDuplicates = true;
+      } else {
+        return;
+      }
+    }
+
+    var lines = 0;
+
+    if (node.start != null) {
+      // user node
+      lines = fn.call(self, node);
+    } else {
+      // generated node
+      if (!leading) lines++; // definently include a single line
+
+      var needs = t.needsWhitespaceAfter;
+      if (leading) needs = t.needsWhitespaceBefore;
+      if (needs(node, opts.statement)) lines++;
+    }
+
+    self.newline(lines, ignoreDuplicates);
+  };
+
   if (this[node.type]) {
     this.printLeadingComments(node, parent);
 
+    newline(true, this.needsNewlineBefore);
+
     if (opts.before) opts.before();
     this.mark(node, "start");
 
@@ -291,6 +321,8 @@ CodeGenerator.prototype.print = function (node, parent, opts) {
     this.mark(node, "end");
     if (opts.after) opts.after();
 
+    newline(false, this.needsNewlineAfter);
+
     this.printTrailingComments(node, parent);
   } else {
     throw new ReferenceError("unknown node " + node.type + " " + JSON.stringify(node));
@@ -302,8 +334,6 @@ CodeGenerator.prototype.generateComment = function (comment) {
   if (comment.type === "Line") {
     val = "//" + val;
   } else {
-    this.newline(true);
-    val = val.replace(/\n(\s+)/g, "\n ");
     val = "/*" + val + "*/";
   }
   return val;
@@ -322,11 +352,19 @@ CodeGenerator.prototype.getComments = function (key, node, parent) {
     return [];
   }
 
+  var comments = [];
+  var nodes    = [node];
+  var self     = this;
+
   if (t.isExpressionStatement(node)) {
-    return [].concat(this._getComments(key, node), this._getComments(key, node.argument));
+    nodes.push(node.argument);
   }
 
-  return this._getComments(key, node);
+  _.each(nodes, function (node) {
+    comments = comments.concat(self._getComments(key, node));
+  });
+
+  return comments;
 };
 
 CodeGenerator.prototype._getComments = function (key, node) {
@@ -340,23 +378,40 @@ CodeGenerator.prototype._printComments = function (comments) {
 
   _.each(comments, function (comment, i) {
     // whitespace before
-    if (self.hasWhitespaceBetween(comment, comments[i - 1])) {
+    self.newline(self.needsNewlineBefore(comment));
-      self.newline();
+
+    var val = self.generateComment(comment);
+
+    if (self.column && !self.isLast(["\n", " ", "[", "{"])) {
+      self._push(" ");
     }
 
-    self.push(self.generateComment(comment));
+    //
-    self.newline();
+
+    var offset = comment.loc.start.column;
+    if (offset) {
+      var newlineRegex = new RegExp("\\n\\s{1," + offset + "}", "g");
+      val = val.replace(newlineRegex, "\n");
+    }
+
+    var indent = Math.max(self.indentSize(), self.column);
+    val = val.replace(/\n/g, "\n" + util.repeat(indent));
+
+    if (self.column == 0) {
+      val = self.getIndent() + val;
+    }
+
+    //
+
+    self._push(val);
 
     // whitespace after
-    if (self.hasWhitespaceBetween(comment, comments[i + 1])) {
+    self.newline(self.needsNewlineAfter(comment));
-      self.newline();
-    }
   });
 };
 
-CodeGenerator.prototype.printJoin = function (print, nodes, sep, opts) {
+CodeGenerator.prototype.printJoin = function (print, nodes, opts) {
   opts = opts || {};
-  sep  = sep || "\n";
 
   var self = this;
   var len  = nodes.length;
@@ -364,19 +419,18 @@ CodeGenerator.prototype.printJoin = function (print, nodes, sep, opts) {
   if (opts.indent) self.indent();
 
   _.each(nodes, function (node, i) {
-    if (opts.print) {
+    print(node, {
-      opts.print(node, i);
+      statement: opts.statement,
-    } else {
+      after: function () {
-      print(node);
+        if (opts.iterator) {
-    }
+          opts.iterator(node, i);
+        }
 
-    if (opts.iterator) {
+        if (opts.separator && i < len - 1) {
-      opts.iterator(node, i);
+          self.push(opts.separator);
-    }
+        }
-
+      }
-    if (i < len - 1) {
+    });
-      self.push(sep);
-    }
   });
 
   if (opts.indent) self.dedent();