Improve prefer-immediate-return: keep comments when fixing (#249)

Author: vilchik-elena

src/rules/prefer-immediate-return.ts

@@ -73,37 +73,33 @@ const rule: Rule.RuleModule = {
             context.report({
               message: formatMessage(last, returnedIdentifier.name),
               node: declaredIdentifier.init,
-              fix: fixer => fix(fixer, last, lastButOne, declaredIdentifier.init),
+              fix: fixer =>
+                fix(fixer, last, lastButOne, declaredIdentifier.init, returnedIdentifier),
             });
           }
         }
       }
     }
 
+    // eslint-disable-next-line max-params
     function fix(
       fixer: TSESLint.RuleFixer,
       last: TSESTree.Statement,
       lastButOne: TSESTree.Statement,
-      expression: TSESTree.Expression,
+      expressionToReturn: TSESTree.Expression,
+      returnedExpression: TSESTree.Expression,
     ): any {
-      const throwOrReturnKeyword = context.getSourceCode().getFirstToken(last);
-
-      if (lastButOne.range && last.range && throwOrReturnKeyword) {
-        const expressionText = context.getSourceCode().getText(expression);
-        const fixedRangeStart = lastButOne.range[0];
-        const fixedRangeEnd = last.range[1];
-        const semicolonToken = context.getSourceCode().getLastToken(last);
-        const semicolon = semicolonToken && semicolonToken.value === ';' ? ';' : '';
-        return [
-          fixer.removeRange([fixedRangeStart, fixedRangeEnd]),
-          fixer.insertTextAfterRange(
-            [1, fixedRangeStart],
-            `${throwOrReturnKeyword.value} ${expressionText}${semicolon}`,
-          ),
-        ];
-      } else {
-        return null;
-      }
+      const expressionText = context.getSourceCode().getText(expressionToReturn);
+      const rangeToRemoveStart = lastButOne.range[0];
+      const commentsBetweenStatements = context.getSourceCode().getCommentsAfter(lastButOne);
+      const rangeToRemoveEnd =
+        commentsBetweenStatements.length > 0
+          ? commentsBetweenStatements[0].range[0]
+          : last.range[0];
+      return [
+        fixer.removeRange([rangeToRemoveStart, rangeToRemoveEnd]),
+        fixer.replaceText(returnedExpression, expressionText),
+      ];
     }
 
     function getOnlyReturnedVariable(node: TSESTree.Statement) {

tests/rules/prefer-immediate-return.test.ts

@@ -394,5 +394,24 @@ ruleTester.run('prefer-immediate-return', rule, {
       }
       `,
     },
+    {
+      code: `
+        function var_returned() {
+          // comment1
+          var x /* commentInTheMiddle1 */ = 42; // commentOnTheLine1
+          // comment2
+          return /* commentInTheMiddle2 */ x;   // commentOnTheLine2
+          // comment3
+        }`,
+      errors: 1,
+      output: `
+        function var_returned() {
+          // comment1
+          // commentOnTheLine1
+          // comment2
+          return /* commentInTheMiddle2 */ 42;   // commentOnTheLine2
+          // comment3
+        }`,
+    },
   ],
 });