Make folding if to return/assignment work with if-else if-else #KT-13452 Fixed

idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt

@@ -219,4 +219,10 @@ internal fun KotlinType.isFlexibleRecursive(): Boolean {
     return arguments.any { !it.isStarProjection && it.type.isFlexibleRecursive() }
 }
 
+val KtIfExpression.branches: List<KtExpression?> get() = ifBranchesOrThis()
+
+private fun KtExpression.ifBranchesOrThis(): List<KtExpression?> {
+    if (this !is KtIfExpression) return listOf(this)
+    return listOf(then) + `else`?.ifBranchesOrThis().orEmpty()
+}
 

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

@@ -20,27 +20,38 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
 import org.jetbrains.kotlin.idea.intentions.SelfTargetingRangeIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.BranchedFoldingUtils
+import org.jetbrains.kotlin.idea.intentions.branches
 import org.jetbrains.kotlin.psi.KtIfExpression
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.createExpressionByPattern
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 
 class FoldIfToAssignmentIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpression::class.java, "Replace 'if' expression with assignment") {
     override fun applicabilityRange(element: KtIfExpression): TextRange? {
-        val thenAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(element.then) ?: return null
-        val elseAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(element.`else`) ?: return null
-        if (!BranchedFoldingUtils.checkAssignmentsMatch(thenAssignment, elseAssignment)) return null
+        val branches = element.branches
+
+        if (branches.lastOrNull()?.getStrictParentOfType<KtIfExpression>()?.`else` == null) return null
+        if (branches.size < 2) return null
+
+        for (i in 0..branches.size - 2) {
+            val assignment1 = BranchedFoldingUtils.getFoldableBranchedAssignment(branches[i]) ?: return null
+            val assignment2 = BranchedFoldingUtils.getFoldableBranchedAssignment(branches[i + 1]) ?: return null
+            if (!BranchedFoldingUtils.checkAssignmentsMatch(assignment1, assignment2)) return null
+        }
+
         return element.ifKeyword.textRange
     }
 
     override fun applyTo(element: KtIfExpression, editor: Editor?) {
-        var thenAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(element.then!!)!!
-        val elseAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(element.`else`!!)!!
+        val thenAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(element.then!!)!!
 
         val op = thenAssignment.operationReference.text
         val leftText = thenAssignment.left!!.text
 
-        thenAssignment.replace(thenAssignment.right!!)
-        elseAssignment.replace(elseAssignment.right!!)
+        element.branches.forEach {
+            val assignment = BranchedFoldingUtils.getFoldableBranchedAssignment(it!!)!!
+            assignment.replace(assignment.right!!)
+        }
 
         element.replace(KtPsiFactory(element).createExpressionByPattern("$0 $1 $2", leftText, op, element))
     }

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

@@ -20,23 +20,29 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
 import org.jetbrains.kotlin.idea.intentions.SelfTargetingRangeIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.BranchedFoldingUtils
+import org.jetbrains.kotlin.idea.intentions.branches
 import org.jetbrains.kotlin.psi.KtIfExpression
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.createExpressionByPattern
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 
 class FoldIfToReturnIntention : SelfTargetingRangeIntention<KtIfExpression>(KtIfExpression::class.java, "Replace 'if' expression with return") {
+
     override fun applicabilityRange(element: KtIfExpression): TextRange? {
-        if (BranchedFoldingUtils.getFoldableBranchedReturn(element.then) == null) return null
-        if (BranchedFoldingUtils.getFoldableBranchedReturn(element.`else`) == null) return null
+        val branches = element.branches
+
+        val lastElse = branches.lastOrNull()?.getStrictParentOfType<KtIfExpression>()?.`else` ?: return null
+        if (BranchedFoldingUtils.getFoldableBranchedReturn(lastElse) == null) return null
+
+        if (branches.any { BranchedFoldingUtils.getFoldableBranchedReturn(it) == null }) return null
         return element.ifKeyword.textRange
     }
 
     override fun applyTo(element: KtIfExpression, editor: Editor?) {
-        val thenReturn = BranchedFoldingUtils.getFoldableBranchedReturn(element.then!!)!!
-        val elseReturn = BranchedFoldingUtils.getFoldableBranchedReturn(element.`else`!!)!!
-
-        thenReturn.replace(thenReturn.returnedExpression!!)
-        elseReturn.replace(elseReturn.returnedExpression!!)
+        for (it in element.branches) {
+            val returnExpression = BranchedFoldingUtils.getFoldableBranchedReturn(it)
+            returnExpression!!.replace(returnExpression.returnedExpression!!)
+        }
 
         element.replace(KtPsiFactory(element).createExpressionByPattern("return $0", element))
     }

idea/testData/intentions/branched/folding/ifToAssignment/ifElseIf.kt

@@ -0,0 +1,9 @@
+// IS_APPLICABLE: false
+
+fun test(n: Int) {
+    var a: String = ""
+    <caret>if (n == 1)
+        a = "one"
+    else if (n == 2)
+        a = "two"
+}

idea/testData/intentions/branched/folding/ifToAssignment/ifElseIfElse.kt

@@ -0,0 +1,9 @@
+fun test(n: Int) {
+    val a: String
+    <caret>if (n == 1)
+        a = "one"
+    else if (n == 2)
+        a = "two"
+    else
+        a = "three"
+}

idea/testData/intentions/branched/folding/ifToAssignment/ifElseIfElse.kt.after

@@ -0,0 +1,9 @@
+fun test(n: Int) {
+    val a: String
+    <caret>a = if (n == 1)
+        "one"
+    else if (n == 2)
+        "two"
+    else
+        "three"
+}

idea/testData/intentions/branched/folding/ifToAssignment/ifElseifElseInconsistent.kt

@@ -0,0 +1,12 @@
+// IS_APPLICABLE: false
+
+fun test(n: Int) {
+    val a: String
+    val b: String
+    <caret>if (n == 1)
+        a = "one"
+    else if (n == 2)
+        a = "two"
+    else
+        b = "three"
+}

idea/testData/intentions/branched/folding/ifToReturn/ifElseIf.kt

@@ -0,0 +1,10 @@
+// IS_APPLICABLE: false
+
+fun test(n: Int): String {
+    <caret>if (n == 1)
+        return "one"
+    else if (n == 2)
+        return "two"
+
+    return "three"
+}

idea/testData/intentions/branched/folding/ifToReturn/ifElseIfElse.kt

@@ -0,0 +1,8 @@
+fun test(n: Int): String {
+    <caret>if (n == 1)
+        return "one"
+    else if (n == 2)
+        return "two"
+    else
+        return "three"
+}

idea/testData/intentions/branched/folding/ifToReturn/ifElseIfElse.kt.after

@@ -0,0 +1,8 @@
+fun test(n: Int): String {
+    <caret>return if (n == 1)
+        "one"
+    else if (n == 2)
+        "two"
+    else
+        "three"
+}

idea/testData/intentions/branched/folding/ifToReturn/ifElseIfElseInconsistent.kt

@@ -0,0 +1,12 @@
+// IS_APPLICABLE: false
+
+fun test(n: Int): String {
+    val a: String
+    <caret>if (n == 1)
+        return "one"
+    else if (n == 2)
+        return "two"
+    else
+        a = "three"
+    return a
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -776,6 +776,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/folding/ifToAssignment"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
                 }
 
+                @TestMetadata("ifElseIf.kt")
+                public void testIfElseIf() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToAssignment/ifElseIf.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("ifElseIfElse.kt")
+                public void testIfElseIfElse() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToAssignment/ifElseIfElse.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("ifElseifElseInconsistent.kt")
+                public void testIfElseifElseInconsistent() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToAssignment/ifElseifElseInconsistent.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("innerIfTransformed.kt")
                 public void testInnerIfTransformed() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToAssignment/innerIfTransformed.kt");
@@ -839,6 +857,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/folding/ifToReturn"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
                 }
 
+                @TestMetadata("ifElseIf.kt")
+                public void testIfElseIf() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToReturn/ifElseIf.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("ifElseIfElse.kt")
+                public void testIfElseIfElse() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToReturn/ifElseIfElse.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("ifElseIfElseInconsistent.kt")
+                public void testIfElseIfElseInconsistent() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToReturn/ifElseIfElseInconsistent.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("innerIfTransformed.kt")
                 public void testInnerIfTransformed() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToReturn/innerIfTransformed.kt");