"Lift return / assignment": don't highlight if other returns available

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

@@ -22,10 +22,8 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.BranchedFoldingUtils
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.isIfBranch
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.lineCount
-import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtIfExpression
+import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
-import org.jetbrains.kotlin.psi.KtVisitorVoid
-import org.jetbrains.kotlin.psi.KtWhenExpression
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 
 class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
@@ -36,12 +34,13 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
                     if (expression.lineCount() > LINES_LIMIT) return
                     if (expression.isIfBranch()) return
 
-                    val returnNumber = BranchedFoldingUtils.getFoldableReturnNumber(expression)
+                    val foldableReturns = BranchedFoldingUtils.getFoldableReturns(expression)
-                    if (returnNumber > 0) {
+                    if (foldableReturns?.isNotEmpty() == true) {
+                        val hasOtherReturns = expression.anyDescendantOfType<KtReturnExpression> { it !in foldableReturns }
                         holder.registerProblem(
                                 keyword,
                                 "Return can be lifted out of '${keyword.text}'",
-                                if (returnNumber > 1) ProblemHighlightType.GENERIC_ERROR_OR_WARNING
+                                if (!hasOtherReturns && foldableReturns.size > 1) ProblemHighlightType.GENERIC_ERROR_OR_WARNING
                                 else ProblemHighlightType.INFORMATION,
                                 LiftReturnOutFix(keyword.text)
                         )

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

@@ -95,44 +95,47 @@ object BranchedFoldingUtils {
         return assignments.size
     }
 
-    private fun getFoldableReturnNumber(branches: List<KtExpression?>) =
+    private fun getFoldableReturns(branches: List<KtExpression?>): List<KtReturnExpression>? =
-            branches.fold(0) { prevNumber, branch ->
+            branches.fold<KtExpression?, MutableList<KtReturnExpression>?>(mutableListOf()) { prevList, branch ->
-                when {
+                if (prevList == null) return@fold null
-                    prevNumber == -1 -> -1
+                val foldableBranchedReturn = getFoldableBranchedReturn(branch)
-                    getFoldableBranchedReturn(branch) != null -> prevNumber + 1
+                if (foldableBranchedReturn != null) {
-                    else -> {
+                    prevList.add(foldableBranchedReturn)
-                        val currNumber = getFoldableReturnNumber(branch?.lastBlockStatementOrThis())
-                        if (currNumber == -1) -1 else prevNumber + currNumber
-                    }
                 }
+                else {
+                    val currReturns = getFoldableReturns(branch?.lastBlockStatementOrThis()) ?: return@fold null
+                    prevList += currReturns
+                }
+                prevList
             }
 
-    internal fun getFoldableReturnNumber(expression: KtExpression?): Int = when (expression) {
+    internal fun getFoldableReturns(expression: KtExpression?): List<KtReturnExpression>? = when (expression) {
         is KtWhenExpression -> {
             val entries = expression.entries
             when {
-                !KtPsiUtil.checkWhenExpressionHasSingleElse(expression) -> -1
+                !KtPsiUtil.checkWhenExpressionHasSingleElse(expression) -> null
-                entries.isEmpty() -> -1
+                entries.isEmpty() -> null
-                else -> getFoldableReturnNumber(entries.map { it.expression })
+                else -> getFoldableReturns(entries.map { it.expression })
             }
         }
         is KtIfExpression -> {
             val branches = expression.branches
             when {
-                branches.isEmpty() -> -1
+                branches.isEmpty() -> null
-                branches.lastOrNull()?.getStrictParentOfType<KtIfExpression>()?.`else` == null -> -1
+                branches.lastOrNull()?.getStrictParentOfType<KtIfExpression>()?.`else` == null -> null
-                else -> getFoldableReturnNumber(branches)
+                else -> getFoldableReturns(branches)
             }
         }
         is KtCallExpression -> {
-            if (expression.analyze().getType(expression)?.isNothing() == true) 0 else -1
+            if (expression.analyze().getType(expression)?.isNothing() == true) emptyList() else null
         }
-        is KtBreakExpression, is KtContinueExpression, is KtThrowExpression -> 0
+        is KtBreakExpression, is KtContinueExpression, is KtThrowExpression -> emptyList()
-        else -> -1
+        else -> null
     }
 
-    fun canFoldToReturn(expression: KtExpression?): Boolean =
+    private fun getFoldableReturnNumber(expression: KtExpression?) = getFoldableReturns(expression)?.size ?: -1
-            getFoldableReturnNumber(expression) > 0
+
+    fun canFoldToReturn(expression: KtExpression?): Boolean = getFoldableReturnNumber(expression) > 0
 
     fun foldToAssignment(expression: KtExpression) {
         var lhs: KtExpression? = null

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

@@ -0,0 +1,13 @@
+// HIGHLIGHT: INFORMATION
+
+fun test(n: Int, arg: String?): String {
+    <caret>when (n) {
+        1 -> {
+            if (arg == null) return ""
+            return "** $arg"
+        }
+        else -> {
+            return "Strange"
+        }
+    }
+}

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

@@ -0,0 +1,13 @@
+// HIGHLIGHT: INFORMATION
+
+fun test(n: Int, arg: String?): String {
+    return when (n) {
+        1 -> {
+            if (arg == null) return ""
+            "** $arg"
+        }
+        else -> {
+            "Strange"
+        }
+    }
+}

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

@@ -643,6 +643,12 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("otherReturns.kt")
+            public void testOtherReturns() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("simpleWhen.kt")
             public void testSimpleWhen() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/simpleWhen.kt");