Support cascade if / when in lift return / assignment intentions

So #KT-13458 Fixed
So #KT-13436 Fixed

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

@@ -16,9 +16,13 @@
 
 package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.intentions.branches
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
+import org.jetbrains.kotlin.types.typeUtil.isNothing
 
 object BranchedFoldingUtils {
     fun getFoldableBranchedAssignment(branch: KtExpression?): KtBinaryExpression? {
@@ -45,4 +49,105 @@ object BranchedFoldingUtils {
     fun checkAssignmentsMatch(a1: KtBinaryExpression, a2: KtBinaryExpression): Boolean {
         return a1.left?.text == a2.left?.text && a1.operationToken == a2.operationToken
     }
+
+    fun canFoldToAssignment(expression: KtExpression?): Boolean {
+        val assignments = mutableListOf<KtBinaryExpression>()
+        fun collectAssignmentsAndCheck(e: KtExpression?): Boolean {
+            return when (e) {
+                is KtWhenExpression -> {
+                    val entries = e.entries
+                    KtPsiUtil.checkWhenExpressionHasSingleElse(e) &&
+                    entries.isNotEmpty() &&
+                    entries.all { entry ->
+                        val assignment = getFoldableBranchedAssignment(entry.expression)?.run { assignments.add(this) }
+                        assignment != null || collectAssignmentsAndCheck(entry.expression?.lastBlockStatementOrThis())
+                    }
+                }
+                is KtIfExpression -> {
+                    val branches = e.branches
+                    branches.size > 1 &&
+                    (branches.lastOrNull()?.getStrictParentOfType<KtIfExpression>()?.`else` != null) &&
+                    branches.all { branch ->
+                        val assignment = getFoldableBranchedAssignment(branch)?.run { assignments.add(this) }
+                        assignment != null || collectAssignmentsAndCheck(branch?.lastBlockStatementOrThis())
+                    }
+                }
+                is KtCallExpression -> {
+                    e.analyze().getType(e)?.isNothing() ?: false
+                }
+                else -> false
+            }
+        }
+        if (!collectAssignmentsAndCheck(expression)) return false
+        if (assignments.isEmpty()) return false
+        val firstAssignment = assignments.first()
+        return assignments.all { BranchedFoldingUtils.checkAssignmentsMatch(it, firstAssignment) }
+    }
+
+    fun canFoldToReturn(expression: KtExpression?): Boolean = when (expression) {
+        is KtWhenExpression -> {
+            val entries = expression.entries
+            KtPsiUtil.checkWhenExpressionHasSingleElse(expression) &&
+            entries.isNotEmpty() &&
+            entries.all { entry ->
+                getFoldableBranchedReturn(entry.expression) != null || canFoldToReturn(entry.expression?.lastBlockStatementOrThis())
+            }
+        }
+        is KtIfExpression -> {
+            val branches = expression.branches
+            branches.isNotEmpty() &&
+            (branches.lastOrNull()?.getStrictParentOfType<KtIfExpression>()?.`else` != null) &&
+            branches.all { branch ->
+                getFoldableBranchedReturn(branch) != null || canFoldToReturn(branch?.lastBlockStatementOrThis())
+            }
+        }
+        is KtCallExpression -> {
+            expression.analyze().getType(expression)?.isNothing() ?: false
+        }
+        else -> false
+    }
+
+    fun foldToAssignment(expression: KtExpression) {
+        var lhs: KtExpression? = null
+        var op: String? = null
+        fun KtBinaryExpression.replaceWithRHS() {
+            if (lhs == null || op == null) {
+                lhs = left!!.copy() as KtExpression
+                op = operationReference.text
+            }
+            replace(right!!)
+        }
+        fun lift(e: KtExpression?) {
+            when (e) {
+                is KtWhenExpression -> e.entries.forEach { entry ->
+                    getFoldableBranchedAssignment(entry.expression)?.replaceWithRHS() ?: lift(entry.expression?.lastBlockStatementOrThis())
+                }
+                is KtIfExpression -> e.branches.forEach { branch ->
+                    getFoldableBranchedAssignment(branch)?.replaceWithRHS() ?: lift(branch?.lastBlockStatementOrThis())
+                }
+            }
+        }
+        lift(expression)
+        expression.replace(KtPsiFactory(expression).createExpressionByPattern("$0 $1 $2", lhs!!, op!!, expression))
+    }
+
+    fun foldToReturn(expression: KtExpression) {
+        fun KtReturnExpression.replaceWithReturned() {
+            replace(returnedExpression!!)
+        }
+        fun lift(e: KtExpression?) {
+            when (e) {
+                is KtWhenExpression -> e.entries.forEach { entry ->
+                    getFoldableBranchedReturn(entry.expression)?.replaceWithReturned()
+                    ?: lift(entry.expression?.lastBlockStatementOrThis())
+                }
+                is KtIfExpression -> e.branches.forEach { branch ->
+                    getFoldableBranchedReturn(branch)?.replaceWithReturned() ?:
+                    lift(branch?.lastBlockStatementOrThis())
+                }
+            }
+        }
+        lift(expression)
+        expression.replace(KtPsiFactory(expression).createExpressionByPattern("return $0", expression))
+    }
 }

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

@@ -20,42 +20,17 @@ 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,
         "Lift assignment out of 'if' expression"
 ) {
     override fun applicabilityRange(element: KtIfExpression): TextRange? {
-        val branches = element.branches
+        return if (BranchedFoldingUtils.canFoldToAssignment(element)) element.ifKeyword.textRange else null
-
-        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?) {
-        val thenAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(element.then!!)!!
+        BranchedFoldingUtils.foldToAssignment(element)
-
-        val op = thenAssignment.operationReference.text
-        val leftText = thenAssignment.left!!.text
-
-        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,11 +20,7 @@ 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,
@@ -32,21 +28,10 @@ class FoldIfToReturnIntention : SelfTargetingRangeIntention<KtIfExpression>(
 ) {
 
     override fun applicabilityRange(element: KtIfExpression): TextRange? {
-        val branches = element.branches
+        return if (BranchedFoldingUtils.canFoldToReturn(element)) element.ifKeyword.textRange else null
-
-        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?) {
-        for (it in element.branches) {
+        BranchedFoldingUtils.foldToReturn(element)
-            val returnExpression = BranchedFoldingUtils.getFoldableBranchedReturn(it)
-            returnExpression!!.replace(returnExpression.returnedExpression!!)
-        }
-
-        element.replace(KtPsiFactory(element).createExpressionByPattern("return $0", element))
     }
 }

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

@@ -20,53 +20,17 @@ 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.psi.*
+import org.jetbrains.kotlin.psi.KtWhenExpression
-import java.util.*
 
 class FoldWhenToAssignmentIntention : SelfTargetingRangeIntention<KtWhenExpression>(
         KtWhenExpression::class.java,
         "Lift assignment out of 'when' expression"
 ) {
     override fun applicabilityRange(element: KtWhenExpression): TextRange? {
-        if (!KtPsiUtil.checkWhenExpressionHasSingleElse(element)) return null
+        return if (BranchedFoldingUtils.canFoldToAssignment(element)) element.whenKeyword.textRange else null
-
-        val entries = element.entries
-
-        if (entries.isEmpty()) return null
-
-        val assignments = ArrayList<KtBinaryExpression>()
-        for (entry in entries) {
-            val assignment = BranchedFoldingUtils.getFoldableBranchedAssignment(entry.expression) ?: return null
-            assignments.add(assignment)
-        }
-
-        assert(!assignments.isEmpty())
-
-        val firstAssignment = assignments.first()
-        for (assignment in assignments) {
-            if (!BranchedFoldingUtils.checkAssignmentsMatch(assignment, firstAssignment)) return null
-        }
-
-        return element.whenKeyword.textRange
     }
 
     override fun applyTo(element: KtWhenExpression, editor: Editor?) {
-        assert(!element.entries.isEmpty())
+        BranchedFoldingUtils.foldToAssignment(element)
-
-        val firstAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(element.entries.first().expression!!)!!
-
-        val op = firstAssignment.operationReference.text
-        val lhs = firstAssignment.left as KtNameReferenceExpression
-
-        val assignment = KtPsiFactory(element).createExpressionByPattern("$0 $1 $2", lhs, op, element)
-        val newWhenExpression = (assignment as KtBinaryExpression).right as KtWhenExpression
-
-        for (entry in newWhenExpression.entries) {
-            val currAssignment = BranchedFoldingUtils.getFoldableBranchedAssignment(entry.expression!!)!!
-            val currRhs = currAssignment.right!!
-            currAssignment.replace(currRhs)
-        }
-
-        element.replace(assignment)
     }
 }

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

@@ -20,38 +20,18 @@ 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.psi.*
+import org.jetbrains.kotlin.psi.KtWhenExpression
 
 class FoldWhenToReturnIntention : SelfTargetingRangeIntention<KtWhenExpression>(
         KtWhenExpression::class.java,
         "Lift return out of 'when' expression"
 ) {
+
     override fun applicabilityRange(element: KtWhenExpression): TextRange? {
-        if (!KtPsiUtil.checkWhenExpressionHasSingleElse(element)) return null
+        return if (BranchedFoldingUtils.canFoldToReturn(element)) element.whenKeyword.textRange else null
-
-        val entries = element.entries
-
-        if (entries.isEmpty()) return null
-
-        for (entry in entries) {
-            if (BranchedFoldingUtils.getFoldableBranchedReturn(entry.expression) == null) return null
-        }
-
-        return element.whenKeyword.textRange
     }
 
     override fun applyTo(element: KtWhenExpression, editor: Editor?) {
-        assert(!element.entries.isEmpty())
+        BranchedFoldingUtils.foldToReturn(element)
-
-        val newReturnExpression = KtPsiFactory(element).createExpressionByPattern("return $0", element) as KtReturnExpression
-        val newWhenExpression = newReturnExpression.returnedExpression as KtWhenExpression
-
-        for (entry in newWhenExpression.entries) {
-            val currReturn = BranchedFoldingUtils.getFoldableBranchedReturn(entry.expression!!)!!
-            val currExpr = currReturn.returnedExpression!!
-            currReturn.replace(currExpr)
-        }
-
-        element.replace(newReturnExpression)
     }
 }

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

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+fun test(x: Any) {
+    var res: String
+    <caret>if (x is String)
+        when {
+            x.length > 3 -> res = "long string"
+            else -> res = "short string"
+        }
+    else if (x is Int)
+        when {
+            x > 999 || x < -99 -> res = "long int"
+            else -> res = "short int"
+        }
+    else if (x is Long)
+        TODO()
+    else
+        res = "I don't know"
+}

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

@@ -0,0 +1,18 @@
+// WITH_RUNTIME
+fun test(x: Any) {
+    var res: String
+    <caret>res = if (x is String)
+        when {
+            x.length > 3 -> "long string"
+            else -> "short string"
+        }
+    else if (x is Int)
+        when {
+            x > 999 || x < -99 -> "long int"
+            else -> "short int"
+        }
+    else if (x is Long)
+        TODO()
+    else
+        "I don't know"
+}

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

@@ -0,0 +1,17 @@
+// WITH_RUNTIME
+fun test(x: Any): String {
+    <caret>if (x is String)
+        when {
+            x.length > 3 -> return "long string"
+            else -> return "short string"
+        }
+    else if (x is Int)
+        when {
+            x > 999 || x < -99 -> return "long int"
+            else -> return "short int"
+        }
+    else if (x is Long)
+        TODO()
+    else
+        return "I don't know"
+}

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

@@ -0,0 +1,17 @@
+// WITH_RUNTIME
+fun test(x: Any): String {
+    <caret>return if (x is String)
+        when {
+            x.length > 3 -> "long string"
+            else -> "short string"
+        }
+    else if (x is Int)
+        when {
+            x > 999 || x < -99 -> "long int"
+            else -> "short int"
+        }
+    else if (x is Long)
+        TODO()
+    else
+        "I don't know"
+}

idea/testData/intentions/branched/folding/whenToAssignment/cascadeWhen.kt

@@ -0,0 +1,16 @@
+// WITH_RUNTIME
+fun test(x: Any) {
+    var res: String
+    <caret>when (x) {
+        is String ->
+            if (x.length > 3) res = "long string"
+            else res = "short string"
+        is Int ->
+            if (x > 999 || x < -99) res = "long int"
+            else res = "short int"
+        is Long ->
+            TODO()
+        else ->
+            res = "I don't know"
+    }
+}

idea/testData/intentions/branched/folding/whenToAssignment/cascadeWhen.kt.after

@@ -0,0 +1,16 @@
+// WITH_RUNTIME
+fun test(x: Any) {
+    var res: String
+    <caret>res = when (x) {
+        is String ->
+            if (x.length > 3) "long string"
+            else "short string"
+        is Int ->
+            if (x > 999 || x < -99) "long int"
+            else "short int"
+        is Long ->
+            TODO()
+        else ->
+            "I don't know"
+    }
+}

idea/testData/intentions/branched/folding/whenToReturn/cascadeWhen.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+fun test(x: Any): String {
+    <caret>when (x) {
+        is String ->
+            if (x.length > 3) return "long string"
+            else return "short string"
+        is Int ->
+            if (x > 999 || x < -99) return "long int"
+            else return "short int"
+        is Long ->
+            TODO()
+        else ->
+            return "I don't know"
+    }
+}

idea/testData/intentions/branched/folding/whenToReturn/cascadeWhen.kt.after

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+fun test(x: Any): String {
+    <caret>return when (x) {
+        is String ->
+            if (x.length > 3) "long string"
+            else "short string"
+        is Int ->
+            if (x > 999 || x < -99) "long int"
+            else "short int"
+        is Long ->
+            TODO()
+        else ->
+            "I don't know"
+    }
+}

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

@@ -975,6 +975,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/folding/ifToAssignment"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
                 }
 
+                @TestMetadata("cascadeIf.kt")
+                public void testCascadeIf() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToAssignment/cascadeIf.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("ifElseIf.kt")
                 public void testIfElseIf() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToAssignment/ifElseIf.kt");
@@ -1056,6 +1062,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/folding/ifToReturn"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
                 }
 
+                @TestMetadata("cascadeIf.kt")
+                public void testCascadeIf() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToReturn/cascadeIf.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("ifElseIf.kt")
                 public void testIfElseIf() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/ifToReturn/ifElseIf.kt");
@@ -1128,6 +1140,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/folding/whenToAssignment"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
                 }
 
+                @TestMetadata("cascadeWhen.kt")
+                public void testCascadeWhen() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/whenToAssignment/cascadeWhen.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("innerWhenTransformed.kt")
                 public void testInnerWhenTransformed() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/whenToAssignment/innerWhenTransformed.kt");
@@ -1173,6 +1191,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/folding/whenToReturn"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
                 }
 
+                @TestMetadata("cascadeWhen.kt")
+                public void testCascadeWhen() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/whenToReturn/cascadeWhen.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("innerWhenTransformed.kt")
                 public void testInnerWhenTransformed() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/folding/whenToReturn/innerWhenTransformed.kt");