render comments in custom ASTs gracefully (#3393)

fixes #3246

3 files changed, 42 insertions, 16 deletions

diff --git a/lib/output.js b/lib/output.js
index a3b555ff..2d0dc53f 100644
--- a/lib/output.js
+++ b/lib/output.js
@@ -451,16 +451,11 @@ function OutputStream(options) {
 
     function prepend_comments(node) {
         var self = this;
-        var start = node.start;
-        if (!start) return;
-        if (start.comments_before && start.comments_before._dumped === self) return;
-        var comments = start.comments_before;
-        if (!comments) {
-            comments = start.comments_before = [];
-        }
-        comments._dumped = self;
+        var scan = node instanceof AST_Exit && node.value;
+        var comments = dump(node);
+        if (!comments) return;
 
-        if (node instanceof AST_Exit && node.value) {
+        if (scan) {
             var tw = new TreeWalker(function(node) {
                 var parent = tw.parent();
                 if (parent instanceof AST_Exit
@@ -471,11 +466,8 @@ function OutputStream(options) {
                     || 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);
-                    }
+                    var before = dump(node);
+                    if (before) comments = comments.concat(before);
                 } else {
                     return true;
                 }
@@ -518,13 +510,29 @@ function OutputStream(options) {
             }
         });
         if (!last_nlb) {
-            if (start.nlb) {
+            if (node.start.nlb) {
                 print("\n");
                 indent();
             } else {
                 space();
             }
         }
+
+        function dump(node) {
+            var token = node.start;
+            if (!token) {
+                if (!scan) return;
+                node.start = token = new AST_Token();
+            }
+            var comments = token.comments_before;
+            if (!comments) {
+                if (!scan) return;
+                token.comments_before = comments = [];
+            }
+            if (comments._dumped === self) return;
+            comments._dumped = self;
+            return comments;
+        }
     }
 
     function append_comments(node, tail) {
diff --git a/test/mocha/comments.js b/test/mocha/comments.js
index 1ca1432a..3907f2fa 100644
--- a/test/mocha/comments.js
+++ b/test/mocha/comments.js
@@ -260,6 +260,23 @@ describe("comments", function() {
         ].join("\n"));
     });
 
+    it("Should handle programmatic AST insertions gracefully", function() {
+        var ast = UglifyJS.parse([
+            "function f() {",
+            "    //foo",
+            "    bar;",
+            "    return;",
+            "}",
+        ].join("\n"));
+        ast.body[0].body[0] = new UglifyJS.AST_Throw({value: ast.body[0].body[0].body});
+        ast.body[0].body[1].value = new UglifyJS.AST_Number({value: 42});
+        assert.strictEqual(ast.print_to_string({comments: "all"}), [
+            "function f(){",
+            "//foo",
+            "throw bar;return 42}",
+        ].join("\n"));
+    });
+
     describe("comment before constant", function() {
         var js = 'function f() { /*c1*/ var /*c2*/ foo = /*c3*/ false; return foo; }';
 
diff --git a/test/run-tests.js b/test/run-tests.js
index aa415ec9..c1e46078 100755
--- a/test/run-tests.js
+++ b/test/run-tests.js
@@ -219,8 +219,9 @@ function run_compress_tests() {
             var input_code = make_code(input);
             var input_formatted = make_code(test.input, {
                 beautify: true,
+                comments: "all",
+                keep_quoted_props: true,
                 quote_style: 3,
-                keep_quoted_props: true
             });
             try {
                 U.parse(input_code);