KotlinAnnotatedElementsSearcher: handle local scope; correctly handle import aliases

compiler/frontend/src/org/jetbrains/kotlin/psi/stubs/StubInterfaces.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.psi.*
 public interface KotlinFileStub : PsiFileStub<JetFile> {
     public fun getPackageFqName(): FqName
     public fun isScript(): Boolean
+    public fun findImportByAlias(alias: String): KotlinImportDirectiveStub?
 }
 
 public interface KotlinPlaceHolderStub<T : JetElement> : StubElement<T>

compiler/frontend/src/org/jetbrains/kotlin/psi/stubs/impl/KotlinFileStubImpl.kt

@@ -25,6 +25,7 @@ import com.intellij.util.io.StringRef
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.JetFile
 import org.jetbrains.kotlin.psi.stubs.KotlinFileStub
+import org.jetbrains.kotlin.psi.stubs.KotlinImportDirectiveStub
 import org.jetbrains.kotlin.psi.stubs.elements.JetStubElementTypes
 
 public class KotlinFileStubImpl(jetFile: JetFile?, private val packageName: StringRef, private val isScript: Boolean)// SCRIPT: PsiJetFileStubImpl knows about scripting
@@ -43,4 +44,9 @@ public class KotlinFileStubImpl(jetFile: JetFile?, private val packageName: Stri
     override fun getClasses(): Array<PsiClass> {
         return childrenStubs.filterIsInstance<PsiClassStub<*>>().map { it.psi }.toTypedArray()
     }
+
+    override fun findImportByAlias(alias: String): KotlinImportDirectiveStub? {
+        val importList = childrenStubs.firstOrNull { it.stubType == JetStubElementTypes.IMPORT_LIST } ?: return null
+        return importList.childrenStubs.filterIsInstance<KotlinImportDirectiveStub>().firstOrNull { it.getAliasName() == alias }
+    }
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/stubindex/JetAnnotationsIndex.java

@@ -46,4 +46,9 @@ public class JetAnnotationsIndex extends StringStubIndexExtension<JetAnnotationE
     public Collection<JetAnnotationEntry> get(String s, Project project, @NotNull GlobalSearchScope scope) {
         return super.get(s, project, JetSourceFilterScope.kotlinSourcesAndLibraries(scope, project));
     }
+
+    @Override
+    public int getVersion() {
+        return super.getVersion() + 1;
+    }
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/stubindex/StubIndexServiceImpl.java

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.stubindex;
 
 import com.intellij.psi.stubs.IndexSink;
+import com.intellij.psi.stubs.StubElement;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.psi.JetClassOrObject;
 import org.jetbrains.kotlin.psi.stubs.*;
@@ -123,5 +124,24 @@ public class StubIndexServiceImpl implements StubIndexService {
     @Override
     public void indexAnnotation(KotlinAnnotationEntryStub stub, IndexSink sink) {
         sink.occurrence(JetAnnotationsIndex.getInstance().getKey(), stub.getShortName());
+
+        KotlinFileStub fileStub = getContainingFileStub(stub);
+        if (fileStub != null) {
+            KotlinImportDirectiveStub aliasImportStub = fileStub.findImportByAlias(stub.getShortName());
+            if (aliasImportStub != null) {
+                sink.occurrence(JetAnnotationsIndex.getInstance().getKey(), aliasImportStub.getImportedFqName().shortName().asString());
+            }
+        }
+    }
+
+    private static KotlinFileStub getContainingFileStub(StubElement stub) {
+        StubElement parent = stub.getParentStub();
+        while (parent != null) {
+            if (parent instanceof KotlinFileStub) {
+                return (KotlinFileStub) parent;
+            }
+            parent = parent.getParentStub();
+        }
+        return null;
     }
 }

idea/src/org/jetbrains/kotlin/idea/search/ideaExtensions/KotlinAnnotatedElementsSearcher.kt

@@ -20,24 +20,26 @@ import com.intellij.openapi.diagnostic.Logger
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiModifierListOwner
-import com.intellij.psi.impl.search.AnnotatedElementsSearcher
 import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.LocalSearchScope
 import com.intellij.psi.search.SearchScope
 import com.intellij.psi.search.searches.AnnotatedElementsSearch
 import com.intellij.psi.util.PsiUtilCore
 import com.intellij.util.Processor
+import com.intellij.util.QueryExecutor
 import com.intellij.util.indexing.FileBasedIndex
 import org.jetbrains.kotlin.asJava.LightClassUtil
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.stubindex.JetAnnotationsIndex
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getChildrenOfType
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
-public class KotlinAnnotatedElementsSearcher : AnnotatedElementsSearcher() {
+public class KotlinAnnotatedElementsSearcher : QueryExecutor<PsiModifierListOwner, AnnotatedElementsSearch.Parameters> {
 
     override fun execute(p: AnnotatedElementsSearch.Parameters, consumer: Processor<PsiModifierListOwner>): Boolean {
         val annClass = p.getAnnotationClass()
@@ -86,18 +88,10 @@ public class KotlinAnnotatedElementsSearcher : AnnotatedElementsSearcher() {
             return runReadAction(fun(): Collection<PsiElement> {
                 if (useScope is GlobalSearchScope) {
                     val name = annClass.getName() ?: return emptyList()
-                    val annotationEntries = JetAnnotationsIndex.getInstance().get(name, annClass.getProject(), useScope)
-
-                    // Add annotations 'test' as often used alias when we search Test annotation
-                    if (name == "Test") {
-                        annotationEntries.addAll(JetAnnotationsIndex.getInstance().get(name.toLowerCase(), annClass.getProject(), useScope))
-                    }
-
-                    return annotationEntries
+                    return JetAnnotationsIndex.getInstance().get(name, annClass.getProject(), useScope)
                 }
 
-                // TODO getJetAnnotationCandidates works only with global search scope
-                return emptyList()
+                return (useScope as LocalSearchScope).scope.flatMap { it.getChildrenOfType<JetAnnotationEntry>().asIterable() }
             })
         }
 

idea/testData/search/junit/testJunit4FancyAlias.kt

@@ -1,10 +1,10 @@
-import org.junit.Test as test
+import org.junit.Test as unittest
 import org.junit.Test
 
 class MyTestClass {
-    test fun test1() {}
+    unittest fun test1() {}
 
-    @Deprecated @test fun test2() {}
+    @Deprecated @unittest fun test2() {}
 
     Test fun test3() {}
 }

idea/tests/org/jetbrains/kotlin/search/JUnitMembersSearcherTest.java

@@ -47,6 +47,10 @@ public class JUnitMembersSearcherTest extends AbstractSearcherTest {
         doJUnit4test();
     }
 
+    public void testJunit4FancyAlias() throws IOException {
+        doJUnit4test();
+    }
+
     private void doJUnit3test() throws IOException {
         myFixture.configureByFile(getFileName());
         List<String> directives = InTextDirectivesUtils.findListWithPrefixes(FileUtil.loadFile(new File(getPathToFile()), true), "// CLASS: ");