diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/AddLabeledReturnInLambdaIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/AddLabeledReturnInLambdaIntention.kt index 6479971f25e..752022c0483 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/AddLabeledReturnInLambdaIntention.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/AddLabeledReturnInLambdaIntention.kt @@ -13,35 +13,36 @@ import org.jetbrains.kotlin.psi.* import org.jetbrains.kotlin.psi.psiUtil.* import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression -class AddLabeledReturnInLambdaIntention : SelfTargetingRangeIntention( - KtLambdaExpression::class.java, +class AddLabeledReturnInLambdaIntention : SelfTargetingRangeIntention( + KtBlockExpression::class.java, "Add labeled return to last expression in a lambda" ), LowPriorityAction { - override fun applicabilityRange(element: KtLambdaExpression): TextRange? { + override fun applicabilityRange(element: KtBlockExpression): TextRange? { if (!isApplicableTo(element)) return null val labelName = createLabelName(element) ?: return null text = "Add return@$labelName" - return element.bodyExpression?.statements?.last()?.textRange + return element.statements.lastOrNull()?.textRange } - override fun applyTo(element: KtLambdaExpression, editor: Editor?) { + override fun applyTo(element: KtBlockExpression, editor: Editor?) { if (!isApplicableTo(element)) return val labelName = createLabelName(element) ?: return - val lastStatement = element.bodyExpression?.statements?.last() ?: return + val lastStatement = element.statements.lastOrNull() ?: return val newExpression = KtPsiFactory(element.project).createExpressionByPattern("return@$labelName $0", lastStatement) lastStatement.replace(newExpression) } - private fun isApplicableTo(element: KtLambdaExpression): Boolean { - val block = element.bodyExpression ?: return false - val lastStatement = block.statements.last() - return lastStatement !is KtReturnExpression && lastStatement.isUsedAsExpression(lastStatement.analyze()) + private fun isApplicableTo(block: KtBlockExpression): Boolean { + val lastStatement = block.statements.lastOrNull() + return lastStatement !is KtReturnExpression && lastStatement?.isUsedAsExpression(lastStatement.analyze()) == true } - private fun createLabelName(element: KtLambdaExpression): String? { - val block = element.bodyExpression ?: return null - val callExpression = element.getStrictParentOfType() ?: return null - val valueArgument = callExpression.valueArguments.findArgumentWithGivenBlock(block) ?: return null + private fun createLabelName(block: KtBlockExpression): String? { + val lambdaExpression = block.getStrictParentOfType() ?: return null + val callExpression = lambdaExpression.getStrictParentOfType() ?: return null + val valueArgument = callExpression.valueArguments.find { + it.getArgumentExpression()?.unpackFunctionLiteral(allowParentheses = false) === lambdaExpression + } ?: return null val lambdaLabelName = (valueArgument.getArgumentExpression() as? KtLabeledExpression)?.getLabelName() return lambdaLabelName ?: callExpression.getCallNameExpression()?.text } diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/RemoveLabeledReturnInLambdaIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/RemoveLabeledReturnInLambdaIntention.kt index ef411747be1..a1736c25358 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/RemoveLabeledReturnInLambdaIntention.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/RemoveLabeledReturnInLambdaIntention.kt @@ -19,8 +19,11 @@ class RemoveLabeledReturnInLambdaIntention : SelfTargetingIntention() ?: return false if (block.statements.lastOrNull() != element) return false - val callExpression = block.getStrictParentOfType() ?: return false - val lambdaArgument = callExpression.lambdaArguments.findArgumentWithGivenBlock(block) ?: return false + val lambdaExpression = block.getStrictParentOfType() ?: return false + val callExpression = lambdaExpression.getStrictParentOfType() ?: return false + val lambdaArgument = callExpression.lambdaArguments.find { + it.getArgumentExpression().unpackFunctionLiteral(allowParentheses = false) === lambdaExpression + } ?: return false val callName = (lambdaArgument.getArgumentExpression() as? KtLabeledExpression)?.getLabelName() ?: callExpression.getCallNameExpression()?.text ?: return false if (labelName != callName) return false diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt index c471b7f7795..37d09a5a9a0 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt +++ b/idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt @@ -294,15 +294,3 @@ fun KtCallExpression.isArrayOfMethod(): Boolean { return (descriptor.containingDeclaration as? PackageFragmentDescriptor)?.fqName == KotlinBuiltIns.BUILT_INS_PACKAGE_FQ_NAME && ARRAY_OF_METHODS.contains(descriptor.name) } - -fun List.findArgumentWithGivenBlock( - block: KtBlockExpression -): T? = firstOrNull { - val argumentExpression = it.getArgumentExpression() - val lambda = when (argumentExpression) { - is KtLambdaExpression -> argumentExpression - is KtLabeledExpression -> argumentExpression.baseExpression as? KtLambdaExpression - else -> null - } - lambda?.bodyExpression === block -} diff --git a/idea/testData/intentions/addLabeledReturnInLambda/multipleBlocks.kt b/idea/testData/intentions/addLabeledReturnInLambda/multipleBlocks.kt new file mode 100644 index 00000000000..d2bec68fce9 --- /dev/null +++ b/idea/testData/intentions/addLabeledReturnInLambda/multipleBlocks.kt @@ -0,0 +1,12 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Add return@find" + +fun foo() { + listOf(1,2,3).find { + if (it > 0) { + true + } else { + false + } + } +} \ No newline at end of file diff --git a/idea/testData/intentions/addLabeledReturnInLambda/multipleBlocks.kt.after b/idea/testData/intentions/addLabeledReturnInLambda/multipleBlocks.kt.after new file mode 100644 index 00000000000..e6a83848e49 --- /dev/null +++ b/idea/testData/intentions/addLabeledReturnInLambda/multipleBlocks.kt.after @@ -0,0 +1,12 @@ +// WITH_RUNTIME +// INTENTION_TEXT: "Add return@find" + +fun foo() { + listOf(1,2,3).find { + if (it > 0) { + return@find true + } else { + false + } + } +} \ No newline at end of file diff --git a/idea/testData/intentions/removeLabeledReturnInLambda/multipleBlocks.kt b/idea/testData/intentions/removeLabeledReturnInLambda/multipleBlocks.kt new file mode 100644 index 00000000000..fcbd4f86f99 --- /dev/null +++ b/idea/testData/intentions/removeLabeledReturnInLambda/multipleBlocks.kt @@ -0,0 +1,11 @@ +// WITH_RUNTIME + +fun foo() { + listOf(1,2,3).find { + if (it > 0) { + return@find true + } else { + false + } + } +} \ No newline at end of file diff --git a/idea/testData/intentions/removeLabeledReturnInLambda/multipleBlocks.kt.after b/idea/testData/intentions/removeLabeledReturnInLambda/multipleBlocks.kt.after new file mode 100644 index 00000000000..046e402b151 --- /dev/null +++ b/idea/testData/intentions/removeLabeledReturnInLambda/multipleBlocks.kt.after @@ -0,0 +1,11 @@ +// WITH_RUNTIME + +fun foo() { + listOf(1,2,3).find { + if (it > 0) { + true + } else { + false + } + } +} \ No newline at end of file diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java index 2113f8b55c2..d9d5cd73144 100644 --- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java +++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java @@ -390,6 +390,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest { doTest(fileName); } + @TestMetadata("multipleBlocks.kt") + public void testMultipleBlocks() throws Exception { + String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addLabeledReturnInLambda/multipleBlocks.kt"); + doTest(fileName); + } + @TestMetadata("multipleHighOrderFun.kt") public void testMultipleHighOrderFun() throws Exception { String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/addLabeledReturnInLambda/multipleHighOrderFun.kt"); @@ -13305,6 +13311,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest { doTest(fileName); } + @TestMetadata("multipleBlocks.kt") + public void testMultipleBlocks() throws Exception { + String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/removeLabeledReturnInLambda/multipleBlocks.kt"); + doTest(fileName); + } + @TestMetadata("normal.kt") public void testNormal() throws Exception { String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/removeLabeledReturnInLambda/normal.kt");