Do not insert returns before Nothing in "Replace return with 'if'"

So #KT-22159 Fixed

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

@@ -19,15 +19,16 @@ package org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions
 import com.intellij.codeInsight.intention.LowPriorityAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.copied
 import org.jetbrains.kotlin.idea.intentions.SelfTargetingRangeIntention
-import org.jetbrains.kotlin.psi.KtIfExpression
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtPsiFactory
-import org.jetbrains.kotlin.psi.KtReturnExpression
-import org.jetbrains.kotlin.psi.createExpressionByPattern
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.lastBlockStatementOrThis
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.types.typeUtil.isNothing
 
 class UnfoldReturnToIfIntention : SelfTargetingRangeIntention<KtReturnExpression>(KtReturnExpression::class.java, "Replace return with 'if' expression"), LowPriorityAction {
     override fun applicabilityRange(element: KtReturnExpression): TextRange? {
@@ -37,14 +38,29 @@ class UnfoldReturnToIfIntention : SelfTargetingRangeIntention<KtReturnExpression
 
     override fun applyTo(element: KtReturnExpression, editor: Editor?) {
         val ifExpression = element.returnedExpression as KtIfExpression
+        val thenExpr = ifExpression.then!!.lastBlockStatementOrThis()
+        val elseExpr = ifExpression.`else`!!.lastBlockStatementOrThis()
+
         val newIfExpression = ifExpression.copied()
-        val thenExpr = newIfExpression.then!!.lastBlockStatementOrThis()
+        val newThenExpr = newIfExpression.then!!.lastBlockStatementOrThis()
-        val elseExpr = newIfExpression.`else`!!.lastBlockStatementOrThis()
+        val newElseExpr = newIfExpression.`else`!!.lastBlockStatementOrThis()
 
         val psiFactory = KtPsiFactory(element)
-        thenExpr.replace(psiFactory.createExpressionByPattern("return $0", thenExpr))
+        val context = element.analyze()
-        elseExpr.replace(psiFactory.createExpressionByPattern("return $0", elseExpr))
 
+        newThenExpr.replace(createReturnExpression(thenExpr, psiFactory, context))
+        newElseExpr.replace(createReturnExpression(elseExpr, psiFactory, context))
         element.replace(newIfExpression)
     }
+
+    companion object {
+        fun createReturnExpression(expr: KtExpression, psiFactory: KtPsiFactory, context: BindingContext): KtExpression {
+            val returnText = when (expr) {
+                is KtBreakExpression, is KtContinueExpression, is KtReturnExpression, is KtThrowExpression -> ""
+                else -> if (expr.getResolvedCall(context)?.resultingDescriptor?.returnType?.isNothing() == true) "" else "return "
+            }
+            return psiFactory.createExpressionByPattern("$returnText$0", expr)
+        }
+    }
+
 }

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

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions
 import com.intellij.codeInsight.intention.LowPriorityAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.copied
 import org.jetbrains.kotlin.idea.intentions.SelfTargetingRangeIntention
 import org.jetbrains.kotlin.psi.*
@@ -35,14 +36,17 @@ class UnfoldReturnToWhenIntention : SelfTargetingRangeIntention<KtReturnExpressi
     }
 
     override fun applyTo(element: KtReturnExpression, editor: Editor?) {
+        val psiFactory = KtPsiFactory(element)
+        val context = element.analyze()
+
         val whenExpression = element.returnedExpression as KtWhenExpression
         val newWhenExpression = whenExpression.copied()
 
-        for (entry in newWhenExpression.entries) {
+        whenExpression.entries.zip(newWhenExpression.entries).forEach { (entry, newEntry) ->
             val expr = entry.expression!!.lastBlockStatementOrThis()
-            expr.replace(KtPsiFactory(element).createExpressionByPattern("return $0", expr))
+            val newExpr = newEntry.expression!!.lastBlockStatementOrThis()
+            newExpr.replace(UnfoldReturnToIfIntention.createReturnExpression(expr, psiFactory, context))
         }
-
         element.replace(newWhenExpression)
     }
 }

idea/testData/intentions/branched/unfolding/returnToIf/ifWithBreak.kt

@@ -0,0 +1,10 @@
+fun test(b: Boolean): Int {
+    while (true) {
+        <caret>return if (b) {
+            1
+        } else {
+            break
+        }
+    }
+    return 0
+}

idea/testData/intentions/branched/unfolding/returnToIf/ifWithBreak.kt.after

@@ -0,0 +1,10 @@
+fun test(b: Boolean): Int {
+    while (true) {
+        if (b) {
+            return 1
+        } else {
+            break
+        }
+    }
+    return 0
+}

idea/testData/intentions/branched/unfolding/returnToIf/ifWithContinue.kt

@@ -0,0 +1,12 @@
+fun test(b: Boolean): Int {
+    var i = 0
+    while (i == 0) {
+        <caret>return if (b) {
+            1
+        } else {
+            i++
+            continue
+        }
+    }
+    return 0
+}

idea/testData/intentions/branched/unfolding/returnToIf/ifWithContinue.kt.after

@@ -0,0 +1,12 @@
+fun test(b: Boolean): Int {
+    var i = 0
+    while (i == 0) {
+        if (b) {
+            return 1
+        } else {
+            i++
+            continue
+        }
+    }
+    return 0
+}

idea/testData/intentions/branched/unfolding/returnToIf/ifWithInnerReturn.kt

@@ -0,0 +1,9 @@
+fun test(b: Boolean): Int {
+    while (true) {
+        <caret>return if (b) {
+            1
+        } else {
+            return 0
+        }
+    }
+}

idea/testData/intentions/branched/unfolding/returnToIf/ifWithInnerReturn.kt.after

@@ -0,0 +1,9 @@
+fun test(b: Boolean): Int {
+    while (true) {
+        if (b) {
+            return 1
+        } else {
+            return 0
+        }
+    }
+}

idea/testData/intentions/branched/unfolding/returnToIf/ifWithNothing.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun test(b: Boolean): Int {
+    <caret>return if (b) {
+        1
+    } else {
+        TODO()
+    }
+}

idea/testData/intentions/branched/unfolding/returnToIf/ifWithNothing.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+fun test(b: Boolean): Int {
+    if (b) {
+        return 1
+    } else {
+        TODO()
+    }
+}

idea/testData/intentions/branched/unfolding/returnToIf/ifWithThrow.kt

@@ -0,0 +1,7 @@
+fun test(b: Boolean): Int {
+    <caret>return if (b) {
+        throw AssertionError()
+    } else {
+        1
+    }
+}

idea/testData/intentions/branched/unfolding/returnToIf/ifWithThrow.kt.after

@@ -0,0 +1,7 @@
+fun test(b: Boolean): Int {
+    if (b) {
+        throw AssertionError()
+    } else {
+        return 1
+    }
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithBreak.kt

@@ -0,0 +1,9 @@
+fun test(b: Boolean): Int {
+    loop@ while (true) {
+        <caret>return when (b) {
+            true -> 1
+            else -> break@loop
+        }
+    }
+    return 0
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithBreak.kt.after

@@ -0,0 +1,9 @@
+fun test(b: Boolean): Int {
+    loop@ while (true) {
+        when (b) {
+            true -> return 1
+            else -> break@loop
+        }
+    }
+    return 0
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithContinue.kt

@@ -0,0 +1,13 @@
+fun test(b: Boolean): Int {
+    var i = 0
+    loop@while (i == 0) {
+        <caret>return when (b) {
+            true -> 1
+            else -> {
+                i++
+                continue@loop
+            }
+        }
+    }
+    return 0
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithContinue.kt.after

@@ -0,0 +1,13 @@
+fun test(b: Boolean): Int {
+    var i = 0
+    loop@while (i == 0) {
+        when (b) {
+            true -> return 1
+            else -> {
+                i++
+                continue@loop
+            }
+        }
+    }
+    return 0
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithInnerReturn.kt

@@ -0,0 +1,8 @@
+fun test(b: Boolean): Int {
+    while (true) {
+        <caret>return when (b) {
+            true -> 1
+            else -> return 0
+        }
+    }
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithInnerReturn.kt.after

@@ -0,0 +1,8 @@
+fun test(b: Boolean): Int {
+    while (true) {
+        when (b) {
+            true -> return 1
+            else -> return 0
+        }
+    }
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithNothing.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun test(b: Boolean): Int {
+    <caret>return when (b) {
+        true -> 1
+        else -> TODO()
+    }
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithNothing.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+fun test(b: Boolean): Int {
+    when (b) {
+        true -> return 1
+        else -> TODO()
+    }
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithThrow.kt

@@ -0,0 +1,6 @@
+fun foo(b: Boolean): Int {
+    <caret>return when (b) {
+        true -> throw AssertionError()
+        else -> 1
+    }
+}

idea/testData/intentions/branched/unfolding/returnToWhen/whenWithThrow.kt.after

@@ -0,0 +1,6 @@
+fun foo(b: Boolean): Int {
+    when (b) {
+        true -> throw AssertionError()
+        else -> return 1
+    }
+}

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

@@ -2467,6 +2467,36 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/unfolding/returnToIf"), Pattern.compile("^([\\w\\-_]+)\\.(kt|kts)$"), TargetBackend.ANY, true);
                 }
 
+                @TestMetadata("ifWithBreak.kt")
+                public void testIfWithBreak() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToIf/ifWithBreak.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("ifWithContinue.kt")
+                public void testIfWithContinue() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToIf/ifWithContinue.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("ifWithInnerReturn.kt")
+                public void testIfWithInnerReturn() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToIf/ifWithInnerReturn.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("ifWithNothing.kt")
+                public void testIfWithNothing() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToIf/ifWithNothing.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("ifWithThrow.kt")
+                public void testIfWithThrow() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToIf/ifWithThrow.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("innerIfTransformed.kt")
                 public void testInnerIfTransformed() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToIf/innerIfTransformed.kt");
@@ -2511,6 +2541,36 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToWhen/simpleWhenWithBlocks.kt");
                     doTest(fileName);
                 }
+
+                @TestMetadata("whenWithBreak.kt")
+                public void testWhenWithBreak() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToWhen/whenWithBreak.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("whenWithContinue.kt")
+                public void testWhenWithContinue() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToWhen/whenWithContinue.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("whenWithInnerReturn.kt")
+                public void testWhenWithInnerReturn() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToWhen/whenWithInnerReturn.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("whenWithNothing.kt")
+                public void testWhenWithNothing() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToWhen/whenWithNothing.kt");
+                    doTest(fileName);
+                }
+
+                @TestMetadata("whenWithThrow.kt")
+                public void testWhenWithThrow() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/unfolding/returnToWhen/whenWithThrow.kt");
+                    doTest(fileName);
+                }
             }
         }