make comments output more robust

fixes #112
fixes #218
fixes #372

lib/output.js

@@ -52,6 +52,7 @@ function is_some_comments(comment) {
 
 function OutputStream(options) {
 
+    var readonly = !options;
     options = defaults(options, {
         ascii_only       : false,
         beautify         : false,
@@ -427,6 +428,80 @@ function OutputStream(options) {
         return OUTPUT;
     };
 
+    function add_comments(node) {
+        var self = this;
+        var start = node.start;
+        if (!(start.comments_before && start.comments_before._dumped === self)) {
+            var comments = start.comments_before;
+            if (!comments) {
+                comments = start.comments_before = [];
+            }
+            comments._dumped = self;
+
+            if (node instanceof AST_Exit && node.value) {
+                var tw = new TreeWalker(function(node) {
+                    var parent = tw.parent();
+                    if (parent instanceof AST_Exit
+                        || parent instanceof AST_Binary && parent.left === node
+                        || parent.TYPE == "Call" && parent.expression === node
+                        || parent instanceof AST_Conditional && parent.condition === node
+                        || parent instanceof AST_Dot && parent.expression === node
+                        || parent instanceof AST_Sequence && parent.expressions[0] === node
+                        || parent instanceof AST_Sub && parent.expression === node
+                        || parent instanceof AST_UnaryPostfix) {
+                        var text = node.start.comments_before;
+                        if (text && text._dumped !== self) {
+                            text._dumped = self;
+                            comments = comments.concat(text);
+                        }
+                    } else {
+                        return true;
+                    }
+                });
+                tw.push(node);
+                node.value.walk(tw);
+            }
+
+            if (current_pos == 0) {
+                if (comments.length > 0 && options.shebang && comments[0].type == "comment5") {
+                    print("#!" + comments.shift().value + "\n");
+                    indent();
+                }
+                var preamble = options.preamble;
+                if (preamble) {
+                    print(preamble.replace(/\r\n?|[\n\u2028\u2029]|\s*$/g, "\n"));
+                }
+            }
+
+            comments = comments.filter(comment_filter, node);
+
+            // Keep single line comments after nlb, after nlb
+            if (current_col != 0
+                && !options.beautify
+                && comments.length > 0
+                && comments[0].nlb
+                && /comment[134]/.test(comments[0].type)) {
+                print("\n");
+            }
+
+            comments.forEach(function(c){
+                if (/comment[134]/.test(c.type)) {
+                    print("//" + c.value + "\n");
+                    indent();
+                }
+                else if (c.type == "comment2") {
+                    print("/*" + c.value + "*/");
+                    if (start.nlb) {
+                        print("\n");
+                        indent();
+                    } else {
+                        space();
+                    }
+                }
+            });
+        }
+    }
+
     var stack = [];
     return {
         get             : get,
@@ -464,7 +539,7 @@ function OutputStream(options) {
         with_square     : with_square,
         add_mapping     : add_mapping,
         option          : function(opt) { return options[opt] },
-        comment_filter  : comment_filter,
+        add_comments    : readonly ? noop : add_comments,
         line            : function() { return current_line },
         col             : function() { return current_col },
         pos             : function() { return current_pos },
@@ -500,7 +575,7 @@ function OutputStream(options) {
             use_asm = active_scope;
         }
         function doit() {
-            self.add_comments(stream);
+            stream.add_comments(self);
             self.add_source_map(stream);
             generator(self, stream);
         }
@@ -519,77 +594,10 @@ function OutputStream(options) {
 
     AST_Node.DEFMETHOD("print_to_string", function(options){
         var s = OutputStream(options);
-        if (!options) s._readonly = true;
         this.print(s);
         return s.get();
     });
 
-    /* -----[ comments ]----- */
-
-    AST_Node.DEFMETHOD("add_comments", function(output){
-        if (output._readonly) return;
-        var self = this;
-        var start = self.start;
-        if (start && !start._comments_dumped) {
-            start._comments_dumped = true;
-            var comments = start.comments_before || [];
-
-            // XXX: ugly fix for https://github.com/mishoo/UglifyJS2/issues/112
-            //               and https://github.com/mishoo/UglifyJS2/issues/372
-            if (self instanceof AST_Exit && self.value) {
-                self.value.walk(new TreeWalker(function(node){
-                    if (node.start && node.start.comments_before) {
-                        comments = comments.concat(node.start.comments_before);
-                        node.start.comments_before = [];
-                    }
-                    if (node instanceof AST_Function ||
-                        node instanceof AST_Array ||
-                        node instanceof AST_Object)
-                    {
-                        return true; // don't go inside.
-                    }
-                }));
-            }
-
-            if (output.pos() == 0) {
-                if (comments.length > 0 && output.option("shebang") && comments[0].type == "comment5") {
-                    output.print("#!" + comments.shift().value + "\n");
-                    output.indent();
-                }
-                var preamble = output.option("preamble");
-                if (preamble) {
-                    output.print(preamble.replace(/\r\n?|[\n\u2028\u2029]|\s*$/g, "\n"));
-                }
-            }
-
-            comments = comments.filter(output.comment_filter, self);
-
-            // Keep single line comments after nlb, after nlb
-            if (!output.option("beautify") && comments.length > 0 &&
-                /comment[134]/.test(comments[0].type) &&
-                output.col() !== 0 && comments[0].nlb)
-            {
-                output.print("\n");
-            }
-
-            comments.forEach(function(c){
-                if (/comment[134]/.test(c.type)) {
-                    output.print("//" + c.value + "\n");
-                    output.indent();
-                }
-                else if (c.type == "comment2") {
-                    output.print("/*" + c.value + "*/");
-                    if (start.nlb) {
-                        output.print("\n");
-                        output.indent();
-                    } else {
-                        output.space();
-                    }
-                }
-            });
-        }
-    });
-
     /* -----[ PARENTHESES ]----- */
 
     function PARENS(nodetype, func) {

lib/utils.js

@@ -324,7 +324,7 @@ function first_in_statement(stack) {
         if (p instanceof AST_Statement && p.body === node)
             return true;
         if ((p instanceof AST_Sequence      && p.expressions[0] === node) ||
-            (p instanceof AST_Call          && p.expression === node && !(p instanceof AST_New) ) ||
+            (p.TYPE == "Call"               && p.expression === node ) ||
             (p instanceof AST_Dot           && p.expression === node ) ||
             (p instanceof AST_Sub           && p.expression === node ) ||
             (p instanceof AST_Conditional   && p.condition === node  ) ||

test/mocha/comment.js

@@ -47,4 +47,70 @@ describe("Comment", function() {
             }, fail, tests[i]);
         }
     });
+
+    it("Should handle comment within return correctly", function() {
+        var result = uglify.minify([
+            "function unequal(x, y) {",
+            "    return (",
+            "        // Either one",
+            "        x < y",
+            "        ||",
+            "        y < x",
+            "    );",
+            "}",
+        ].join("\n"), {
+            compress: false,
+            mangle: false,
+            output: {
+                beautify: true,
+                comments: "all",
+            },
+        });
+        if (result.error) throw result.error;
+        assert.strictEqual(result.code, [
+            "function unequal(x, y) {",
+            "    // Either one",
+            "    return x < y || y < x;",
+            "}",
+        ].join("\n"));
+    });
+
+    it("Should handle comment folded into return correctly", function() {
+        var result = uglify.minify([
+            "function f() {",
+            "    /* boo */ x();",
+            "    return y();",
+            "}",
+        ].join("\n"), {
+            mangle: false,
+            output: {
+                beautify: true,
+                comments: "all",
+            },
+        });
+        if (result.error) throw result.error;
+        assert.strictEqual(result.code, [
+            "function f() {",
+            "    /* boo */",
+            "    return x(), y();",
+            "}",
+        ].join("\n"));
+    });
+
+    it("Should not drop comments after first OutputStream", function() {
+        var code = "/* boo */\nx();";
+        var ast = uglify.parse(code);
+        var out1 = uglify.OutputStream({
+            beautify: true,
+            comments: "all",
+        });
+        ast.print(out1);
+        var out2 = uglify.OutputStream({
+            beautify: true,
+            comments: "all",
+        });
+        ast.print(out2);
+        assert.strictEqual(out1.get(), code);
+        assert.strictEqual(out2.get(), out1.get());
+    });
 });