From 74bf4a3a7f65c8220fcb441dc5b457312ec0ead1 Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Tue, 28 Apr 2015 12:12:25 +0300
Subject: [PATCH] Refine findElementOfClassAtRange to make it find labels

---
 .../kotlin/idea/codeInsight/CodeInsightUtils.java | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/codeInsight/CodeInsightUtils.java b/idea/src/org/jetbrains/kotlin/idea/codeInsight/CodeInsightUtils.java
index c146fc69985..74a9eb15426 100644
--- a/idea/src/org/jetbrains/kotlin/idea/codeInsight/CodeInsightUtils.java
+++ b/idea/src/org/jetbrains/kotlin/idea/codeInsight/CodeInsightUtils.java
@@ -110,8 +110,12 @@ public class CodeInsightUtils {
 
     @Nullable
     public static PsiElement findElementOfClassAtRange(@NotNull PsiFile file, int startOffset, int endOffset, Class<JetExpression> aClass) {
-        PsiElement element1 = getElementAtOffsetIgnoreWhitespaceBefore(file, startOffset);
-        PsiElement element2 = getElementAtOffsetIgnoreWhitespaceAfter(file, endOffset);
+        // When selected range is this@Fo<select>o</select> we'd like to return `@Foo`
+        // But it's PSI looks like: (AT IDENTIFIER):JetLabel
+        // So if we search parent starting exactly at IDENTIFIER then we find nothing
+        // Solution is to retrieve label if we are on AT or IDENTIFIER
+        PsiElement element1 = getParentLabelOrElement(getElementAtOffsetIgnoreWhitespaceBefore(file, startOffset));
+        PsiElement element2 = getParentLabelOrElement(getElementAtOffsetIgnoreWhitespaceAfter(file, endOffset));
 
         if (element1 == null || element2 == null) return null;
 
@@ -127,6 +131,13 @@ public class CodeInsightUtils {
         return jetExpression;
     }
 
+    private static PsiElement getParentLabelOrElement(@Nullable PsiElement element) {
+        if (element != null && element.getParent() instanceof JetLabelReferenceExpression) {
+            return element.getParent();
+        }
+        return element;
+    }
+
     @Nullable
     public static PsiElement[] findElementsOfClassInRange(@NotNull PsiFile file, int startOffset, int endOffset, Class<? extends PsiElement> aClass) {
         PsiElement element1 = getElementAtOffsetIgnoreWhitespaceBefore(file, startOffset);