[LL FIR] Optimize getPackage in combined Kotlin symbol providers

- `getPackage` can also benefit from a combined index access.
- Care has to be taken with allowed/disallowed `kotlin` packages. Since
  we're not delegating to individual symbol providers after the index
  access, `allowKotlinPackage` has to be taken into account in the
  combined symbol provider explicitly.

^KT-61791 fixed

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

@@ -8,9 +8,11 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.providers
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.low.level.api.fir.caches.NullableCaffeineCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
-import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
+import org.jetbrains.kotlin.analysis.providers.createPackageProvider
 import org.jetbrains.kotlin.analysis.providers.mergeDeclarationProviders
+import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirCompositeCachedSymbolNamesProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolNamesProvider
@@ -38,12 +40,17 @@ import org.jetbrains.kotlin.psi.KtCallableDeclaration
  *   footprint.
  *
  * [declarationProvider] must have a scope which combines the scopes of the individual [providers].
+ *
+ * [packageProviderForKotlinPackages] should be the package provider combined from all [providers] which allow `kotlin` packages (see
+ * [LLFirProvider.SymbolProvider.allowKotlinPackage]). It may be `null` if no such provider exists. See [getPackage] for a use case.
  */
 internal class LLFirCombinedKotlinSymbolProvider private constructor(
     session: FirSession,
     project: Project,
     providers: List<LLFirKotlinSymbolProvider>,
     private val declarationProvider: KotlinDeclarationProvider,
+    private val packageProvider: KotlinPackageProvider,
+    private val packageProviderForKotlinPackages: KotlinPackageProvider?,
 ) : LLFirSelectingCombinedSymbolProvider<LLFirKotlinSymbolProvider>(session, project, providers) {
     override val symbolNamesProvider: FirSymbolNamesProvider = FirCompositeCachedSymbolNamesProvider.fromSymbolProviders(session, providers)
 
@@ -111,13 +118,42 @@ internal class LLFirCombinedKotlinSymbolProvider private constructor(
             }
     }
 
-    override fun getPackage(fqName: FqName): FqName? = providers.firstNotNullOfOrNull { it.getPackage(fqName) }
+    override fun getPackage(fqName: FqName): FqName? {
+        val hasPackage = if (fqName.startsWith(StandardNames.BUILT_INS_PACKAGE_NAME)) {
+            // If a package is a `kotlin` package, `packageProvider` might find it via the scope of an individual symbol provider that
+            // disallows `kotlin` packages. Hence, the combined `getPackage` would erroneously find a package it shouldn't be able to find,
+            // because calling that individual symbol provider directly would result in `null` (as it disallows `kotlin` packages). The
+            // `packageProviderForKotlinPackages` solves this issue by including only scopes from symbol providers which allow `kotlin`
+            // packages.
+            packageProviderForKotlinPackages?.doesKotlinOnlyPackageExist(fqName) == true
+        } else {
+            packageProvider.doesKotlinOnlyPackageExist(fqName)
+        }
+        return fqName.takeIf { hasPackage }
+    }
 
     companion object {
         fun merge(session: LLFirSession, project: Project, providers: List<LLFirKotlinSymbolProvider>): FirSymbolProvider? =
             if (providers.size > 1) {
                 val declarationProvider = project.mergeDeclarationProviders(providers.map { it.declarationProvider })
-                LLFirCombinedKotlinSymbolProvider(session, project, providers, declarationProvider)
+
+                // TODO (marco): Implement a package provider merger.
+                val combinedScope = providers.createCombinedScope()
+                val packageProvider = project.createPackageProvider(combinedScope)
+
+                val packageProviderForKotlinPackages = providers
+                    .filter { it.allowKotlinPackage }
+                    .takeIf { it.isNotEmpty() }
+                    ?.let { project.createPackageProvider(it.createCombinedScope()) }
+
+                LLFirCombinedKotlinSymbolProvider(
+                    session,
+                    project,
+                    providers,
+                    declarationProvider,
+                    packageProvider,
+                    packageProviderForKotlinPackages,
+                )
             } else providers.singleOrNull()
     }
 }

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

@@ -28,6 +28,12 @@ import org.jetbrains.kotlin.psi.KtProperty
 internal abstract class LLFirKotlinSymbolProvider(session: FirSession) : FirSymbolProvider(session) {
     abstract val declarationProvider: KotlinDeclarationProvider
 
+    /**
+     * Whether the [LLFirKotlinSymbolProvider] should be able to find symbols defined in `kotlin` packages. This is usually not the case for
+     * source sessions, unless the `allowKotlinPackage` flag is enabled in the session's `languageVersionSettings`.
+     */
+    abstract val allowKotlinPackage: Boolean
+
     /**
      * This function is optimized for a known [classLikeDeclaration].
      */

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

@@ -108,6 +108,8 @@ internal class LLFirProvider(
 
         override val symbolNamesProvider: FirSymbolNamesProvider get() = providerHelper.symbolNameCache
 
+        override val allowKotlinPackage get() = providerHelper.allowKotlinPackage
+
         override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
             if (!providerHelper.symbolNameCache.mayHaveTopLevelClassifier(classId)) return null
             return getFirClassifierByFqName(classId)?.symbol

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

@@ -70,7 +70,7 @@ internal class LLFirProviderHelper(
         )
     )
 
-    private val allowKotlinPackage = canContainKotlinPackage ||
+    val allowKotlinPackage: Boolean = canContainKotlinPackage ||
             firSession.languageVersionSettings.getFlag(AnalysisFlags.allowKotlinPackage)
 
     private val classifierByClassId =

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

@@ -5,6 +5,14 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.providers
 
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import java.util.*
 
-internal fun <T: Any> Optional<T>.getOrNull(): T? = orElse(null)
+internal fun <T: Any> Optional<T>.getOrNull(): T? = orElse(null)
+
+internal fun List<FirSymbolProvider>.createCombinedScope(): GlobalSearchScope {
+    if (isEmpty()) return GlobalSearchScope.EMPTY_SCOPE
+    return GlobalSearchScope.union(map { it.session.llFirModuleData.ktModule.contentScope })
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/stubBased/deserialization/StubBasedFirDeserializedSymbolProvider.kt

@@ -58,6 +58,8 @@ internal open class StubBasedFirDeserializedSymbolProvider(
 
     override val symbolNamesProvider: FirSymbolNamesProvider = LLFirKotlinSymbolNamesProvider.cached(session, declarationProvider)
 
+    override val allowKotlinPackage: Boolean get() = true
+
     private val typeAliasCache: FirCache<ClassId, FirTypeAliasSymbol?, StubBasedFirDeserializationContext?> =
         session.firCachesFactory.createCacheWithPostCompute(
             createValue = { classId, context -> findAndDeserializeTypeAlias(classId, context) },