Lift return out: report also on 'return' keywords #KT-27173 Fixed

idea/src/org/jetbrains/kotlin/idea/inspections/LiftReturnOrAssignmentInspection.kt

@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.idea.intentions.branchedTransformations.isElseIf
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.lineCount
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
 
@@ -39,25 +39,38 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
                     val hasOtherReturns = expression.anyDescendantOfType<KtReturnExpression> { it !in foldableReturns }
                     val isSerious = !hasOtherReturns && foldableReturns.size > 1
                     val verb = if (isSerious) "should" else "can"
+                    val description = "Return $verb be lifted out of '${keyword.text}'"
                     holder.registerProblemWithoutOfflineInformation(
-                        keyword,
+                        expression,
-                        "Return $verb be lifted out of '${keyword.text}'",
+                        description,
                         isOnTheFly,
                         if (isSerious) ProblemHighlightType.GENERIC_ERROR_OR_WARNING
                         else ProblemHighlightType.INFORMATION,
+                        keyword.textRange?.shiftRight(-expression.startOffset),
                         LiftReturnOutFix(keyword.text)
                     )
+                    foldableReturns.forEach {
+                        holder.registerProblemWithoutOfflineInformation(
+                            expression,
+                            description,
+                            isOnTheFly,
+                            ProblemHighlightType.INFORMATION,
+                            it.returnKeyword.textRange?.shiftRight(-expression.startOffset),
+                            LiftReturnOutFix(keyword.text)
+                        )
+                    }
                     return
                 }
                 val assignmentNumber = BranchedFoldingUtils.getFoldableAssignmentNumber(expression)
                 if (assignmentNumber > 0) {
                     val verb = if (assignmentNumber > 1) "should" else "can"
                     holder.registerProblemWithoutOfflineInformation(
-                        keyword,
+                        expression,
                         "Assignment $verb be lifted out of '${keyword.text}'",
                         isOnTheFly,
                         if (assignmentNumber > 1) ProblemHighlightType.GENERIC_ERROR_OR_WARNING
                         else ProblemHighlightType.INFORMATION,
+                        keyword.textRange?.shiftRight(-expression.startOffset),
                         LiftAssignmentOutFix(keyword.text)
                     )
                 }
@@ -87,7 +100,8 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
         override fun getFamilyName() = name
 
         override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
-            BranchedFoldingUtils.foldToReturn(descriptor.psiElement.getParentOfType(true)!!)
+            val replaced = BranchedFoldingUtils.foldToReturn(descriptor.psiElement as KtExpression)
+            replaced.findExistingEditor()?.caretModel?.moveToOffset(replaced.startOffset)
         }
     }
 
@@ -97,7 +111,7 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
         override fun getFamilyName() = name
 
         override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
-            BranchedFoldingUtils.foldToAssignment(descriptor.psiElement.getParentOfType(true)!!)
+            BranchedFoldingUtils.foldToAssignment(descriptor.psiElement as KtExpression)
         }
     }
 

idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 import org.jetbrains.kotlin.cfg.WhenChecker
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.intentions.branches
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
@@ -206,7 +207,7 @@ object BranchedFoldingUtils {
         expression.replace(psiFactory.createExpressionByPattern("$0 $1 $2", lhs!!, op!!, expression))
     }
 
-    fun foldToReturn(expression: KtExpression) {
+    fun foldToReturn(expression: KtExpression): KtExpression {
         fun KtReturnExpression.replaceWithReturned() {
             replace(returnedExpression!!)
         }
@@ -226,7 +227,7 @@ object BranchedFoldingUtils {
             }
         }
         lift(expression)
-        expression.replace(KtPsiFactory(expression).createExpressionByPattern("return $0", expression))
+        return expression.replaced(KtPsiFactory(expression).createExpressionByPattern("return $0", expression))
     }
 
     private fun KtTryExpression.tryBlockAndCatchBodies(): List<KtExpression?> = listOf(tryBlock) + catchClauses.map { it.catchBody }

idea/testData/inspectionsLocal/liftOut/ifToReturn/onReturn.kt

@@ -0,0 +1,7 @@
+// HIGHLIGHT: INFORMATION
+fun test(n: Int): String {
+    if (n == 1)
+        <caret>return "one"
+    else
+        return "two"
+}

idea/testData/inspectionsLocal/liftOut/ifToReturn/onReturn.kt.after

@@ -0,0 +1,7 @@
+// HIGHLIGHT: INFORMATION
+fun test(n: Int): String {
+    <caret>return if (n == 1)
+        "one"
+    else
+        "two"
+}

idea/testData/inspectionsLocal/liftOut/ifToReturn/onReturn2.kt

@@ -0,0 +1,7 @@
+// HIGHLIGHT: INFORMATION
+fun test(n: Int): String {
+    if (n == 1)
+        return "one"
+    else
+        <caret>return "two"
+}

idea/testData/inspectionsLocal/liftOut/ifToReturn/onReturn2.kt.after

@@ -0,0 +1,7 @@
+// HIGHLIGHT: INFORMATION
+fun test(n: Int): String {
+    <caret>return if (n == 1)
+        "one"
+    else
+        "two"
+}

idea/testData/inspectionsLocal/liftOut/tryToReturn/onReturn.kt

@@ -0,0 +1,8 @@
+// HIGHLIGHT: INFORMATION
+fun test(): String {
+    try {
+        <caret>return "success"
+    } catch (e: Exception) {
+        throw e
+    }
+}

idea/testData/inspectionsLocal/liftOut/tryToReturn/onReturn.kt.after

@@ -0,0 +1,8 @@
+// HIGHLIGHT: INFORMATION
+fun test(): String {
+    <caret>return try {
+        "success"
+    } catch (e: Exception) {
+        throw e
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToReturn/onReturn.kt

@@ -0,0 +1,7 @@
+// HIGHLIGHT: INFORMATION
+fun test(n: Int): String {
+    when (n) {
+        1 -> <caret>return "one"
+        else -> return "two"
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToReturn/onReturn.kt.after

@@ -0,0 +1,7 @@
+// HIGHLIGHT: INFORMATION
+fun test(n: Int): String {
+    <caret>return when (n) {
+        1 -> "one"
+        else -> "two"
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToReturn/onReturn2.kt

@@ -0,0 +1,7 @@
+// HIGHLIGHT: INFORMATION
+fun test(n: Int): String {
+    when (n) {
+        1 -> return "one"
+        else -> <caret>return "two"
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToReturn/onReturn2.kt.after

@@ -0,0 +1,7 @@
+// HIGHLIGHT: INFORMATION
+fun test(n: Int): String {
+    <caret>return when (n) {
+        1 -> "one"
+        else -> "two"
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -2865,6 +2865,16 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 runTest("idea/testData/inspectionsLocal/liftOut/ifToReturn/innerIfTransformed.kt");
             }
 
+            @TestMetadata("onReturn.kt")
+            public void testOnReturn() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToReturn/onReturn.kt");
+            }
+
+            @TestMetadata("onReturn2.kt")
+            public void testOnReturn2() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToReturn/onReturn2.kt");
+            }
+
             @TestMetadata("simpleIf.kt")
             public void testSimpleIf() throws Exception {
                 runTest("idea/testData/inspectionsLocal/liftOut/ifToReturn/simpleIf.kt");
@@ -2996,6 +3006,11 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 runTest("idea/testData/inspectionsLocal/liftOut/tryToReturn/inner.kt");
             }
 
+            @TestMetadata("onReturn.kt")
+            public void testOnReturn() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/tryToReturn/onReturn.kt");
+            }
+
             @TestMetadata("withoutReturn.kt")
             public void testWithoutReturn() throws Exception {
                 runTest("idea/testData/inspectionsLocal/liftOut/tryToReturn/withoutReturn.kt");
@@ -3112,6 +3127,16 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 runTest("idea/testData/inspectionsLocal/liftOut/whenToReturn/localReturns.kt");
             }
 
+            @TestMetadata("onReturn.kt")
+            public void testOnReturn() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/whenToReturn/onReturn.kt");
+            }
+
+            @TestMetadata("onReturn2.kt")
+            public void testOnReturn2() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/whenToReturn/onReturn2.kt");
+            }
+
             @TestMetadata("otherReturns.kt")
             public void testOtherReturns() throws Exception {
                 runTest("idea/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt");