[analysis api] optimize search for declarations in IDE

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirDeserializedDeclarationSourceProvider.kt

@@ -67,7 +67,7 @@ internal object FirDeserializedDeclarationSourceProvider {
         classByClassId(klass.symbol.classId, klass.scope(project), project)
 
     private fun provideSourceForTypeAlias(alias: FirTypeAlias, project: Project): PsiElement? {
-        val candidates = project.createDeclarationProvider(alias.scope(project)).getTypeAliasesByClassId(alias.symbol.classId)
+        val candidates = project.createDeclarationProvider(alias.scope(project)).getAllTypeAliasesByClassId(alias.symbol.classId)
         return candidates.firstOrNull(KtElement::isCompiled)
     }
 
@@ -120,7 +120,7 @@ internal object FirDeserializedDeclarationSourceProvider {
     private fun classByClassId(classId: ClassId, scope: GlobalSearchScope, project: Project): KtClassOrObject? {
         val correctedClassId = classIdMapping[classId] ?: classId
         return project.createDeclarationProvider(scope)
-            .getClassesByClassId(correctedClassId)
+            .getAllClassesByClassId(correctedClassId)
             .firstOrNull(KtElement::isCompiled)
     }
 

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/KotlinDeclarationProvider.kt

@@ -19,8 +19,10 @@ import org.jetbrains.kotlin.psi.*
  * May be called frequently, so for implementations it is better to cache results.
  */
 public abstract class KotlinDeclarationProvider {
-    public abstract fun getClassesByClassId(classId: ClassId): Collection<KtClassOrObject>
+    public abstract fun getClassLikeDeclarationByClassId(classId: ClassId): KtClassLikeDeclaration?
-    public abstract fun getTypeAliasesByClassId(classId: ClassId): Collection<KtTypeAlias>
+
+    public abstract fun getAllClassesByClassId(classId: ClassId): Collection<KtClassOrObject>
+    public abstract fun getAllTypeAliasesByClassId(classId: ClassId): Collection<KtTypeAlias>
 
     public abstract fun getClassNamesInPackage(packageFqName: FqName): Set<Name>
     public abstract fun getTypeAliasNamesInPackage(packageFqName: FqName): Set<Name>

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticDeclarationProvider.kt

@@ -32,14 +32,19 @@ public class KotlinStaticDeclarationProvider internal constructor(
     private val KtElement.inScope: Boolean
         get() = containingKtFile.virtualFile in scope
 
-    override fun getClassesByClassId(classId: ClassId): Collection<KtClassOrObject> =
+    override fun getClassLikeDeclarationByClassId(classId: ClassId): KtClassLikeDeclaration? {
+        return getAllClassesByClassId(classId).firstOrNull()
+            ?: getAllTypeAliasesByClassId(classId).firstOrNull()
+    }
+
+    override fun getAllClassesByClassId(classId: ClassId): Collection<KtClassOrObject> =
         index.classMap[classId.packageFqName]
             ?.filter { ktClassOrObject ->
                 ktClassOrObject.getClassId() == classId && ktClassOrObject.inScope
             }
             ?: emptyList()
 
-    override fun getTypeAliasesByClassId(classId: ClassId): Collection<KtTypeAlias> =
+    override fun getAllTypeAliasesByClassId(classId: ClassId): Collection<KtTypeAlias> =
         index.typeAliasMap[classId.packageFqName]
             ?.filter { ktTypeAlias ->
                 ktTypeAlias.getClassId() == classId && ktTypeAlias.inScope

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdePredicateBasedProvider.kt

@@ -158,7 +158,7 @@ internal class LLFirIdePredicateBasedProvider(
         val directAnnotations = annotationsOnDeclaration(declaration)
         val metaAnnotations = directAnnotations
             .asSequence()
-            .mapNotNull { declarationProvider.getClassesByClassId(ClassId.topLevel(it)).singleOrNull() }
+            .mapNotNull { declarationProvider.getAllClassesByClassId(ClassId.topLevel(it)).singleOrNull() }
             .flatMap { annotationsResolver.annotationsOnDeclaration(it) }
             .toSet()
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt

@@ -35,10 +35,9 @@ internal class LLFirProviderHelper(
         if (classId.isLocal) return null
         return executeOrReturnDefaultValueOnPCE(null) {
             cache.classifierByClassId.computeIfAbsent(classId) {
-                val ktClass = when (val klass = declarationProvider.getClassesByClassId(classId).firstOrNull()) {
+                val ktClass = declarationProvider.getClassLikeDeclarationByClassId(classId)
-                    null -> declarationProvider.getTypeAliasesByClassId(classId).firstOrNull()
+                    ?: return@computeIfAbsent Optional.empty()
-                    else -> if (klass.getClassId() == null) null else klass
+                if (ktClass.getClassId() == null) return@computeIfAbsent Optional.empty()
-                } ?: return@computeIfAbsent Optional.empty()
                 val firFile = firFileBuilder.buildRawFirFileWithCaching(ktClass.containingKtFile)
                 val classifier = FirElementFinder.findElementIn<FirClassLikeDeclaration>(firFile) { classifier ->
                     classifier.symbol.classId == classId

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/KotlinAsJavaFirSupport.kt

@@ -34,7 +34,7 @@ class KotlinAsJavaFirSupport(private val project: Project) : KotlinAsJavaSupport
         searchScope: GlobalSearchScope
     ): Collection<KtClassOrObject> = project.createDeclarationProvider(searchScope).run {
         getClassNamesInPackage(packageFqName).flatMap {
-            getClassesByClassId(ClassId.topLevel(packageFqName.child(it)))
+            getAllClassesByClassId(ClassId.topLevel(packageFqName.child(it)))
         }
     }
 
@@ -66,7 +66,7 @@ class KotlinAsJavaFirSupport(private val project: Project) : KotlinAsJavaSupport
 
     override fun findClassOrObjectDeclarations(fqName: FqName, searchScope: GlobalSearchScope): Collection<KtClassOrObject> =
         fqName.toClassIdSequence().flatMap {
-            project.createDeclarationProvider(searchScope).getClassesByClassId(it)
+            project.createDeclarationProvider(searchScope).getAllClassesByClassId(it)
         }
             .filter { it.isFromSourceOrLibraryBinary(project) }
             .toSet()