diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiFactory.kt b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiFactory.kt index 119e3d9efe6..cb866b31d8a 100644 --- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiFactory.kt +++ b/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 { - return createFunction("fun foo() {}").getBodyExpression()!! + public fun createEmptyBody(): JetBlockExpression { + 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 { - return createFunction("fun foo() {\n" + bodyText + "\n}").getBodyExpression()!! + public fun createFunctionBody(bodyText: String): JetBlockExpression { + return createFunction("fun foo() {\n" + bodyText + "\n}").getBodyExpression() as JetBlockExpression } public fun createEmptyClassObject(): JetClassObject { diff --git a/idea/src/org/jetbrains/jet/plugin/intentions/ConvertToBlockBodyAction.kt b/idea/src/org/jetbrains/jet/plugin/intentions/ConvertToBlockBodyAction.kt index 25fc5bae536..9daf3216d67 100644 --- a/idea/src/org/jetbrains/jet/plugin/intentions/ConvertToBlockBodyAction.kt +++ b/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 { - val body = declaration.getBodyExpression()!! + class object { + fun convert(declaration: JetDeclarationWithBody): JetDeclarationWithBody { + val body = declaration.getBodyExpression()!! - fun generateBody(returnsValue: Boolean): JetExpression { - val bodyType = expressionType(body) - val needReturn = returnsValue && - (bodyType == null || (!KotlinBuiltIns.isUnit(bodyType) && !KotlinBuiltIns.isNothing(bodyType))) + fun generateBody(returnsValue: Boolean): JetExpression { + val bodyType = expressionType(body) + val needReturn = returnsValue && + (bodyType == null || (!KotlinBuiltIns.isUnit(bodyType) && !KotlinBuiltIns.isNothing(bodyType))) - val oldBodyText = body.getText()!! - val newBodyText = if (needReturn) "return ${oldBodyText}" else oldBodyText - return JetPsiFactory(declaration).createFunctionBody(newBodyText) - } + val oldBodyText = body.getText()!! + val newBodyText = if (needReturn) "return ${oldBodyText}" else oldBodyText + return JetPsiFactory(declaration).createFunctionBody(newBodyText) + } - val newBody = when (declaration) { - is JetNamedFunction -> { - val returnType = functionReturnType(declaration)!! - if (!declaration.hasDeclaredReturnType() && !KotlinBuiltIns.isUnit(returnType)) { - specifyTypeExplicitly(declaration, returnType) + val newBody = when (declaration) { + is JetNamedFunction -> { + val returnType = functionReturnType(declaration)!! + if (!declaration.hasDeclaredReturnType() && !KotlinBuiltIns.isUnit(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() - body.replace(newBody) + else -> throw RuntimeException("Unknown declaration type: $declaration") } - is JetPropertyAccessor -> { - val newBody = generateBody(declaration.isGetter()) - declaration.getEqualsToken()!!.delete() - body.replace(newBody) - } - - else -> throw RuntimeException("Unknown declaration type: $declaration") + declaration.getEqualsToken()!!.delete() + body.replace(newBody) + return declaration } - return newBody.getParent() as JetDeclarationWithBody - } + private fun findDeclaration(element: PsiElement): JetDeclarationWithBody? { + val declaration = element.getStrictParentOfType() + 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? { - val declaration = element.getStrictParentOfType() - if (declaration == null || declaration is JetFunctionLiteral || declaration.hasBlockBody()) return null - val body = declaration.getBodyExpression() - if (body == null) return null + return when (declaration) { + is JetNamedFunction -> { + val returnType = functionReturnType(declaration) + if (returnType == null) return null + if (!declaration.hasDeclaredReturnType() && returnType.isError()) return null // do not convert when type is implicit and unknown + declaration + } - return when (declaration) { - is JetNamedFunction -> { - val returnType = functionReturnType(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") } - - is JetPropertyAccessor -> declaration - - else -> throw RuntimeException("Unknown declaration type: $declaration") } } } diff --git a/idea/src/org/jetbrains/jet/plugin/quickfix/createFromUsage/createVariable/CreateLocalVariableActionFactory.kt b/idea/src/org/jetbrains/jet/plugin/quickfix/createFromUsage/createVariable/CreateLocalVariableActionFactory.kt index 38562f8cf6b..c30f5c2e01d 100644 --- a/idea/src/org/jetbrains/jet/plugin/quickfix/createFromUsage/createVariable/CreateLocalVariableActionFactory.kt +++ b/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) diff --git a/idea/src/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.java b/idea/src/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.java index 81ac9d2ded6..23e6cb7403a 100644 --- a/idea/src/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.java +++ b/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 callback = new Pass() { @Override public void pass(OccurrencesChooser.ReplaceChoice replaceChoice) { - boolean replaceOccurrence = container != expression.getParent(); + boolean replaceOccurrence = + container != expression.getParent() || container instanceof JetNamedFunction; List 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 allReplaces, final PsiElement commonContainer, - final PsiElement commonParent, final boolean replaceOccurrence, + final List allReplaces, + final PsiElement commonContainer, + final PsiElement commonParent, + final boolean replaceOccurrence, final Ref propertyRef, final ArrayList references, final Ref 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 newAllReplaces = KotlinPackage.map( + allReplaces, + new Function1() { + @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 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); - while (elem != null && !(elem instanceof JetExpression && - actualExpressionText.equals(elem.getText()))) { - elem = elem.getParent(); - } + PsiElement elem = findElementByOffsetAndText(diff, actualExpressionText, newElement); 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); - elem = emptyBody.findElementAt(diff); - while (elem != null && !(elem instanceof JetExpression && - actualExpressionText.equals(elem.getText()))) { - elem = elem.getParent(); - } + emptyBody = (JetBlockExpression) anchor.replace(emptyBody); + elem = findElementByOffsetAndText(diff, actualExpressionText, emptyBody); 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) { - //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) { + 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; } }; } diff --git a/idea/testData/refactoring/introduceVariable/FunctionAddBlock.kt.after b/idea/testData/refactoring/introduceVariable/FunctionAddBlock.kt.after index 31ca620fef4..92136161933 100644 --- a/idea/testData/refactoring/introduceVariable/FunctionAddBlock.kt.after +++ b/idea/testData/refactoring/introduceVariable/FunctionAddBlock.kt.after @@ -1,3 +1,4 @@ fun x(): Int { val i = 1 + return i } \ No newline at end of file diff --git a/idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt b/idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt index 3462b44ccf5..6d65b432f02 100644 --- a/idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt +++ b/idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt @@ -1 +1,6 @@ -fun x(): Int = println(1) \ No newline at end of file +fun foo(a: Int) = + if (a > 1) { + (a + 1) * (a - 1) + } else { + a * (a + 1) + } \ No newline at end of file diff --git a/idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt.after b/idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt.after index 9ab41811fd3..8e07278b67e 100644 --- a/idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt.after +++ b/idea/testData/refactoring/introduceVariable/FunctionAddBlockInner.kt.after @@ -1,4 +1,8 @@ -fun x(): Int { - val i = 1 - println(i) +fun foo(a: Int): Int { + val i = a + 1 + return if (a > 1) { + i * (a - 1) + } else { + a * i + } } \ No newline at end of file