render comments in custom ASTs gracefully (#3393)

fixes #3246

1 files changed, 23 insertions, 15 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) {