Find Usages: Fix processing of label references in 'return' expressions

#KT-7516 Fixed

idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getAssignmentByLHS
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.isReallySuccess
@@ -83,6 +84,16 @@ fun PsiReference.matchesTarget(candidateTarget: PsiElement): Boolean {
 
     val targets = unwrappedTargets
     if (unwrappedCandidate in targets) return true
+
+    if (element is KtLabelReferenceExpression && (element.parent as? KtContainerNode)?.parent is KtReturnExpression) {
+        targets.forEach {
+            if (it !is KtFunctionLiteral && !(it is KtNamedFunction && it.name.isNullOrEmpty())) return@forEach
+
+            val calleeReference = (it as KtFunction).getCalleeByLambdaArgument()?.mainReference ?: return@forEach
+            if (calleeReference.matchesTarget(candidateTarget)) return true
+        }
+    }
+
     // TODO: Investigate why PsiCompiledElement identity changes
     if (unwrappedCandidate is PsiCompiledElement && targets.any { it.isEquivalentTo(unwrappedCandidate) }) return true
 
@@ -197,4 +208,13 @@ fun KtReference.canBeResolvedViaImport(target: DeclarationDescriptor): Boolean {
     if (CallTypeAndReceiver.detect(referenceExpression).receiver != null) return false
     if (element.parent is KtThisExpression || element.parent is KtSuperExpression) return false // TODO: it's a bad design of PSI tree, we should change it
     return true
+}
+
+fun KtFunction.getCalleeByLambdaArgument(): KtSimpleNameExpression? {
+    val argument = getParentOfTypeAndBranch<KtValueArgument> { getArgumentExpression() } ?: return null
+    val callExpression = when (argument) {
+                             is KtLambdaArgument -> argument.parent as? KtCallExpression
+                             else -> (argument.parent as? KtValueArgumentList)?.parent as? KtCallExpression
+                         } ?: return null
+    return callExpression.calleeExpression as? KtSimpleNameExpression
 }

idea/testData/findUsages/kotlin/findFunctionUsages/labeledReturns.0.kt

@@ -0,0 +1,12 @@
+// PSI_ELEMENT: org.jetbrains.kotlin.psi.KtNamedFunction
+// OPTIONS: usages
+
+fun <R> <caret>foo(f: () -> R) = f()
+
+fun test() {
+    foo {
+        return@foo false
+    }
+
+    foo(fun(): Boolean { return@foo false })
+}

idea/testData/findUsages/kotlin/findFunctionUsages/labeledReturns.results.txt

@@ -0,0 +1,4 @@
+Function call 11 foo(fun(): Boolean { return@foo false })
+Function call 7 foo {
+Unclassified usage 11 foo(fun(): Boolean { return@foo false })
+Unclassified usage 8 return@foo false

idea/tests/org/jetbrains/kotlin/findUsages/FindUsagesTestGenerated.java

@@ -778,6 +778,12 @@ public class FindUsagesTestGenerated extends AbstractFindUsagesTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("labeledReturns.0.kt")
+            public void testLabeledReturns() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/findUsages/kotlin/findFunctionUsages/labeledReturns.0.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("localClassMember.0.kt")
             public void testLocalClassMember() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/findUsages/kotlin/findFunctionUsages/localClassMember.0.kt");