Don't use plain text generation

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

@@ -566,10 +566,14 @@ public class JetPsiFactory(private val project: Project) {
         }
     }
 
-    public fun createFunctionBody(bodyText: String): JetBlockExpression {
+    public fun createBlock(bodyText: String): JetBlockExpression {
         return createFunction("fun foo() {\n" + bodyText + "\n}").getBodyExpression() as JetBlockExpression
     }
 
+    public fun createSingleStatementBlock(statement: JetExpression): JetBlockExpression {
+        return createDeclarationByPattern<JetNamedFunction>("fun foo() {\n$0\n}", statement).getBodyExpression() as JetBlockExpression
+    }
+
     public fun createComment(text: String): PsiComment {
         val file = createFile(text)
         val comments = file.getChildren().filterIsInstance<PsiComment>()

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/jetPsiUtil.kt

@@ -573,12 +573,16 @@ public inline fun <reified T : PsiElement> PsiElement.forEachDescendantOfType(no
 }
 
 public inline fun <reified T : PsiElement> PsiElement.anyDescendantOfType(noinline predicate: (T) -> Boolean): Boolean {
-    var result = false
+    return findDescendantOfType(predicate) != null
+}
+
+public inline fun <reified T : PsiElement> PsiElement.findDescendantOfType(noinline predicate: (T) -> Boolean): T? {
+    var result: T? = null
     this.accept(object : PsiRecursiveElementVisitor(){
         override fun visitElement(element: PsiElement) {
-            if (result) return
+            if (result != null) return
             if (element is T && predicate(element)) {
-                result = true
+                result = element
                 return
             }
             super.visitElement(element)
@@ -587,7 +591,7 @@ public inline fun <reified T : PsiElement> PsiElement.anyDescendantOfType(noinli
     return result
 }
 
-public inline fun <reified T : PsiElement> PsiElement.collectDescendantsOfType(noinline predicate: (T) -> Boolean = { true }): Collection<T> {
+public inline fun <reified T : PsiElement> PsiElement.collectDescendantsOfType(noinline predicate: (T) -> Boolean = { true }): List<T> {
     val result = ArrayList<T>()
     forEachDescendantOfType<T> {
         if (predicate(it)) {

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

@@ -39,7 +39,7 @@ public class AddBracesIntention : JetSelfTargetingIntention<JetExpression>(javaC
         }
 
         val psiFactory = JetPsiFactory(element)
-        expression.replace(psiFactory.createFunctionBody(expression.getText()))
+        expression.replace(psiFactory.createSingleStatementBlock(expression))
 
         if (element is JetDoWhileExpression) { // remove new line between '}' and while
             (element.getBody()!!.getParent().getNextSibling() as? PsiWhiteSpace)?.delete()

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

@@ -56,9 +56,9 @@ public class ConvertToBlockBodyIntention : JetSelfTargetingIntention<JetDeclarat
                 val needReturn = returnsValue &&
                                  (bodyType == null || (!KotlinBuiltIns.isUnit(bodyType) && !KotlinBuiltIns.isNothing(bodyType)))
 
-                val oldBodyText = body.getText()!!
+                val factory = JetPsiFactory(declaration)
-                val newBodyText = if (needReturn) "return ${oldBodyText}" else oldBodyText
+                val statement = if (needReturn) factory.createExpressionByPattern("return $0", body) else body
-                return JetPsiFactory(declaration).createFunctionBody(newBodyText)
+                return factory.createSingleStatementBlock(statement)
             }
 
             val newBody = when (declaration) {

idea/src/org/jetbrains/kotlin/idea/intentions/declarations/ConvertMemberToExtensionIntention.kt

@@ -85,7 +85,7 @@ public class ConvertMemberToExtensionIntention : JetSelfTargetingRangeIntention<
         when (extension) {
             is JetFunction -> {
                 if (!extension.hasBody()) {
-                    extension.add(psiFactory.createFunctionBody(THROW_UNSUPPORTED_OPERATION_EXCEPTION))
+                    extension.add(psiFactory.createBlock(THROW_UNSUPPORTED_OPERATION_EXCEPTION))
                     selectBody(extension)
                 }
             }

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/CallableBuilder.kt

@@ -708,7 +708,7 @@ class CallableBuilder(val config: CallableBuilderConfiguration) {
                 throw IncorrectOperationException("Failed to parse file template", e)
             }
 
-            oldBody.replace(JetPsiFactory(func).createFunctionBody(bodyText))
+            oldBody.replace(JetPsiFactory(func).createBlock(bodyText))
         }
 
         private fun setupCallTypeArguments(callElement: JetCallElement, typeParameters: List<TypeParameterDescriptor>) {

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceUtil.kt

@@ -16,14 +16,21 @@
 
 package org.jetbrains.kotlin.idea.refactoring.introduce
 
-import org.jetbrains.kotlin.psi.psiUtil.*
+import com.intellij.openapi.editor.Editor
-import org.jetbrains.kotlin.idea.refactoring.*
+import com.intellij.openapi.editor.ScrollType
-import com.intellij.openapi.editor.*
+import com.intellij.openapi.project.Project
-import com.intellij.psi.*
+import com.intellij.openapi.util.Key
-import com.intellij.psi.util.*
+import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.idea.codeInsight.*
+import com.intellij.psi.PsiFile
-import com.intellij.openapi.project.*
+import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.idea.codeInsight.CodeInsightUtils
 import org.jetbrains.kotlin.idea.core.refactoring.chooseContainerElementIfNecessary
+import org.jetbrains.kotlin.idea.refactoring.JetRefactoringBundle
+import org.jetbrains.kotlin.idea.refactoring.JetRefactoringUtil
+import org.jetbrains.kotlin.psi.JetExpression
+import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
+import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
+import org.jetbrains.kotlin.psi.psiUtil.getOutermostParentContainedIn
 
 fun showErrorHint(project: Project, editor: Editor, message: String, title: String) {
     CodeInsightUtils.showErrorHint(project, editor, message, title, null)
@@ -122,3 +129,21 @@ fun selectElementsWithTargetParent(
     editor.getScrollingModel().scrollToCaret(ScrollType.MAKE_VISIBLE)
     selectSingleExpression()
 }
+
+fun PsiElement.findExpressionByCopyableDataAndClearIt(key: Key<Boolean>): JetExpression {
+    val result = findDescendantOfType<JetExpression> { it.getCopyableUserData(key) != null }!!
+    result.putCopyableUserData(key, null)
+    return result
+}
+
+fun PsiElement.findElementByCopyableDataAndClearIt(key: Key<Boolean>): PsiElement {
+    val result = findDescendantOfType<PsiElement> { it.getCopyableUserData(key) != null }!!
+    result.putCopyableUserData(key, null)
+    return result
+}
+
+fun PsiElement.findExpressionsByCopyableDataAndClearIt(key: Key<Boolean>): List<JetExpression> {
+    val results = collectDescendantsOfType<JetExpression> { it.getCopyableUserData(key) != null }
+    results.forEach { it.putCopyableUserData(key, null) }
+    return results
+}

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.java

@@ -21,6 +21,7 @@ import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.command.CommandProcessor;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.Pass;
 import com.intellij.openapi.util.Ref;
 import com.intellij.psi.PsiDocumentManager;
@@ -45,6 +46,7 @@ import org.jetbrains.kotlin.idea.intentions.RemoveCurlyBracesFromTemplateIntenti
 import org.jetbrains.kotlin.idea.refactoring.JetNameValidatorImpl;
 import org.jetbrains.kotlin.idea.refactoring.JetRefactoringBundle;
 import org.jetbrains.kotlin.idea.refactoring.JetRefactoringUtil;
+import org.jetbrains.kotlin.idea.refactoring.introduce.IntroducePackage;
 import org.jetbrains.kotlin.idea.refactoring.introduce.KotlinIntroduceHandlerBase;
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers;
 import org.jetbrains.kotlin.idea.util.ShortenReferences;
@@ -65,10 +67,7 @@ import org.jetbrains.kotlin.types.JetType;
 import org.jetbrains.kotlin.types.TypeUtils;
 import org.jetbrains.kotlin.types.checker.JetTypeChecker;
 
-import java.util.ArrayList;
+import java.util.*;
-import java.util.Collections;
-import java.util.LinkedHashSet;
-import java.util.List;
 
 import static org.jetbrains.kotlin.psi.PsiPackage.JetPsiFactory;
 
@@ -272,28 +271,24 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                     final JetExpression originalBody = originalDeclaration.getBodyExpression();
                     assert originalBody != null : "Original body is not found: " + originalDeclaration;
 
+                    Key<Boolean> EXPRESSION_KEY = Key.create("EXPRESSION_KEY");
+                    Key<Boolean> REPLACE_KEY = Key.create("REPLACE_KEY");
+                    Key<Boolean> COMMON_PARENT_KEY = Key.create("COMMON_PARENT_KEY");
+                    expression.putCopyableUserData(EXPRESSION_KEY, true);
+                    for (JetExpression replace : allReplaces) {
+                        replace.putCopyableUserData(REPLACE_KEY, true);
+                    }
+                    commonParent.putCopyableUserData(COMMON_PARENT_KEY, true);
+
                     JetDeclarationWithBody newDeclaration = ConvertToBlockBodyIntention.Companion.convert(originalDeclaration);
 
                     JetBlockExpression newCommonContainer = (JetBlockExpression) newDeclaration.getBodyExpression();
                     assert newCommonContainer != null : "New body is not found: " + newDeclaration;
 
-                    JetExpression resultExpression = (JetExpression) newCommonContainer.getStatements().get(0);
+                    JetExpression newExpression = IntroducePackage.findExpressionByCopyableDataAndClearIt(newCommonContainer, EXPRESSION_KEY);
-                    if (resultExpression instanceof JetReturnExpression && !(originalBody instanceof JetReturnExpression)) {
+                    PsiElement newCommonParent = IntroducePackage.findElementByCopyableDataAndClearIt(newCommonContainer, COMMON_PARENT_KEY);
-                        resultExpression = ((JetReturnExpression) resultExpression).getReturnedExpression();
+                    List<JetExpression> newAllReplaces = IntroducePackage.findExpressionsByCopyableDataAndClearIt(newCommonContainer, REPLACE_KEY);
-                    }
-                    final JetExpression finalResultExpression = resultExpression;
 
-                    JetExpression newExpression = (JetExpression) findElementCounterpart(expression, originalBody, resultExpression);
-                    PsiElement newCommonParent = findElementCounterpart(commonParent, originalBody, resultExpression);
-                    List<JetExpression> newAllReplaces = KotlinPackage.map(
-                            allReplaces,
-                            new Function1<JetExpression, JetExpression>() {
-                                @Override
-                                public JetExpression invoke(JetExpression expression) {
-                                    return (JetExpression) findElementCounterpart(expression, originalBody, finalResultExpression);
-                                }
-                            }
-                    );
                     run(newExpression, newCommonContainer, newCommonParent, newAllReplaces);
                 }
                 else {
@@ -438,14 +433,6 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                 }
             }
 
-            private PsiElement findElementCounterpart(PsiElement oldElement, PsiElement oldContainer, PsiElement newContainer) {
-                return findElementByOffsetAndText(
-                        oldElement.getTextOffset() - oldContainer.getTextOffset(),
-                        oldElement.getText(),
-                        newContainer
-                );
-            }
-
             private PsiElement findElementByOffsetAndText(int offset, String text, PsiElement newContainer) {
                 PsiElement elem = newContainer.findElementAt(offset);
                 while (elem != null && !(elem instanceof JetExpression && text.equals(elem.getText()))) {

idea/testData/intentions/convertToBlockBody/addSpace.kt

@@ -0,0 +1,3 @@
+fun <caret>foo(): Int=bar()
+
+fun bar() = 1

idea/testData/intentions/convertToBlockBody/addSpace.kt.after

@@ -0,0 +1,5 @@
+fun <caret>foo(): Int {
+    return bar()
+}
+
+fun bar() = 1

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

@@ -3417,6 +3417,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)
     public static class ConvertToBlockBody extends AbstractIntentionTest {
+        @TestMetadata("addSpace.kt")
+        public void testAddSpace() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/intentions/convertToBlockBody/addSpace.kt");
+            doTest(fileName);
+        }
+
         public void testAllFilesPresentInConvertToBlockBody() throws Exception {
             JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/convertToBlockBody"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
         }