KotlinAnnotatedElementsSearcher optimization

When possible, try to resolve references in types using only the PSI
without building the full resolve session.

idea/idea-analysis/src/org/jetbrains/kotlin/idea/search/PsiBasedClassResolver.kt

@@ -0,0 +1,242 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.search
+
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiModifier
+import com.intellij.psi.search.PsiShortNamesCache
+import org.jetbrains.annotations.TestOnly
+import org.jetbrains.kotlin.idea.caches.resolve.getNullableModuleInfo
+import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.resolve.frontendService
+import org.jetbrains.kotlin.idea.stubindex.KotlinTypeAliasShortNameIndex
+import org.jetbrains.kotlin.idea.util.application.runReadAction
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtAnnotationEntry
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+import org.jetbrains.kotlin.psi.KtUserType
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.resolve.lazy.DefaultImportProvider
+import java.util.concurrent.atomic.AtomicInteger
+
+/**
+ * Can quickly check whether a short name reference in a given file can resolve to the class/interface/type alias
+ * with the given qualified name.
+ */
+class PsiBasedClassResolver @TestOnly constructor(private val targetClassFqName: String) {
+    private val targetShortName = targetClassFqName.substringAfterLast('.')
+    private val targetPackage = targetClassFqName.substringBeforeLast('.', "")
+    /**
+     * Qualified names of packages which contain classes with the same short name as the target class.
+     */
+    private val conflictingPackages = mutableListOf<String>()
+    /**
+     * Qualified names of packages which contain typealiases with the same short name as the target class
+     * (which may or may not resolve to the target class).
+     */
+    private val packagesWithTypeAliases = mutableListOf<String>()
+    private var forceAmbiguity: Boolean = false
+    private var forceAmbiguityForNonAnnotations: Boolean = false
+
+    companion object {
+        @TestOnly val attempts = AtomicInteger()
+        @TestOnly val trueHits = AtomicInteger()
+        @TestOnly val falseHits = AtomicInteger()
+    }
+
+    constructor(target: PsiClass): this(target.qualifiedName ?: "") {
+        if (target.qualifiedName == null || target.containingClass != null || targetPackage.isEmpty()) {
+            forceAmbiguity = true
+            return
+        }
+
+        runReadAction {
+            findPotentialClassConflicts(target)
+            findPotentialTypeAliasConflicts(target)
+        }
+    }
+
+    private fun findPotentialClassConflicts(target: PsiClass) {
+        val candidates = PsiShortNamesCache.getInstance(target.project).getClassesByName(targetShortName, target.project.allScope())
+        for (candidate in candidates) {
+            // An inner class can be referenced by short name in subclasses without an explicit import
+            if (candidate.containingClass != null && !candidate.hasModifierProperty(PsiModifier.PRIVATE)) {
+                if (candidate.isAnnotationType) {
+                    forceAmbiguity = true
+                }
+                else {
+                    forceAmbiguityForNonAnnotations = true
+                }
+                break
+            }
+
+            if (candidate.qualifiedName == target.qualifiedName) {
+                // File with same FQ name in another module, don't bother with analyzing dependencies
+                if (candidate.navigationElement.containingFile != target.navigationElement.containingFile) {
+                    forceAmbiguity = true
+                    break
+                }
+            }
+            else {
+                candidate.qualifiedName?.substringBeforeLast('.', "")?.let { candidatePackage ->
+                    if (candidatePackage == "")
+                        forceAmbiguity = true
+                    else
+                        conflictingPackages.add(candidatePackage)
+                }
+            }
+        }
+    }
+
+    private fun findPotentialTypeAliasConflicts(target: PsiClass) {
+        val candidates = KotlinTypeAliasShortNameIndex.getInstance().get(targetShortName, target.project, target.project.allScope())
+        for (candidate in candidates) {
+            packagesWithTypeAliases.add(candidate.containingKtFile.packageFqName.asString())
+        }
+    }
+
+    @TestOnly fun addConflict(fqName: String) {
+        conflictingPackages.add(fqName.substringBeforeLast('.'))
+    }
+
+    /**
+     * Checks if a reference with the short name of [targetClassFqName] in the given file will resolve
+     * to the target class.
+     *
+     * @return true if it will definitely resolve to that class, false if it will definitely resolve to something else,
+     * null if full resolve is required to answer that question.
+     */
+    fun canBeTargetReference(ref: KtSimpleNameExpression): Boolean? {
+        attempts.incrementAndGet()
+        // The names can be different if the target was imported via an import alias
+        if (ref.getReferencedName() != targetShortName) {
+            return null
+        }
+
+        // Names in expressions can conflict with local declarations and methods of implicit receivers,
+        // so we can't find out what they refer to without a full resolve.
+        val userType = ref.getStrictParentOfType<KtUserType>() ?: return null
+        if (forceAmbiguityForNonAnnotations && userType.getParentOfTypeAndBranch<KtAnnotationEntry> { typeReference } == null) {
+            return null
+        }
+
+        if (forceAmbiguity) {
+            return null
+        }
+
+        val qualifiedCheckResult = checkQualifiedReferenceToTarget(ref)
+        if (qualifiedCheckResult != null) {
+            return qualifiedCheckResult.returnValue
+        }
+
+        val file = ref.containingKtFile
+        var result: Result = Result.NothingFound
+        val filePackage = file.packageFqName.asString()
+        if (filePackage == targetPackage) {
+            result = result.changeTo(Result.Found)
+        }
+        else if (filePackage in conflictingPackages) {
+            result = result.changeTo(Result.FoundOther)
+        }
+        else if (filePackage in packagesWithTypeAliases) {
+            return null
+        }
+
+        if (file.getNullableModuleInfo() != null) {  // it can be null in tests
+            val defaultImportProvider = file.getResolutionFacade().frontendService<DefaultImportProvider>()
+            for (importPath in defaultImportProvider.defaultImports) {
+                result = analyzeSingleImport(result, importPath.fqName, importPath.isAllUnder, importPath.alias?.asString())
+                if (result == Result.Ambiguity) return null
+            }
+        }
+
+        for (importDirective in file.importDirectives) {
+            result = analyzeSingleImport(result, importDirective.importedFqName, importDirective.isAllUnder, importDirective.aliasName)
+            if (result == Result.Ambiguity) return null
+        }
+
+        if (result.returnValue == true) {
+            trueHits.incrementAndGet()
+        }
+        else if (result.returnValue == false) {
+            falseHits.incrementAndGet()
+        }
+        return result.returnValue
+    }
+
+    private fun analyzeSingleImport(result: Result, importedFqName: FqName?, isAllUnder: Boolean, aliasName: String?): Result {
+        val qName = importedFqName
+        if (!isAllUnder) {
+            if (qName?.asString() == targetClassFqName &&
+                (aliasName == null || aliasName == targetShortName)) {
+                return result.changeTo(Result.Found)
+            }
+            else if (qName?.shortName()?.asString() == targetShortName &&
+                     qName.parent().asString() in conflictingPackages &&
+                     aliasName == null) {
+                return result.changeTo(Result.FoundOther)
+            }
+            else if (qName?.shortName()?.asString() == targetShortName &&
+                     qName.parent().asString() in packagesWithTypeAliases &&
+                     aliasName == null) {
+                return Result.Ambiguity
+            }
+            else if (aliasName == targetShortName) {
+                return result.changeTo(Result.FoundOther)
+            }
+        }
+        else {
+            if (qName?.asString() == targetPackage) {
+                return result.changeTo(Result.Found)
+            }
+            else if (qName?.asString() in conflictingPackages) {
+                return result.changeTo(Result.FoundOther)
+            }
+            else if (qName?.asString() in packagesWithTypeAliases) {
+                return Result.Ambiguity
+            }
+        }
+        return result
+    }
+
+    private fun checkQualifiedReferenceToTarget(ref: KtSimpleNameExpression): Result? {
+        // A qualified name can resolve to the target element even if it's not imported,
+        // but it can also resolve to something else e.g. if the file defines a class with the same name
+        // as the top-level package of the target class.
+        val qualifier = (ref.parent as? KtUserType)?.qualifier
+        if (qualifier != null) {
+            if (qualifier.text == targetPackage) return Result.Ambiguity
+            return Result.FoundOther
+        }
+        return null
+    }
+
+    enum class Result(val returnValue: Boolean?) {
+        NothingFound(false),
+        Found(true),
+        FoundOther(false),
+        Ambiguity(null)
+    }
+
+    fun Result.changeTo(newResult: Result): Result {
+        if (this == Result.NothingFound || this.returnValue == newResult.returnValue) {
+            return newResult
+        }
+        return Result.Ambiguity
+    }
+}

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

@@ -31,6 +31,7 @@ import com.intellij.util.indexing.FileBasedIndex
 import org.jetbrains.kotlin.asJava.LightClassUtil
 import org.jetbrains.kotlin.asJava.toLightClass
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.search.PsiBasedClassResolver
 import org.jetbrains.kotlin.idea.stubindex.KotlinAnnotationsIndex
 import org.jetbrains.kotlin.idea.stubindex.KotlinSourceFilterScope
 import org.jetbrains.kotlin.idea.util.application.runReadAction
@@ -71,13 +72,13 @@ class KotlinAnnotatedElementsSearcher : QueryExecutor<PsiModifierListOwner, Anno
         private val LOG = Logger.getInstance("#com.intellij.psi.impl.search.AnnotatedMembersSearcher")
 
         fun processAnnotatedMembers(annClass: PsiClass,
-                                           useScope: SearchScope,
-                                           preFilter: (KtAnnotationEntry) -> Boolean = { true },
-                                           consumer: (KtDeclaration) -> Boolean): Boolean {
+                                    useScope: SearchScope,
+                                    preFilter: (KtAnnotationEntry) -> Boolean = { true },
+                                    consumer: (KtDeclaration) -> Boolean): Boolean {
             assert(annClass.isAnnotationType) { "Annotation type should be passed to annotated members search" }
 
-            val annotationFQN = annClass.qualifiedName
-            assert(annotationFQN != null)
+            val psiBasedClassResolver = PsiBasedClassResolver(annClass)
+            val annotationFQN = annClass.qualifiedName!!
 
             val candidates = getKotlinAnnotationCandidates(annClass, useScope)
             for (elt in candidates) {
@@ -89,11 +90,18 @@ class KotlinAnnotatedElementsSearcher : QueryExecutor<PsiModifierListOwner, Anno
 
                     val declaration = elt.getStrictParentOfType<KtDeclaration>() ?: return true
 
-                    val context = elt.analyze(BodyResolveMode.PARTIAL)
-                    val annotationDescriptor = context.get(BindingContext.ANNOTATION, elt) ?: return true
+                    val psiBasedResolveResult = elt.calleeExpression?.constructorReferenceExpression?.let { ref ->
+                        psiBasedClassResolver.canBeTargetReference(ref)
+                    }
 
-                    val descriptor = annotationDescriptor.type.constructor.declarationDescriptor ?: return true
-                    if (!(DescriptorUtils.getFqName(descriptor).asString() == annotationFQN)) return true
+                    if (psiBasedResolveResult == false) return true
+                    if (psiBasedResolveResult == null) {
+                        val context = elt.analyze(BodyResolveMode.PARTIAL)
+                        val annotationDescriptor = context.get(BindingContext.ANNOTATION, elt) ?: return true
+
+                        val descriptor = annotationDescriptor.type.constructor.declarationDescriptor ?: return true
+                        if (DescriptorUtils.getFqName(descriptor).asString() != annotationFQN) return true
+                    }
 
                     if (!consumer(declaration)) return false
 

idea/testData/search/annotations/testAmbiguousNestedNonAnnotationClass.kt

@@ -0,0 +1,13 @@
+package foo
+
+annotation class Xyzzy
+
+object Foo {
+    class Xyzzy
+}
+
+@Xyzzy fun foo()
+
+// ANNOTATION: foo.Xyzzy
+// SEARCH: method:foo
+// OPTIMIZED_TRUE: 1

idea/testData/search/annotations/testAmbiguousNestedPrivateAnnotationClass.kt

@@ -0,0 +1,13 @@
+package foo
+
+annotation class Test
+
+object Foo {
+    private annotation class Test
+}
+
+@Test fun foo()
+
+// ANNOTATION: foo.Test
+// SEARCH: method:foo
+// OPTIMIZED_TRUE: 1

idea/testData/search/annotations/testDefaultImport.kt

@@ -0,0 +1,5 @@
+@Suppress fun shazam(): Nothing = TODO()
+
+// ANNOTATION: kotlin.Suppress
+// SEARCH: method:shazam
+// OPTIMIZED_TRUE: 1

idea/testData/search/annotations/testNestedClassAsAnnotation.kt

@@ -0,0 +1,16 @@
+package foo.bar
+
+class A
+
+object F {
+    class foo {
+        class bar {
+            annotation class A
+
+            @foo.bar.A fun test(): Nothing = TODO()
+        }
+    }
+}
+
+// ANNOTATION: foo.bar.F.foo.bar.A
+// SEARCH: method:test

idea/testData/search/annotations/testNestedPrivateAnnotationClass.kt

@@ -0,0 +1,12 @@
+package foo
+
+class A {
+    private annotation class B
+
+    class C {
+        @B
+        fun foo() {}
+    }
+}
+// ANNOTATION: foo.A.B
+// SEARCH: method:foo

idea/testData/search/annotations/testTypeAlias.kt

@@ -0,0 +1,6 @@
+typealias Suppress = kotlin.Suppress
+
+@Suppress fun shazam(): Nothing = TODO()
+
+// ANNOTATION: kotlin.Suppress
+// SEARCH: method:shazam

idea/tests/org/jetbrains/kotlin/idea/search/PsiBasedClassResolverTest.kt

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.search
+
+import com.intellij.psi.PsiFileFactory
+import org.jetbrains.kotlin.idea.KotlinFileType
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCaseBase
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.junit.Assert
+
+class PsiBasedClassResolverTest : KotlinLightCodeInsightFixtureTestCaseBase() {
+    fun testSimple() {
+        Assert.assertFalse(doTest("@Test fun foo()")!!)
+    }
+
+    fun testImported() {
+        Assert.assertTrue(doTest("import org.junit.Test\n\n@Test fun foo()")!!)
+    }
+
+    fun testCallReference() {
+        Assert.assertNull(doTest("import org.junit.Test\n\nval c = Test()", "Test()"))
+    }
+
+    fun testImportStar() {
+        Assert.assertTrue(doTest("import org.junit.*\n\n@Test fun foo()")!!)
+    }
+
+    fun testImportConflict() {
+        Assert.assertFalse(doTest("import org.testng.Test\n\n@Test fun foo()")!!)
+    }
+
+    fun testImportConflictOtherName() {
+        Assert.assertTrue(doTest("import org.junit.Test\nimport org.testng.NotTest\n\n@Test fun foo()")!!)
+    }
+
+    fun testImportStarConflict() {
+        Assert.assertFalse(doTest("import org.testng.*\n\n@Test fun foo()")!!)
+    }
+
+    fun testImportBoth() {
+        Assert.assertNull(doTest("import org.testng.*\nimport org.junit.Test\n@Test fun foo()"))
+    }
+
+    fun testImportBoth2() {
+        Assert.assertNull(doTest("import org.testng.Test\nimport org.junit.*\n@Test fun foo()"))
+    }
+
+    fun testImportAs() {
+        Assert.assertNull(doTest("import org.junit.*\nimport java.lang.String as Test\n@Test fun foo()"))
+    }
+
+    fun testImportAsSameName() {
+        Assert.assertTrue(doTest("import org.junit.Test as Test\n@Test fun foo()")!!)
+    }
+
+    fun testThisPackage() {
+        Assert.assertTrue(doTest("package org.junit\n\n@Test fun foo()")!!)
+    }
+
+    fun testConflictPackageAmbiguity() {
+        Assert.assertNull(doTest("package org.testng\n\nimport org.junit.Test\n\n@Test fun foo()"))
+    }
+
+    fun testFqName() {
+        Assert.assertNull(doTest("val c = org.junit.Test()", "Test"))
+    }
+
+    fun testFqNameAnn() {
+        Assert.assertNull(doTest("@org.junit.Test fun foo() { }", "@org.junit.Test"))
+    }
+
+    fun testOtherFqNameAnn() {
+        Assert.assertFalse(doTest("import org.junit.Test\n@foo.bar.Test fun foo() { }", "@foo.bar.Test")!!)
+    }
+
+    fun testFqNameWithConflictingImport() {
+        Assert.assertNull(doTest("import org.testng.Test\n\nval c = org.junit.Test()", "org.junit.Test"))
+    }
+
+    fun testUserType() {
+        Assert.assertTrue(doTest("import org.junit.Test\nclass MyTest : Test()", "Test()")!!)
+    }
+
+    fun testUserTypeQualified() {
+        Assert.assertNull(doTest("class MyTest : org.junit.Test()", "Test()"))
+    }
+
+    private fun doTest(fileText: String, marker: String = "@Test"): Boolean? {
+        val resolver = PsiBasedClassResolver("org.junit.Test")
+        resolver.addConflict("org.testng.Test")
+        val file = PsiFileFactory.getInstance(project).createFileFromText("foo.kt", KotlinFileType.INSTANCE,
+                                                                          fileText) as KtFile
+        val index = fileText.indexOf(marker)
+        val ref = file.findElementAt(index + marker.indexOf("Test"))!!.getParentOfType<KtSimpleNameExpression>(false)!!
+        return resolver.canBeTargetReference(ref)
+    }
+}

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

@@ -19,6 +19,10 @@ package org.jetbrains.kotlin.search;
 import com.intellij.openapi.util.io.FileUtil;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.search.searches.AnnotatedMembersSearch;
+import com.intellij.testFramework.LightProjectDescriptor;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.idea.search.PsiBasedClassResolver;
+import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor;
 import org.jetbrains.kotlin.idea.test.PluginTestCaseBase;
 import org.jetbrains.kotlin.test.InTextDirectivesUtils;
 
@@ -27,6 +31,11 @@ import java.io.IOException;
 import java.util.List;
 
 public class AnnotatedMembersSearchTest extends AbstractSearcherTest {
+    @NotNull
+    @Override
+    protected LightProjectDescriptor getProjectDescriptor() {
+        return KotlinWithJdkAndRuntimeLightProjectDescriptor.INSTANCE;
+    }
 
     public void testAnnotationsOnClass() throws IOException {
         doTest();
@@ -40,14 +49,51 @@ public class AnnotatedMembersSearchTest extends AbstractSearcherTest {
         doTest();
     }
 
+    public void testNestedClassAsAnnotation() throws IOException {
+        doTest();
+    }
+
+    public void testAmbiguousNestedNonAnnotationClass() throws IOException {
+        doTest();
+    }
+
+    public void testAmbiguousNestedPrivateAnnotationClass() throws IOException {
+        doTest();
+    }
+
+    public void testNestedPrivateAnnotationClass() throws IOException {
+        doTest();
+    }
+
+    public void testTypeAlias() throws IOException {
+        doTest();
+    }
+
+    public void testDefaultImport() throws IOException {
+        doTest();
+    }
+
     private void doTest() throws IOException {
         myFixture.configureByFile(getFileName());
-        List<String> directives = InTextDirectivesUtils.findListWithPrefixes(FileUtil.loadFile(new File(getPathToFile()), true),
+        String fileText = FileUtil.loadFile(new File(getPathToFile()), true);
+        List<String> directives = InTextDirectivesUtils.findListWithPrefixes(fileText,
                                                                              "// ANNOTATION: ");
         assertFalse("Specify ANNOTATION directive in test file", directives.isEmpty());
         String annotationClassName = directives.get(0);
         PsiClass psiClass = getPsiClass(annotationClassName);
+        PsiBasedClassResolver.Companion.getTrueHits().set(0);
+        PsiBasedClassResolver.Companion.getFalseHits().set(0);
+
         checkResult(AnnotatedMembersSearch.search(psiClass, getProjectScope()));
+
+        Integer optimizedTrue = InTextDirectivesUtils.getPrefixedInt(fileText, "// OPTIMIZED_TRUE:");
+        if (optimizedTrue != null) {
+            assertEquals(optimizedTrue.intValue(), PsiBasedClassResolver.Companion.getTrueHits().get());
+        }
+        Integer optimizedFalse = InTextDirectivesUtils.getPrefixedInt(fileText, "// OPTIMIZED_FALSE:");
+        if (optimizedFalse != null) {
+            assertEquals(optimizedFalse.intValue(), PsiBasedClassResolver.Companion.getFalseHits().get());
+        }
     }
 
     @Override