Make add / remove labeled return: extract 'getParentLambdaLabelName'

Related to KT-20439

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

@@ -10,7 +10,6 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
 
 class AddLabeledReturnInLambdaIntention : SelfTargetingRangeIntention<KtBlockExpression>(
@@ -19,14 +18,13 @@ class AddLabeledReturnInLambdaIntention : SelfTargetingRangeIntention<KtBlockExp
 ), LowPriorityAction {
     override fun applicabilityRange(element: KtBlockExpression): TextRange? {
         if (!isApplicableTo(element)) return null
-        val labelName = createLabelName(element) ?: return null
+        val labelName = element.getParentLambdaLabelName() ?: return null
         text = "Add return@$labelName"
         return element.statements.lastOrNull()?.textRange
     }
 
     override fun applyTo(element: KtBlockExpression, editor: Editor?) {
-        if (!isApplicableTo(element)) return
+        val labelName = element.getParentLambdaLabelName() ?: return
-        val labelName = createLabelName(element) ?: return
         val lastStatement = element.statements.lastOrNull() ?: return
         val newExpression = KtPsiFactory(element.project).createExpressionByPattern("return@$labelName $0", lastStatement)
         lastStatement.replace(newExpression)
@@ -36,14 +34,4 @@ class AddLabeledReturnInLambdaIntention : SelfTargetingRangeIntention<KtBlockExp
         val lastStatement = block.statements.lastOrNull()
         return lastStatement !is KtReturnExpression && lastStatement?.isUsedAsExpression(lastStatement.analyze()) == true
     }
-
-    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

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.idea.intentions
 import com.intellij.codeInsight.intention.LowPriorityAction
 import com.intellij.openapi.editor.Editor
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getCallNameExpression
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 
 class RemoveLabeledReturnInLambdaIntention : SelfTargetingIntention<KtReturnExpression>(
@@ -19,13 +18,7 @@ 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 lambdaExpression = block.getStrictParentOfType<KtLambdaExpression>() ?: return false
+        val callName = block.getParentLambdaLabelName() ?: 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
         text = "Remove return@$labelName"
         return true

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

@@ -30,6 +30,8 @@ import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getCallNameExpression
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.CollectionLiteralResolver
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
@@ -294,3 +296,13 @@ fun KtCallExpression.isArrayOfMethod(): Boolean {
     return (descriptor.containingDeclaration as? PackageFragmentDescriptor)?.fqName == KotlinBuiltIns.BUILT_INS_PACKAGE_FQ_NAME &&
            ARRAY_OF_METHODS.contains(descriptor.name)
 }
+
+fun KtBlockExpression.getParentLambdaLabelName(): String? {
+    val lambdaExpression = 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
+}