From 5e8cd654ec73d4961a0cba652ff492d37fecdc86 Mon Sep 17 00:00:00 2001
From: Alexey Sedunov <alexey.sedunov@jetbrains.com>
Date: Fri, 17 Feb 2017 15:03:50 +0300
Subject: [PATCH] Find Usages: Fix processing of label references in 'return'
 expressions  #KT-7516 Fixed

---
 .../kotlin/idea/references/referenceUtil.kt   | 20 +++++++++++++++++++
 .../findFunctionUsages/labeledReturns.0.kt    | 12 +++++++++++
 .../labeledReturns.results.txt                |  4 ++++
 .../findUsages/FindUsagesTestGenerated.java   |  6 ++++++
 4 files changed, 42 insertions(+)
 create mode 100644 idea/testData/findUsages/kotlin/findFunctionUsages/labeledReturns.0.kt
 create mode 100644 idea/testData/findUsages/kotlin/findFunctionUsages/labeledReturns.results.txt

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt
index 00e3aae75e8..596be75a162 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt
+++ b/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
 }
\ No newline at end of file
diff --git a/idea/testData/findUsages/kotlin/findFunctionUsages/labeledReturns.0.kt b/idea/testData/findUsages/kotlin/findFunctionUsages/labeledReturns.0.kt
new file mode 100644
index 00000000000..351fa6094bf
--- /dev/null
+++ b/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 })
+}
\ No newline at end of file
diff --git a/idea/testData/findUsages/kotlin/findFunctionUsages/labeledReturns.results.txt b/idea/testData/findUsages/kotlin/findFunctionUsages/labeledReturns.results.txt
new file mode 100644
index 00000000000..604a640a343
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/findUsages/FindUsagesTestGenerated.java b/idea/tests/org/jetbrains/kotlin/findUsages/FindUsagesTestGenerated.java
index 8c81e564dcd..51cf5a2147f 100644
--- a/idea/tests/org/jetbrains/kotlin/findUsages/FindUsagesTestGenerated.java
+++ b/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");