[CS2] add parens to chained do IIFE (#4666)

* add parens to chained do iife [Fixes #3736]

* remove debug code

* fixes from code review

lib/coffeescript/rewriter.js

@@ -86,6 +86,7 @@
         this.normalizeLines();
         this.tagPostfixConditionals();
         this.addImplicitBracesAndParens();
+        this.addParensToChainedDoIife();
         this.rescueStowawayComments();
         this.addLocationDataToGeneratedTokens();
         this.enforceValidCSXAttributes();
@@ -734,6 +735,34 @@
         });
       }
 
+      // Add parens around a `do` IIFE followed by a chained `.` so that the
+      // chaining applies to the executed function rather than the function
+      // object (see #3736)
+      addParensToChainedDoIife() {
+        var action, condition, doIndex;
+        condition = function(token, i) {
+          return this.tag(i - 1) === 'OUTDENT';
+        };
+        action = function(token, i) {
+          var ref;
+          if (ref = token[0], indexOf.call(CALL_CLOSERS, ref) < 0) {
+            return;
+          }
+          this.tokens.splice(doIndex, 0, generate('(', '(', this.tokens[doIndex]));
+          return this.tokens.splice(i + 1, 0, generate(')', ')', this.tokens[i]));
+        };
+        doIndex = null;
+        return this.scanTokens(function(token, i, tokens) {
+          var ref;
+          if (!(token[1] === 'do' && ((ref = this.tag(i + 1)) === '->' || ref === '=>') && this.tag(i + 2) === 'INDENT')) {
+            return 1;
+          }
+          doIndex = i;
+          this.detectEnd(i + 2, condition, action);
+          return 2;
+        });
+      }
+
       // Because our grammar is LALR(1), it can’t handle some single-line
       // expressions that lack ending delimiters. The **Rewriter** adds the implicit
       // blocks, so it doesn’t need to. To keep the grammar clean and tidy, trailing

src/rewriter.coffee

@@ -53,6 +53,7 @@ exports.Rewriter = class Rewriter
     @normalizeLines()
     @tagPostfixConditionals()
     @addImplicitBracesAndParens()
+    @addParensToChainedDoIife()
     @rescueStowawayComments()
     @addLocationDataToGeneratedTokens()
     @enforceValidCSXAttributes()
@@ -512,6 +513,23 @@ exports.Rewriter = class Rewriter
         last_column:  prevLocationData.last_column
       return 1
 
+  # Add parens around a `do` IIFE followed by a chained `.` so that the
+  # chaining applies to the executed function rather than the function
+  # object (see #3736)
+  addParensToChainedDoIife: ->
+    condition = (token, i) ->
+      @tag(i - 1) is 'OUTDENT'
+    action = (token, i) ->
+      return unless token[0] in CALL_CLOSERS
+      @tokens.splice doIndex, 0, generate '(', '(', @tokens[doIndex]
+      @tokens.splice i + 1, 0, generate ')', ')', @tokens[i]
+    doIndex = null
+    @scanTokens (token, i, tokens) ->
+      return 1 unless token[1] is 'do' and @tag(i + 1) in ['->', '=>'] and @tag(i + 2) is 'INDENT'
+      doIndex = i
+      @detectEnd i + 2, condition, action
+      return 2
+
   # Because our grammar is LALR(1), it can’t handle some single-line
   # expressions that lack ending delimiters. The **Rewriter** adds the implicit
   # blocks, so it doesn’t need to. To keep the grammar clean and tidy, trailing

test/formatting.coffee

@@ -420,3 +420,19 @@ test "#4576: function chaining on separate rows", ->
     .then ->
       yes
     .then ok
+
+test "#3736: chaining after do IIFE", ->
+  eq 3,
+    do ->
+      a: 3
+    .a
+
+  eq 3,
+    do -> a: 3
+    ?.a
+
+  # preserve existing chaining behavior for non-IIFE `do`
+  b = c: -> 4
+  eq 4,
+    do b
+    .c