Adding block to control structure when needed

compiler/frontend/src/org/jetbrains/kotlin/psi/KtContainerNodeForControlStructureBody.kt

@@ -18,4 +18,7 @@ package org.jetbrains.kotlin.psi
 
 import com.intellij.lang.ASTNode
 
-class KtContainerNodeForControlStructureBody(node: ASTNode) : KtContainerNode(node)
+class KtContainerNodeForControlStructureBody(node: ASTNode) : KtContainerNode(node) {
+    val expression: KtExpression?
+        get() = findChildByClass<KtExpression>(KtExpression::class.java)
+}

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/ReplacementPerformer.kt

@@ -16,9 +16,11 @@
 
 package org.jetbrains.kotlin.idea.replacement
 
+import com.intellij.openapi.util.Key
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
+import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import java.util.*
 
@@ -95,11 +97,36 @@ internal class ExpressionReplacementPerformer(
      * Returns statement in a block to insert statement before it
      */
     private fun findOrCreateBlockToInsertStatement(): KtExpression {
-        //TODO: 1. There can be no block above at all
-        //TODO: 2. Create block for control statements without block
-        //TODO: 3. Sometimes it's not correct because of side effects
-        return elementToBeReplaced.parentsWithSelf
-                .filterIsInstance<KtExpression>()
-                .firstOrNull { it.parent is KtBlockExpression }!!
+        //TODO: Convert expression function body into block body
+        //TODO: Sometimes it's not correct because of side effects
+
+        for (element in elementToBeReplaced.parentsWithSelf) {
+            val parent = element.parent
+            when (element) {
+                is KtContainerNodeForControlStructureBody -> { // control statement without block
+                    return element.expression!!.replaceWithBlock()
+                }
+
+                is KtExpression -> {
+                    if (parent is KtWhenEntry) { // when entry without block
+                        return element.replaceWithBlock()
+                    }
+
+                    if (parent is KtBlockExpression) return element
+                }
+            }
+        }
+
+        //TODO
+        throw UnsupportedOperationException()
     }
-}
+
+    private fun KtExpression.replaceWithBlock(): KtExpression {
+        elementToBeReplaced.putCopyableUserData(ELEMENT_TO_BE_REPLACED_KEY, Unit)
+        val blockExpression = this.replaced(KtPsiFactory(this).createSingleStatementBlock(this))
+        elementToBeReplaced = blockExpression.findDescendantOfType<KtExpression> { it.getCopyableUserData(ELEMENT_TO_BE_REPLACED_KEY) != null }!!
+        return blockExpression.statements.single()
+    }
+}
+
+private val ELEMENT_TO_BE_REPLACED_KEY = Key<Unit>("ELEMENT_TO_BE_REPLACED_KEY")

idea/testData/refactoring/inline/function/returnAtEnd/AddBlockToControlStatement.kt

@@ -0,0 +1,23 @@
+fun <caret>f(p1: Int, p2: Int): Int {
+    println(p1)
+    println(p2)
+    return p1 + p2
+}
+
+fun main(args: Array<String>) {
+    if (args.size > 0)
+        f(1, 2)
+    else
+        f(3, 4)
+
+    for (i in 1..10)
+        f(0, 1)
+
+    when (args.size) {
+        0, 1 -> println(f(1, 1))
+        else -> println(f(2, 2))
+    }
+
+    while (true)
+        println(f(5, 6))
+}

idea/testData/refactoring/inline/function/returnAtEnd/AddBlockToControlStatement.kt.after

@@ -0,0 +1,37 @@
+fun main(args: Array<String>) {
+    if (args.size > 0) {
+        println(1)
+        println(2)
+        1 + 2
+    }
+    else {
+        println(3)
+        println(4)
+        3 + 4
+    }
+
+    for (i in 1..10) {
+        println(0)
+        println(1)
+        0 + 1
+    }
+
+    when (args.size) {
+        0, 1 -> {
+            println(1)
+            println(1)
+            println(1 + 1)
+        }
+        else -> {
+            println(2)
+            println(2)
+            println(2 + 2)
+        }
+    }
+
+    while (true) {
+        println(5)
+        println(6)
+        println(5 + 6)
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/refactoring/inline/InlineTestGenerated.java

@@ -98,6 +98,12 @@ public class InlineTestGenerated extends AbstractInlineTest {
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)
         public static class ReturnAtEnd extends AbstractInlineTest {
+            @TestMetadata("AddBlockToControlStatement.kt")
+            public void testAddBlockToControlStatement() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/inline/function/returnAtEnd/AddBlockToControlStatement.kt");
+                doTest(fileName);
+            }
+
             public void testAllFilesPresentInReturnAtEnd() throws Exception {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/refactoring/inline/function/returnAtEnd"), Pattern.compile("^(\\w+)\\.kt$"), true);
             }