Introduce Variable: Properly convert function body to block when needed

#KT-5353

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiFactory.kt

@@ -214,8 +214,8 @@ public class JetPsiFactory(private val project: Project) {
         return aClass.getPrimaryConstructorModifierList()!!
     }
 
-    public fun createEmptyBody(): JetExpression {
+    public fun createEmptyBody(): JetBlockExpression {
-        return createFunction("fun foo() {}").getBodyExpression()!!
+        return createFunction("fun foo() {}").getBodyExpression() as JetBlockExpression
     }
 
     public fun createAnonymousInitializer(): JetClassInitializer {
@@ -633,8 +633,8 @@ public class JetPsiFactory(private val project: Project) {
         }
     }
 
-    public fun createFunctionBody(bodyText: String): JetExpression {
+    public fun createFunctionBody(bodyText: String): JetBlockExpression {
-        return createFunction("fun foo() {\n" + bodyText + "\n}").getBodyExpression()!!
+        return createFunction("fun foo() {\n" + bodyText + "\n}").getBodyExpression() as JetBlockExpression
     }
 
     public fun createEmptyClassObject(): JetClassObject {

idea/src/org/jetbrains/jet/plugin/intentions/ConvertToBlockBodyAction.kt

@@ -37,61 +37,57 @@ public class ConvertToBlockBodyAction : PsiElementBaseIntentionAction() {
         convert(findDeclaration(element)!!)
     }
 
-    fun convert(declaration: JetDeclarationWithBody): JetDeclarationWithBody {
+    class object {
-        val body = declaration.getBodyExpression()!!
+        fun convert(declaration: JetDeclarationWithBody): JetDeclarationWithBody {
+            val body = declaration.getBodyExpression()!!
 
-        fun generateBody(returnsValue: Boolean): JetExpression {
+            fun generateBody(returnsValue: Boolean): JetExpression {
-            val bodyType = expressionType(body)
+                val bodyType = expressionType(body)
-            val needReturn = returnsValue &&
+                val needReturn = returnsValue &&
-                             (bodyType == null || (!KotlinBuiltIns.isUnit(bodyType) && !KotlinBuiltIns.isNothing(bodyType)))
+                                 (bodyType == null || (!KotlinBuiltIns.isUnit(bodyType) && !KotlinBuiltIns.isNothing(bodyType)))
 
-            val oldBodyText = body.getText()!!
+                val oldBodyText = body.getText()!!
-            val newBodyText = if (needReturn) "return ${oldBodyText}" else oldBodyText
+                val newBodyText = if (needReturn) "return ${oldBodyText}" else oldBodyText
-            return JetPsiFactory(declaration).createFunctionBody(newBodyText)
+                return JetPsiFactory(declaration).createFunctionBody(newBodyText)
-        }
+            }
 
-        val newBody = when (declaration) {
+            val newBody = when (declaration) {
-            is JetNamedFunction -> {
+                is JetNamedFunction -> {
-                val returnType = functionReturnType(declaration)!!
+                    val returnType = functionReturnType(declaration)!!
-                if (!declaration.hasDeclaredReturnType() && !KotlinBuiltIns.isUnit(returnType)) {
+                    if (!declaration.hasDeclaredReturnType() && !KotlinBuiltIns.isUnit(returnType)) {
-                    specifyTypeExplicitly(declaration, returnType)
+                        specifyTypeExplicitly(declaration, returnType)
+                    }
+                    generateBody(!KotlinBuiltIns.isUnit(returnType) && !KotlinBuiltIns.isNothing(returnType))
                 }
 
-                val newBody = generateBody(!KotlinBuiltIns.isUnit(returnType) && !KotlinBuiltIns.isNothing(returnType))
+                is JetPropertyAccessor -> generateBody(declaration.isGetter())
 
-                declaration.getEqualsToken()!!.delete()
+                else -> throw RuntimeException("Unknown declaration type: $declaration")
-                body.replace(newBody)
             }
 
-            is JetPropertyAccessor -> {
+            declaration.getEqualsToken()!!.delete()
-                val newBody = generateBody(declaration.isGetter())
+            body.replace(newBody)
-                declaration.getEqualsToken()!!.delete()
+            return declaration
-                body.replace(newBody)
-            }
-
-            else -> throw RuntimeException("Unknown declaration type: $declaration")
         }
 
-        return newBody.getParent() as JetDeclarationWithBody
+        private fun findDeclaration(element: PsiElement): JetDeclarationWithBody? {
-    }
+            val declaration = element.getStrictParentOfType<JetDeclarationWithBody>()
+            if (declaration == null || declaration is JetFunctionLiteral || declaration.hasBlockBody()) return null
+            val body = declaration.getBodyExpression()
+            if (body == null) return null
 
-    private fun findDeclaration(element: PsiElement): JetDeclarationWithBody? {
+            return when (declaration) {
-        val declaration = element.getStrictParentOfType<JetDeclarationWithBody>()
+                is JetNamedFunction -> {
-        if (declaration == null || declaration is JetFunctionLiteral || declaration.hasBlockBody()) return null
+                    val returnType = functionReturnType(declaration)
-        val body = declaration.getBodyExpression()
+                    if (returnType == null) return null
-        if (body == null) return null
+                    if (!declaration.hasDeclaredReturnType() && returnType.isError()) return null // do not convert when type is implicit and unknown
+                    declaration
+                }
 
-        return when (declaration) {
+                is JetPropertyAccessor -> declaration
-            is JetNamedFunction -> {
+
-                val returnType = functionReturnType(declaration)
+                else -> throw RuntimeException("Unknown declaration type: $declaration")
-                if (returnType == null) return null
-                if (!declaration.hasDeclaredReturnType() && returnType.isError()) return null // do not convert when type is implicit and unknown
-                declaration
             }
-
-            is JetPropertyAccessor -> declaration
-
-            else -> throw RuntimeException("Unknown declaration type: $declaration")
         }
     }
 }

idea/src/org/jetbrains/jet/plugin/quickfix/createFromUsage/createVariable/CreateLocalVariableActionFactory.kt

@@ -54,7 +54,7 @@ object CreateLocalVariableActionFactory: JetSingleIntentionActionFactory() {
                 with (CallableBuilderConfiguration(propertyInfo.singletonOrEmptyList(), assignment ?: refExpr, file!!, editor!!).createBuilder()) {
                     val actualContainer = when (container) {
                         is JetBlockExpression -> container
-                        else -> ConvertToBlockBodyAction().convert(container as JetDeclarationWithBody).getBodyExpression()!!
+                        else -> ConvertToBlockBodyAction.convert(container as JetDeclarationWithBody).getBodyExpression()!!
                     }
                     placement = CallablePlacement.NoReceiver(actualContainer)
                     CommandProcessor.getInstance().executeCommand(project, { build() }, getText(), null)

idea/src/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.java

@@ -52,6 +52,7 @@ import org.jetbrains.jet.lexer.JetTokens;
 import org.jetbrains.jet.plugin.caches.resolve.ResolvePackage;
 import org.jetbrains.jet.plugin.codeInsight.CodeInsightUtils;
 import org.jetbrains.jet.plugin.codeInsight.ShortenReferences;
+import org.jetbrains.jet.plugin.intentions.ConvertToBlockBodyAction;
 import org.jetbrains.jet.plugin.refactoring.JetNameSuggester;
 import org.jetbrains.jet.plugin.refactoring.JetNameValidatorImpl;
 import org.jetbrains.jet.plugin.refactoring.JetRefactoringBundle;
@@ -175,7 +176,8 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
         Pass<OccurrencesChooser.ReplaceChoice> callback = new Pass<OccurrencesChooser.ReplaceChoice>() {
             @Override
             public void pass(OccurrencesChooser.ReplaceChoice replaceChoice) {
-                boolean replaceOccurrence = container != expression.getParent();
+                boolean replaceOccurrence =
+                        container != expression.getParent() || container instanceof JetNamedFunction;
                 List<JetExpression> allReplaces;
                 if (OccurrencesChooser.ReplaceChoice.ALL == replaceChoice) {
                     if (allOccurrences.size() > 1) replaceOccurrence = true;
@@ -240,8 +242,10 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
     private static Runnable introduceVariable(
             final JetExpression expression,
             final String[] suggestedNames,
-            final List<JetExpression> allReplaces, final PsiElement commonContainer,
+            final List<JetExpression> allReplaces,
-            final PsiElement commonParent, final boolean replaceOccurrence,
+            final PsiElement commonContainer,
+            final PsiElement commonParent,
+            final boolean replaceOccurrence,
             final Ref<JetProperty> propertyRef,
             final ArrayList<JetExpression> references,
             final Ref<JetExpression> reference,
@@ -253,6 +257,46 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
         return new Runnable() {
             @Override
             public void run() {
+                if (commonContainer instanceof JetNamedFunction) {
+                    JetDeclarationWithBody originalDeclaration = (JetDeclarationWithBody) commonContainer;
+                    final JetExpression originalBody = originalDeclaration.getBodyExpression();
+                    assert originalBody != null : "Original body is not found: " + originalDeclaration;
+
+                    JetDeclarationWithBody newDeclaration = ConvertToBlockBodyAction.OBJECT$.convert(originalDeclaration);
+
+                    JetBlockExpression newCommonContainer = (JetBlockExpression) newDeclaration.getBodyExpression();
+                    assert newCommonContainer != null : "New body is not found: " + newDeclaration;
+
+                    JetExpression resultExpression = (JetExpression) newCommonContainer.getStatements().get(0);
+                    if (resultExpression instanceof JetReturnExpression && !(originalBody instanceof JetReturnExpression)) {
+                        resultExpression = ((JetReturnExpression) resultExpression).getReturnedExpression();
+                    }
+                    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 {
+                    run(expression, commonContainer, commonParent, allReplaces);
+                }
+            }
+
+            private void run(
+                    JetExpression expression,
+                    PsiElement commonContainer,
+                    PsiElement commonParent,
+                    List<JetExpression> allReplaces
+            ) {
                 String variableText = "val " + suggestedNames[0];
                 if (noTypeInference) {
                     variableText += ": " + IdeDescriptorRenderers.SOURCE_CODE.renderType(expressionType);
@@ -285,10 +329,15 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                     JetExpression emptyBody = psiFactory.createEmptyBody();
                     PsiElement firstChild = emptyBody.getFirstChild();
                     emptyBody.addAfter(psiFactory.createNewLine(), firstChild);
+
                     if (replaceOccurrence && commonContainer != null) {
                         for (JetExpression replace : allReplaces) {
-                            replaceExpression(replace);
+                            JetExpression exprAfterReplace = replaceExpression(replace);
+                            if (anchor == replace) {
+                                anchor = exprAfterReplace;
+                            }
                         }
+
                         PsiElement oldElement = commonContainer;
                         if (commonContainer instanceof JetWhenEntry) {
                             JetExpression body = ((JetWhenEntry)commonContainer).getExpression();
@@ -296,12 +345,6 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                                 oldElement = body;
                             }
                         }
-                        else if (commonContainer instanceof JetNamedFunction) {
-                            JetExpression body = ((JetNamedFunction)commonContainer).getBodyExpression();
-                            if (body != null) {
-                                oldElement = body;
-                            }
-                        }
                         else if (commonContainer instanceof JetContainerNode) {
                             JetContainerNode container = (JetContainerNode)commonContainer;
                             PsiElement[] children = container.getChildren();
@@ -316,11 +359,7 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                         int diff = actualExpression.getTextRange().getStartOffset() - oldElement.getTextRange().getStartOffset();
                         String actualExpressionText = actualExpression.getText();
                         PsiElement newElement = emptyBody.addAfter(oldElement, firstChild);
-                        PsiElement elem = newElement.findElementAt(diff);
+                        PsiElement elem = findElementByOffsetAndText(diff, actualExpressionText, newElement);
-                        while (elem != null && !(elem instanceof JetExpression &&
-                                                 actualExpressionText.equals(elem.getText()))) {
-                            elem = elem.getParent();
-                        }
                         if (elem != null) {
                             reference.set((JetExpression)elem);
                         }
@@ -330,12 +369,8 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                         actualExpression = reference.get();
                         diff = actualExpression.getTextRange().getStartOffset() - emptyBody.getTextRange().getStartOffset();
                         actualExpressionText = actualExpression.getText();
-                        emptyBody = (JetExpression)anchor.replace(emptyBody);
+                        emptyBody = (JetBlockExpression) anchor.replace(emptyBody);
-                        elem = emptyBody.findElementAt(diff);
+                        elem = findElementByOffsetAndText(diff, actualExpressionText, emptyBody);
-                        while (elem != null && !(elem instanceof JetExpression &&
-                                                 actualExpressionText.equals(elem.getText()))) {
-                            elem = elem.getParent();
-                        }
                         if (elem != null) {
                             reference.set((JetExpression)elem);
                         }
@@ -343,25 +378,14 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                     else {
                         property = (JetProperty)emptyBody.addAfter(property, firstChild);
                         emptyBody.addAfter(psiFactory.createNewLine(), firstChild);
-                        emptyBody = (JetExpression)anchor.replace(emptyBody);
+                        emptyBody = (JetBlockExpression) anchor.replace(emptyBody);
                     }
                     for (PsiElement child : emptyBody.getChildren()) {
                         if (child instanceof JetProperty) {
                             property = (JetProperty)child;
                         }
                     }
-                    if (commonContainer instanceof JetNamedFunction) {
+                    if (commonContainer instanceof JetContainerNode) {
-                        //we should remove equals sign
-                        JetNamedFunction function = (JetNamedFunction)commonContainer;
-                        if (!function.hasDeclaredReturnType()) {
-                            //todo: add return type
-                        }
-
-                        PsiElement equalsToken = function.getEqualsToken();
-                        assert equalsToken != null : "Function without block body was expected: " + function.getText();
-                        equalsToken.delete();
-                    }
-                    else if (commonContainer instanceof JetContainerNode) {
                         JetContainerNode node = (JetContainerNode)commonContainer;
                         if (node.getParent() instanceof JetIfExpression) {
                             PsiElement next = node.getNextSibling();
@@ -390,7 +414,23 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                 }
             }
 
-            private void replaceExpression(JetExpression replace) {
+            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()))) {
+                    elem = elem.getParent();
+                }
+                return elem;
+            }
+
+            private JetExpression replaceExpression(JetExpression replace) {
                 boolean isActualExpression = expression == replace;
 
                 JetExpression replacement = psiFactory.createExpression(suggestedNames[0]);
@@ -407,6 +447,8 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                 }
                 references.add(result);
                 if (isActualExpression) reference.set(result);
+
+                return result;
             }
         };
     }

idea/testData/refactoring/introduceVariable/FunctionAddBlock.kt.after

@@ -1,3 +1,4 @@
 fun x(): Int {
     val i = 1
+    return i
 }

idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt

@@ -1 +1,6 @@
-fun x(): Int = println(<selection>1</selection>)
+fun foo(a: Int) =
+        if (a > 1) {
+            (<selection>a + 1</selection>) * (a - 1)
+        } else {
+            a * (a + 1)
+        }

idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt.after

@@ -1,4 +1,8 @@
-fun x(): Int {
+fun foo(a: Int): Int {
-    val i = 1
+    val i = a + 1
-    println(i)
+    return if (a > 1) {
+        i * (a - 1)
+    } else {
+        a * i
+    }
 }