make comments output more robust (#2633)

- improve handling of comments right after `return`
- retain comments after `OutputStream`
- preserve trailing comments
- fix handling of new line before comments
- handle comments around parentheses

fixes #88
fixes #112
fixes #218
fixes #372
fixes #2629

2 files changed, 175 insertions, 3 deletions

diff --git a/test/mocha/comment-filter.js b/test/mocha/comment-filter.js
index 0e4f3dff..25233d11 100644
--- a/test/mocha/comment-filter.js
+++ b/test/mocha/comment-filter.js
@@ -14,7 +14,7 @@ describe("comment filters", function() {
 
     it("Should be able to filter commments with the 'some' option", function() {
         var ast = UglifyJS.parse("// foo\n/*@preserve*/\n// bar\n/*@license*/\n//@license with the wrong comment type\n/*@cc_on something*/");
-        assert.strictEqual(ast.print_to_string({comments: "some"}), "/*@preserve*/\n/*@license*/\n/*@cc_on something*/\n");
+        assert.strictEqual(ast.print_to_string({comments: "some"}), "/*@preserve*/\n/*@license*/\n/*@cc_on something*/");
     });
 
     it("Should be able to filter comments by passing a function", function() {
@@ -55,12 +55,12 @@ describe("comment filters", function() {
             return true;
         };
 
-        assert.strictEqual(ast.print_to_string({comments: f}), "#!Random comment\n//test1\n/*test2*/\n");
+        assert.strictEqual(ast.print_to_string({comments: f}), "#!Random comment\n//test1\n/*test2*/");
     });
 
     it("Should never be able to filter comment5 when using 'some' as filter", function() {
         var ast = UglifyJS.parse("#!foo\n//foo\n/*@preserve*/\n/* please hide me */");
-        assert.strictEqual(ast.print_to_string({comments: "some"}), "#!foo\n/*@preserve*/\n");
+        assert.strictEqual(ast.print_to_string({comments: "some"}), "#!foo\n/*@preserve*/");
     });
 
     it("Should have no problem on multiple calls", function() {
diff --git a/test/mocha/comment.js b/test/mocha/comment.js
index 6b5428d4..3bad9e01 100644
--- a/test/mocha/comment.js
+++ b/test/mocha/comment.js
@@ -47,4 +47,176 @@ 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());
+    });
+
+    it("Should retain trailing comments", function() {
+        var code = [
+            "if (foo /* lost comment */ && bar /* lost comment */) {",
+            "    // this one is kept",
+            "    {/* lost comment */}",
+            "    !function() {",
+            "        // lost comment",
+            "    }();",
+            "    function baz() {/* lost comment */}",
+            "    // lost comment",
+            "}",
+            "// comments right before EOF are lost as well",
+        ].join("\n");
+        var result = uglify.minify(code, {
+            compress: false,
+            mangle: false,
+            output: {
+                beautify: true,
+                comments: "all",
+            },
+        });
+        if (result.error) throw result.error;
+        assert.strictEqual(result.code, code);
+    });
+
+    it("Should correctly preserve new lines around comments", function() {
+        var tests = [
+            [
+                "// foo",
+                "// bar",
+                "x();",
+            ].join("\n"),
+            [
+                "// foo",
+                "/* bar */",
+                "x();",
+            ].join("\n"),
+            [
+                "// foo",
+                "/* bar */ x();",
+            ].join("\n"),
+            [
+                "/* foo */",
+                "// bar",
+                "x();",
+            ].join("\n"),
+            [
+                "/* foo */ // bar",
+                "x();",
+            ].join("\n"),
+            [
+                "/* foo */",
+                "/* bar */",
+                "x();",
+            ].join("\n"),
+            [
+                "/* foo */",
+                "/* bar */ x();",
+            ].join("\n"),
+            [
+                "/* foo */ /* bar */",
+                "x();",
+            ].join("\n"),
+            "/* foo */ /* bar */ x();",
+        ].forEach(function(code) {
+            var result = uglify.minify(code, {
+                compress: false,
+                mangle: false,
+                output: {
+                    beautify: true,
+                    comments: "all",
+                },
+            });
+            if (result.error) throw result.error;
+            assert.strictEqual(result.code, code);
+        });
+    });
+
+    it("Should preserve new line before comment without beautify", function() {
+        var code = [
+            "function f(){",
+            "/* foo */bar()}",
+        ].join("\n");
+        var result = uglify.minify(code, {
+            compress: false,
+            mangle: false,
+            output: {
+                comments: "all",
+            },
+        });
+        if (result.error) throw result.error;
+        assert.strictEqual(result.code, code);
+    });
+
+    it("Should preserve comments around IIFE", function() {
+        var result = uglify.minify("/*a*/(/*b*/function(){/*c*/}/*d*/)/*e*/();", {
+            compress: false,
+            mangle: false,
+            output: {
+                comments: "all",
+            },
+        });
+        if (result.error) throw result.error;
+        assert.strictEqual(result.code, "/*a*/ /*b*/(function(){/*c*/}/*d*/ /*e*/)();");
+    });
 });