Lift return out: do not report for zero returns

Also, do not highlight (but suggest fix) for exactly one return
Related to KT-14900

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

@@ -31,15 +31,18 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean, session: LocalInspectionToolSession) =
             object : KtVisitorVoid() {
                 private fun visitIfOrWhen(expression: KtExpression, keyword: PsiElement) {
-                    if (BranchedFoldingUtils.canFoldToReturn(expression)) {
+                    val returnNumber = BranchedFoldingUtils.getFoldableReturnNumber(expression)
+                    if (returnNumber > 0) {
                         holder.registerProblem(
                                 keyword,
                                 "Return can be lifted out of '${keyword.text}'",
-                                ProblemHighlightType.GENERIC_ERROR_OR_WARNING,
+                                if (returnNumber > 1) ProblemHighlightType.GENERIC_ERROR_OR_WARNING
+                                else ProblemHighlightType.INFORMATION,
                                 LiftReturnOutFix(keyword.text)
                         )
+                        return
                     }
-                    else if (BranchedFoldingUtils.canFoldToAssignment(expression)) {
+                    if (BranchedFoldingUtils.canFoldToAssignment(expression)) {
                         holder.registerProblem(
                                 keyword,
                                 "Assignment can be lifted out of '${keyword.text}'",

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

@@ -82,30 +82,45 @@ object BranchedFoldingUtils {
         return assignments.all { BranchedFoldingUtils.checkAssignmentsMatch(it, firstAssignment) }
     }
 
-    fun canFoldToReturn(expression: KtExpression?): Boolean = when (expression) {
+    private fun getFoldableReturnNumber(branches: List<KtExpression?>) =
+            branches.fold(0) { prevNumber, branch ->
+                when {
+                    prevNumber == -1 -> -1
+                    getFoldableBranchedReturn(branch) != null -> prevNumber + 1
+                    else -> {
+                        val currNumber = getFoldableReturnNumber(branch?.lastBlockStatementOrThis())
+                        if (currNumber == -1) -1 else prevNumber + currNumber
+                    }
+                }
+            }
+
+    internal fun getFoldableReturnNumber(expression: KtExpression?): Int = when (expression) {
         is KtWhenExpression -> {
             val entries = expression.entries
-            KtPsiUtil.checkWhenExpressionHasSingleElse(expression) &&
-            entries.isNotEmpty() &&
-            entries.all { entry ->
-                getFoldableBranchedReturn(entry.expression) != null || canFoldToReturn(entry.expression?.lastBlockStatementOrThis())
+            when {
+                !KtPsiUtil.checkWhenExpressionHasSingleElse(expression) -> -1
+                entries.isEmpty() -> -1
+                else -> getFoldableReturnNumber(entries.map { it.expression })
             }
         }
         is KtIfExpression -> {
             val branches = expression.branches
-            branches.isNotEmpty() &&
-            (branches.lastOrNull()?.getStrictParentOfType<KtIfExpression>()?.`else` != null) &&
-            branches.all { branch ->
-                getFoldableBranchedReturn(branch) != null || canFoldToReturn(branch?.lastBlockStatementOrThis())
+            when {
+                branches.isEmpty() -> -1
+                branches.lastOrNull()?.getStrictParentOfType<KtIfExpression>()?.`else` == null -> -1
+                else -> getFoldableReturnNumber(branches)
             }
         }
         is KtCallExpression -> {
-            expression.analyze().getType(expression)?.isNothing() ?: false
+            if (expression.analyze().getType(expression)?.isNothing() == true) 0 else -1
         }
-        is KtBreakExpression, is KtContinueExpression, is KtThrowExpression -> true
-        else -> false
+        is KtBreakExpression, is KtContinueExpression, is KtThrowExpression -> 0
+        else -> -1
     }
 
+    fun canFoldToReturn(expression: KtExpression?): Boolean =
+            getFoldableReturnNumber(expression) > 0
+
     fun foldToAssignment(expression: KtExpression) {
         var lhs: KtExpression? = null
         var op: String? = null

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

@@ -1,3 +1,5 @@
+// HIGHLIGHT: GENERIC_ERROR_OR_WARNING
+
 fun foo(): Int {
     loop@ while (true) {
         <caret>when (1) {

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

@@ -1,3 +1,5 @@
+// HIGHLIGHT: GENERIC_ERROR_OR_WARNING
+
 fun foo(): Int {
     loop@ while (true) {
         return when (1) {

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

@@ -0,0 +1,8 @@
+// HIGHLIGHT: INFORMATION
+
+fun foo(arg: Int): Int {
+    <caret>when (arg) {
+        0 -> return 0
+        else -> throw Exception()
+    }
+}

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

@@ -0,0 +1,8 @@
+// HIGHLIGHT: INFORMATION
+
+fun foo(arg: Int): Int {
+    return when (arg) {
+        0 -> 0
+        else -> throw Exception()
+    }
+}

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

@@ -0,0 +1,7 @@
+// PROBLEM: none
+
+fun foo(): Int {
+    <caret>when {
+        else -> throw Exception()
+    }
+}

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

@@ -636,6 +636,18 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/simpleWhenWithBlocks.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("whenOneReturn.kt")
+            public void testWhenOneReturn() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/whenOneReturn.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("whenThrowOnly.kt")
+            public void testWhenThrowOnly() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/whenThrowOnly.kt");
+                doTest(fileName);
+            }
         }
     }