preserve trailing comments

fixes #88

lib/ast.js

@@ -87,7 +87,7 @@ function DEFNODE(type, props, methods, base) {
     return ctor;
 };
 
-var AST_Token = DEFNODE("Token", "type value line col pos endline endcol endpos nlb comments_before file raw", {
+var AST_Token = DEFNODE("Token", "type value line col pos endline endcol endpos nlb comments_before comments_after file raw", {
 }, null);
 
 var AST_Node = DEFNODE("Node", "start end", {

lib/output.js

@@ -200,6 +200,7 @@ function OutputStream(options) {
     var might_need_space = false;
     var might_need_semicolon = false;
     var might_add_newline = 0;
+    var need_newline_indented = false;
     var last = "";
     var mapping_token, mapping_name, mappings = options.source_map && [];
 
@@ -258,6 +259,13 @@ function OutputStream(options) {
     function print(str) {
         str = String(str);
         var ch = str.charAt(0);
+        if (need_newline_indented && ch) {
+            need_newline_indented = false;
+            if (ch != "\n") {
+                print("\n");
+                indent();
+            }
+        }
         var prev = last.charAt(last.length - 1);
         if (might_need_semicolon) {
             might_need_semicolon = false;
@@ -428,7 +436,7 @@ function OutputStream(options) {
         return OUTPUT;
     };
 
-    function add_comments(node) {
+    function prepend_comments(node) {
         var self = this;
         var start = node.start;
         if (!(start.comments_before && start.comments_before._dumped === self)) {
@@ -474,23 +482,28 @@ function OutputStream(options) {
             }
 
             comments = comments.filter(comment_filter, node);
-
+            if (comments.length == 0) return;
-            // Keep single line comments after nlb, after nlb
+            var last_nlb = /(^|\n) *$/.test(OUTPUT);
-            if (current_col != 0
+            comments.forEach(function(c, i) {
-                && !options.beautify
+                if (!last_nlb) {
-                && comments.length > 0
+                    if (c.nlb) {
-                && comments[0].nlb
-                && /comment[134]/.test(comments[0].type)) {
                         print("\n");
+                        indent();
+                        last_nlb = true;
+                    } else if (i > 0) {
+                        space();
+                    }
                 }
-
-            comments.forEach(function(c){
                 if (/comment[134]/.test(c.type)) {
                     print("//" + c.value + "\n");
                     indent();
-                }
+                    last_nlb = true;
-                else if (c.type == "comment2") {
+                } else if (c.type == "comment2") {
                     print("/*" + c.value + "*/");
+                    last_nlb = false;
+                }
+            });
+            if (!last_nlb) {
                 if (start.nlb) {
                     print("\n");
                     indent();
@@ -498,6 +511,30 @@ function OutputStream(options) {
                     space();
                 }
             }
+        }
+    }
+
+    function append_comments(node, tail) {
+        var self = this;
+        var token = node.end;
+        if (!token) return;
+        var comments = token[tail ? "comments_before" : "comments_after"];
+        if (comments && comments._dumped !== self) {
+            comments._dumped = self;
+            comments.filter(comment_filter, node).forEach(function(c, i) {
+                if (need_newline_indented || c.nlb) {
+                    print("\n");
+                    indent();
+                    need_newline_indented = false;
+                } else if (i > 0 || !tail) {
+                    space();
+                }
+                if (/comment[134]/.test(c.type)) {
+                    print("//" + c.value);
+                    need_newline_indented = true;
+                } else if (c.type == "comment2") {
+                    print("/*" + c.value + "*/");
+                }
             });
         }
     }
@@ -539,7 +576,8 @@ function OutputStream(options) {
         with_square     : with_square,
         add_mapping     : add_mapping,
         option          : function(opt) { return options[opt] },
-        add_comments    : readonly ? noop : add_comments,
+        prepend_comments: readonly ? noop : prepend_comments,
+        append_comments : readonly ? noop : append_comments,
         line            : function() { return current_line },
         col             : function() { return current_col },
         pos             : function() { return current_pos },
@@ -575,9 +613,10 @@ function OutputStream(options) {
             use_asm = active_scope;
         }
         function doit() {
-            stream.add_comments(self);
+            stream.prepend_comments(self);
             self.add_source_map(stream);
             generator(self, stream);
+            stream.append_comments(self);
         }
         stream.push_node(self);
         if (force_parens || self.needs_parens(stream)) {
@@ -819,14 +858,21 @@ function OutputStream(options) {
         self.body.print(output);
         output.semicolon();
     });
-    function print_bracketed(body, output, allow_directives) {
+    function print_bracketed(self, output, allow_directives) {
-        if (body.length > 0) output.with_block(function(){
+        if (self.body.length > 0) {
-            display_body(body, false, output, allow_directives);
+            output.with_block(function() {
+                display_body(self.body, false, output, allow_directives);
             });
-        else output.print("{}");
+        } else {
+            output.print("{");
+            output.with_indent(output.next_indent(), function() {
+                output.append_comments(self, true);
+            });
+            output.print("}");
+        }
     };
     DEFPRINT(AST_BlockStatement, function(self, output){
-        print_bracketed(self.body, output);
+        print_bracketed(self, output);
     });
     DEFPRINT(AST_EmptyStatement, function(self, output){
         output.semicolon();
@@ -921,7 +967,7 @@ function OutputStream(options) {
             });
         });
         output.space();
-        print_bracketed(self.body, output, true);
+        print_bracketed(self, output, true);
     });
     DEFPRINT(AST_Lambda, function(self, output){
         self._do_print(output);
@@ -1052,7 +1098,7 @@ function OutputStream(options) {
     DEFPRINT(AST_Try, function(self, output){
         output.print("try");
         output.space();
-        print_bracketed(self.body, output);
+        print_bracketed(self, output);
         if (self.bcatch) {
             output.space();
             self.bcatch.print(output);
@@ -1069,12 +1115,12 @@ function OutputStream(options) {
             self.argname.print(output);
         });
         output.space();
-        print_bracketed(self.body, output);
+        print_bracketed(self, output);
     });
     DEFPRINT(AST_Finally, function(self, output){
         output.print("finally");
         output.space();
-        print_bracketed(self.body, output);
+        print_bracketed(self, output);
     });
 
     /* -----[ var/const ]----- */

lib/parse.js

@@ -317,11 +317,7 @@ function tokenizer($TEXT, filename, html5_comments, shebang) {
         }
         if (!is_comment) {
             ret.comments_before = S.comments_before;
-            S.comments_before = [];
+            ret.comments_after = S.comments_before = [];
-            // make note of any newlines in the comments that came before
-            for (var i = 0, len = ret.comments_before.length; i < len; i++) {
-                ret.nlb = ret.nlb || ret.comments_before[i].nlb;
-            }
         }
         S.newline_before = false;
         return new AST_Token(ret);

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() {

test/mocha/comment.js

@@ -113,4 +113,98 @@ describe("Comment", function() {
         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);
+    });
 });