[Low Level FIR] optimize getting callables in LLFirProviderHelper

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

@@ -29,6 +29,8 @@ public abstract class KotlinDeclarationProvider {
     public abstract fun getTopLevelProperties(callableId: CallableId): Collection<KtProperty>
     public abstract fun getTopLevelFunctions(callableId: CallableId): Collection<KtNamedFunction>
 
+    public abstract fun getTopLevelCallableFiles(callableId: CallableId): Collection<KtFile>
+
     public abstract fun getTopLevelCallableNamesInPackage(packageFqName: FqName): Set<Name>
 
     public abstract fun findFilesForFacadeByPackage(packageFqName: FqName): Collection<KtFile>

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

@@ -97,6 +97,11 @@ public class KotlinStaticDeclarationProvider internal constructor(
             }
             ?: emptyList()
 
+    override fun getTopLevelCallableFiles(callableId: CallableId): Collection<KtFile> = buildSet {
+        getTopLevelProperties(callableId).mapTo(this) { it.containingKtFile }
+        getTopLevelFunctions(callableId).mapTo(this) { it.containingKtFile }
+    }
+
 }
 
 public class KotlinStaticDeclarationProviderFactory(

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/LLFirNonUnderContentRootSessionFactory.kt

@@ -74,7 +74,7 @@ internal class LLFirNonUnderContentRootSessionFactory(private val project: Proje
                 components,
                 project.createDeclarationProvider(contentScope),
                 project.createPackageProvider(contentScope),
-                cabContainKotlinPackage = true,
+                canContainKotlinPackage = true,
             )
 
             register(FirProvider::class, provider)

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

@@ -29,7 +29,7 @@ internal class LLFirProvider(
     private val moduleComponents: LLFirModuleResolveComponents,
     private val declarationProvider: KotlinDeclarationProvider,
     packageProvider: KotlinPackageProvider,
-    cabContainKotlinPackage: Boolean,
+    canContainKotlinPackage: Boolean,
 ) : FirProvider() {
     override val symbolProvider: FirSymbolProvider = SymbolProvider()
 
@@ -38,7 +38,7 @@ internal class LLFirProvider(
         moduleComponents.firFileBuilder,
         declarationProvider,
         packageProvider,
-        cabContainKotlinPackage,
+        canContainKotlinPackage,
     )
 
     override val isPhasedFirAllowed: Boolean get() = true

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

@@ -52,10 +52,7 @@ internal class LLFirProviderHelper(
 
 
     private val callablesByCallableId = firSession.firCachesFactory.createCache<CallableId, List<FirCallableSymbol<*>>> { callableId ->
-        val files = Sets.newIdentityHashSet<KtFile>().apply {
+        val files = declarationProvider.getTopLevelCallableFiles(callableId).ifEmpty { return@createCache emptyList() }
-            declarationProvider.getTopLevelFunctions(callableId).mapTo(this) { it.containingKtFile }
-            declarationProvider.getTopLevelProperties(callableId).mapTo(this) { it.containingKtFile }
-        }
         buildList {
             files.forEach { ktFile ->
                 val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionFactory.kt

@@ -98,7 +98,8 @@ internal object LLFirSessionFactory {
                 components,
                 project.createDeclarationProvider(contentScope),
                 project.createPackageProvider(contentScope),
-                cabContainKotlinPackage = true,
+                /* Source modules can contain `kotlin` package only if `-Xallow-kotlin-package` is specified, this is handled in LLFirProvider */
+                canContainKotlinPackage = false,
             )
 
             register(FirProvider::class, provider)
@@ -203,7 +204,7 @@ internal object LLFirSessionFactory {
                 components,
                 project.createDeclarationProvider(contentScope),
                 project.createPackageProvider(contentScope),
-                cabContainKotlinPackage = true,
+                canContainKotlinPackage = true,
             )
 
             register(FirProvider::class, provider)