Make add / remove labeled return working for hierarchical blocks

Related to KT-20439

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>(
-    KtLambdaExpression::class.java,
+class AddLabeledReturnInLambdaIntention : SelfTargetingRangeIntention<KtBlockExpression>(
+    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<KtCallExpression>() ?: return null
-        val valueArgument = callExpression.valueArguments.findArgumentWithGivenBlock(block) ?: return null
+    private fun createLabelName(block: KtBlockExpression): String? {
+        val lambdaExpression = block.getStrictParentOfType<KtLambdaExpression>() ?: return null
+        val callExpression = lambdaExpression.getStrictParentOfType<KtCallExpression>() ?: 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
     }

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

@@ -19,8 +19,11 @@ class RemoveLabeledReturnInLambdaIntention : SelfTargetingIntention<KtReturnExpr
         val labelName = element.getLabelName() ?: return false
         val block = element.getStrictParentOfType<KtBlockExpression>() ?: return false
         if (block.statements.lastOrNull() != element) return false
-        val callExpression = block.getStrictParentOfType<KtCallExpression>() ?: return false
-        val lambdaArgument = callExpression.lambdaArguments.findArgumentWithGivenBlock(block) ?: return false
+        val lambdaExpression = block.getStrictParentOfType<KtLambdaExpression>() ?: return false
+        val callExpression = lambdaExpression.getStrictParentOfType<KtCallExpression>() ?: 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

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 <T : KtValueArgument> List<T>.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
-}

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) {
+            <caret>true
+        } else {
+            false
+        }
+    }
+}

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
+        }
+    }
+}

idea/testData/intentions/removeLabeledReturnInLambda/multipleBlocks.kt

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+
+fun foo() {
+    listOf(1,2,3).find {
+        if (it > 0) {
+            return@find <caret>true
+        } else {
+            false
+        }
+    }
+}

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
+        }
+    }
+}

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");