[LL FIR] KTIJ-24296 Implement symbol name pre-filtering in LLFirProvider

- The optimization reduces time spent in `LLFirProvider.SymbolProvider`
  during my tests by 30-90% in some highlighted files. In other
  performance tests, the optimization makes performance worse, so more
  work is needed. In general, higher workload tests run faster with the
  optimization, while lower workload tests run slower. This is expected
  as the optimization trades short-term performance (for building
  classifier/callable name sets) in favor of long-term speedup.
- The changes should also reduce the size of declaration caches like
  `LLFirProviderHelper.classifierByClassId`.
- Building the name sets has a second use, as a similar optimization can
  be implemented for dependent module providers, which may rely on the
  sets computed here.

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

@@ -108,40 +108,61 @@ internal class LLFirProvider(
 
     @NoMutableState
     private inner class SymbolProvider : FirSymbolProvider(session) {
-        override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> =
+        override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
-            providerHelper.getTopLevelCallableSymbols(packageFqName, name)
+            val classifierNamesInPackage = providerHelper.getTopLevelClassifierNamesInPackage(classId.packageFqName)
+            if (!mayHaveTopLevelClassifier(classId, classifierNamesInPackage.names, classifierNamesInPackage.mayHaveFunctionClass)) {
+                return null
+            }
+
+            return getFirClassifierByFqName(classId)?.symbol
+        }
+
+        override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
+            if (!mayHaveTopLevelCallable(packageFqName, name)) return emptyList()
+            return providerHelper.getTopLevelCallableSymbols(packageFqName, name)
+        }
 
         @FirSymbolProviderInternals
         override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
-            destination += getTopLevelCallableSymbols(packageFqName, name)
+            if (!mayHaveTopLevelCallable(packageFqName, name)) return
+            destination += providerHelper.getTopLevelCallableSymbols(packageFqName, name)
         }
 
-        override fun getTopLevelFunctionSymbols(packageFqName: FqName, name: Name): List<FirNamedFunctionSymbol> =
+        override fun getTopLevelFunctionSymbols(packageFqName: FqName, name: Name): List<FirNamedFunctionSymbol> {
-            providerHelper.getTopLevelFunctionSymbols(packageFqName, name)
+            if (!mayHaveTopLevelCallable(packageFqName, name)) return emptyList()
+            return providerHelper.getTopLevelFunctionSymbols(packageFqName, name)
+        }
 
         @FirSymbolProviderInternals
         override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
-            destination += getTopLevelFunctionSymbols(packageFqName, name)
+            if (!mayHaveTopLevelCallable(packageFqName, name)) return
+            destination += providerHelper.getTopLevelFunctionSymbols(packageFqName, name)
         }
 
-        override fun getTopLevelPropertySymbols(packageFqName: FqName, name: Name): List<FirPropertySymbol> =
+        override fun getTopLevelPropertySymbols(packageFqName: FqName, name: Name): List<FirPropertySymbol> {
-            providerHelper.getTopLevelPropertySymbols(packageFqName, name)
+            if (!mayHaveTopLevelCallable(packageFqName, name)) return emptyList()
+            return providerHelper.getTopLevelPropertySymbols(packageFqName, name)
+        }
 
         @FirSymbolProviderInternals
         override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
-            destination += getTopLevelPropertySymbols(packageFqName, name)
+            if (!mayHaveTopLevelCallable(packageFqName, name)) return
+            destination += providerHelper.getTopLevelPropertySymbols(packageFqName, name)
         }
 
+        private fun mayHaveTopLevelCallable(packageFqName: FqName, name: Name): Boolean =
+            name in computeCallableNamesInPackage(packageFqName)
+
         override fun getPackage(fqName: FqName): FqName? =
             providerHelper.getPackage(fqName)
 
-        // TODO: Consider having proper implementations for sake of optimizations
+        // Computing the set of such package names is expensive and would require a new index. For now, it is not worth the marginal gains.
         override fun computePackageSetWithTopLevelCallables(): Set<String>? = null
-        override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? = null
-        override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? = null
 
-        override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
+        override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String> =
-            return getFirClassifierByFqName(classId)?.symbol
+            providerHelper.getTopLevelClassifierNamesInPackage(packageFqName).names
-        }
+
+        override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name> =
+            providerHelper.getTopLevelCallableNamesInPackage(packageFqName)
     }
 }

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.types.functionTypeService
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -48,6 +49,25 @@ internal class LLFirProviderHelper(
                 ?: error("Classifier $classId was found in file ${ktClass.containingKtFile.virtualFilePath} but was not found in FirFile")
         }
 
+    internal class ClassifierNamesInPackage(
+        val names: Set<String>,
+
+        /**
+         * Some packages (such as `kotlin` or `kotlin.coroutines`) may contain function classes generated by the compiler. To avoid a hash
+         * map access when checking `mayHaveTopLevelClassifier`, [mayHaveFunctionClass] is precomputed for each package name. Also see
+         * [org.jetbrains.kotlin.builtins.functions.FunctionTypeKind].
+         */
+        val mayHaveFunctionClass: Boolean,
+    )
+
+    private val topLevelClassifierNamesByPackage =
+        firSession.firCachesFactory.createCache<FqName, ClassifierNamesInPackage> { packageFqName ->
+            val names = declarationProvider
+                .getTopLevelKotlinClassLikeDeclarationNamesInPackage(packageFqName)
+                .mapTo(mutableSetOf()) { it.asString() }
+            val mayHaveFunctionClass = firSession.functionTypeService.hasKindWithSpecificPackage(packageFqName)
+            ClassifierNamesInPackage(names, mayHaveFunctionClass)
+        }
 
     private val callablesByCallableId = firSession.firCachesFactory.createCache<CallableId, List<FirCallableSymbol<*>>> { callableId ->
         val files = declarationProvider.getTopLevelCallableFiles(callableId).ifEmpty { return@createCache emptyList() }
@@ -59,6 +79,11 @@ internal class LLFirProviderHelper(
         }
     }
 
+    private val topLevelCallableNamesByPackage =
+        firSession.firCachesFactory.createCache<FqName, Set<Name>> { packageFqName ->
+            declarationProvider.getTopLevelCallableNamesInPackage(packageFqName)
+        }
+
     fun getFirClassifierByFqNameAndDeclaration(
         classId: ClassId,
         classLikeDeclaration: KtClassLikeDeclaration?,
@@ -68,6 +93,9 @@ internal class LLFirProviderHelper(
         return classifierByClassId.getValue(classId, classLikeDeclaration)
     }
 
+    fun getTopLevelClassifierNamesInPackage(packageFqName: FqName): ClassifierNamesInPackage =
+        topLevelClassifierNamesByPackage.getValue(packageFqName)
+
     fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
         if (!allowKotlinPackage && packageFqName.isKotlinPackage()) return emptyList()
         val callableId = CallableId(packageFqName, name)
@@ -82,6 +110,9 @@ internal class LLFirProviderHelper(
         return getTopLevelCallableSymbols(packageFqName, name).filterIsInstance<FirPropertySymbol>()
     }
 
+    fun getTopLevelCallableNamesInPackage(packageFqName: FqName): Set<Name> =
+        topLevelCallableNamesByPackage.getValue(packageFqName)
+
     private fun FirFile.collectCallableDeclarationsTo(list: MutableList<FirCallableSymbol<*>>, name: Name) {
         declarations.mapNotNullTo(list) { declaration ->
             if (declaration is FirCallableDeclaration && declaration.symbol.callableId.callableName == name) {