[AA] Add KotlinPackageProviderMerger

- In parallel to Kotlin declaration provider merging, we need a proper
  merging strategy for package providers as well, because resolve
  extensions may define additional package providers.
- Additionally, other non-scope-based package providers may be added in
  the future, and the merger preserves these out of the box.

^KT-61791

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

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir
-import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.CompositeKotlinPackageProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirResolveExtensionTool
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.llResolveExtensionTool
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
@@ -29,6 +28,7 @@ import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.createPackageProvider
 import org.jetbrains.kotlin.analysis.providers.impl.declarationProviders.CompositeKotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.providers.impl.packageProviders.CompositeKotlinPackageProvider
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.resolve.ScopeSession

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/configurators/AnalysisApiBaseTestServiceRegistrar.kt

@@ -99,6 +99,7 @@ object AnalysisApiBaseTestServiceRegistrar: AnalysisApiTestServiceRegistrar()  {
                 KotlinPackageProviderFactory::class.java,
                 KotlinStaticPackageProviderFactory(project, allSourceKtFiles + ktFilesForBinaries)
             )
+            registerService(KotlinPackageProviderMerger::class.java, KotlinStaticPackageProviderMerger(project))
             registerService(KotlinResolutionScopeProvider::class.java, KotlinByModulesResolutionScopeProvider::class.java)
         }
     }

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

@@ -76,6 +76,27 @@ public abstract class KotlinPackageProviderFactory {
     public abstract fun createPackageProvider(searchScope: GlobalSearchScope): KotlinPackageProvider
 }
 
+/**
+ * [KotlinPackageProviderMerger] allows merging multiple [KotlinPackageProvider]s into a more efficient package provider.
+ *
+ * Package providers should not be naively merged by combining scopes and calling [createPackageProvider], because there may be additional
+ * package providers which do not operate based on scopes (e.g. resolve extension package providers).
+ */
+public abstract class KotlinPackageProviderMerger {
+    /**
+     * Merges [packageProviders] if possible, creating a combined package provider that should be more efficient compared to invoking
+     * separate package providers.
+     */
+    public abstract fun mergePackageProviders(packageProviders: List<KotlinPackageProvider>): KotlinPackageProvider
+
+    public companion object {
+        public fun getInstance(project: Project): KotlinPackageProviderMerger = project.getService(KotlinPackageProviderMerger::class.java)
+    }
+}
+
 public fun Project.createPackageProvider(searchScope: GlobalSearchScope): KotlinPackageProvider =
     this.getService(KotlinPackageProviderFactory::class.java)
-        .createPackageProvider(searchScope)
+        .createPackageProvider(searchScope)
+
+public fun Project.mergePackageProviders(packageProviders: List<KotlinPackageProvider>): KotlinPackageProvider =
+    KotlinPackageProviderMerger.getInstance(this).mergePackageProviders(packageProviders)

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

@@ -13,11 +13,9 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.platform.jvm.isJvm
 
-
 public abstract class KotlinPackageProviderBase(
-    protected val
+    protected val project: Project,
-    project: Project,
+    public val searchScope: GlobalSearchScope,
-    protected val searchScope: GlobalSearchScope
 ) : KotlinPackageProvider() {
 
     override fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {

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

@@ -9,13 +9,16 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinPackageProviderFactory
+import org.jetbrains.kotlin.analysis.providers.KotlinPackageProviderMerger
+import org.jetbrains.kotlin.analysis.providers.createPackageProvider
+import org.jetbrains.kotlin.analysis.providers.impl.packageProviders.CompositeKotlinPackageProvider
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtFile
 
 public class KotlinStaticPackageProvider(
     project: Project,
-    scope: GlobalSearchScope,
+    internal val scope: GlobalSearchScope,
     files: Collection<KtFile>
 ) : KotlinPackageProviderBase(project, scope) {
     private val kotlinPackageToSubPackages: Map<FqName, Set<Name>> = run {
@@ -48,4 +51,16 @@ public class KotlinStaticPackageProviderFactory(
     override fun createPackageProvider(searchScope: GlobalSearchScope): KotlinPackageProvider {
         return KotlinStaticPackageProvider(project, searchScope, files)
     }
-}
+}
+
+public class KotlinStaticPackageProviderMerger(private val project: Project) : KotlinPackageProviderMerger() {
+    override fun mergePackageProviders(packageProviders: List<KotlinPackageProvider>): KotlinPackageProvider =
+        packageProviders.mergeSpecificProviders<_, KotlinStaticPackageProvider>(CompositeKotlinPackageProvider.factory) { providers ->
+            val combinedScope = GlobalSearchScope.union(providers.map { it.scope })
+            project.createPackageProvider(combinedScope).apply {
+                check(this is KotlinStaticPackageProvider) {
+                    "`${KotlinStaticPackageProvider::class.simpleName}` can only be merged into a combined package provider of the same type."
+                }
+            }
+        }
+}

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

@@ -3,18 +3,18 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-package org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure
+package org.jetbrains.kotlin.analysis.providers.impl.packageProviders
 
 import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
+import org.jetbrains.kotlin.analysis.providers.impl.KotlinCompositeProvider
+import org.jetbrains.kotlin.analysis.providers.impl.KotlinCompositeProviderFactory
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
 
-class CompositeKotlinPackageProvider
+public class CompositeKotlinPackageProvider private constructor(
-private constructor(
+    override val providers: List<KotlinPackageProvider>,
-    private val providers: List<KotlinPackageProvider>
+) : KotlinPackageProvider(), KotlinCompositeProvider<KotlinPackageProvider> {
-) : KotlinPackageProvider() {
-
     override fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {
         return providers.any { it.doesPackageExist(packageFqName, platform) }
     }
@@ -43,13 +43,12 @@ private constructor(
         return providers.flatMapTo(mutableSetOf()) { it.getPlatformSpecificSubPackagesFqNames(packageFqName, platform, nameFilter) }
     }
 
-    companion object {
+    public companion object {
-        fun create(providers: List<KotlinPackageProvider>): KotlinPackageProvider {
+        public val factory: KotlinCompositeProviderFactory<KotlinPackageProvider> = KotlinCompositeProviderFactory(
-            return when (providers.size) {
+            EmptyKotlinPackageProvider,
-                0 -> EmptyKotlinPackageProvider
+            ::CompositeKotlinPackageProvider,
-                1 -> providers.single()
+        )
-                else -> CompositeKotlinPackageProvider(providers)
+
-            }
+        public fun create(providers: List<KotlinPackageProvider>): KotlinPackageProvider = factory.create(providers)
-        }
     }
 }

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

@@ -3,7 +3,7 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-package org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure
+package org.jetbrains.kotlin.analysis.providers.impl.packageProviders
 
 import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.name.FqName
@@ -11,31 +11,20 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
 
 internal object EmptyKotlinPackageProvider : KotlinPackageProvider() {
-    override fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {
+    override fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean = false
-        return false
-    }
 
-    override fun doesKotlinOnlyPackageExist(packageFqName: FqName): Boolean {
+    override fun doesKotlinOnlyPackageExist(packageFqName: FqName): Boolean = false
-        return false
-    }
 
-    override fun doesPlatformSpecificPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {
+    override fun doesPlatformSpecificPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean = false
-        return false
-    }
 
-    override fun getSubPackageFqNames(packageFqName: FqName, platform: TargetPlatform, nameFilter: (Name) -> Boolean): Set<Name> {
+    override fun getSubPackageFqNames(packageFqName: FqName, platform: TargetPlatform, nameFilter: (Name) -> Boolean): Set<Name> =
-        return emptySet()
+        emptySet()
-    }
 
-    override fun getKotlinOnlySubPackagesFqNames(packageFqName: FqName, nameFilter: (Name) -> Boolean): Set<Name> {
+    override fun getKotlinOnlySubPackagesFqNames(packageFqName: FqName, nameFilter: (Name) -> Boolean): Set<Name> = emptySet()
-        return emptySet()
-    }
 
     override fun getPlatformSpecificSubPackagesFqNames(
         packageFqName: FqName,
         platform: TargetPlatform,
-        nameFilter: (Name) -> Boolean
+        nameFilter: (Name) -> Boolean,
-    ): Set<Name> {
+    ): Set<Name> = emptySet()
-        return emptySet()
-    }
 }

analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt

@@ -144,6 +144,7 @@ public class StandaloneAnalysisAPISessionBuilder(
                 KotlinPackageProviderFactory::class.java,
                 KotlinStaticPackageProviderFactory(project, sourceKtFiles + declarationProviderFactory.getAdditionalCreatedKtFiles())
             )
+            registerService(KotlinPackageProviderMerger::class.java, KotlinStaticPackageProviderMerger(this))
 
             registerService(
                 FirSealedClassInheritorsProcessorFactory::class.java,

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.providers
 
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
@@ -28,6 +29,8 @@ import org.jetbrains.kotlin.psi.KtProperty
 internal abstract class LLFirKotlinSymbolProvider(session: FirSession) : FirSymbolProvider(session) {
     abstract val declarationProvider: KotlinDeclarationProvider
 
+    abstract val packageProvider: KotlinPackageProvider
+
     /**
      * 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`.

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveCompone
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.SyntheticFirClassProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.fir.NoMutableState
 import org.jetbrains.kotlin.fir.ThreadSafeMutableState
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
@@ -106,6 +107,8 @@ internal class LLFirProvider(
     internal inner class SymbolProvider : LLFirKotlinSymbolProvider(session) {
         override val declarationProvider: KotlinDeclarationProvider get() = providerHelper.declarationProvider
 
+        override val packageProvider: KotlinPackageProvider get() = providerHelper.packageProvider
+
         override val symbolNamesProvider: FirSymbolNamesProvider get() = providerHelper.symbolNameCache
 
         override val allowKotlinPackage get() = providerHelper.allowKotlinPackage

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

@@ -8,7 +8,6 @@ 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.caches.getNotNullValueForNotNullContext
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder
-import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.CompositeKotlinPackageProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirResolveExtensionTool
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.llResolveExtensionTool
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
@@ -17,6 +16,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirKotlinSymbolNam
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.createPackageProvider
 import org.jetbrains.kotlin.analysis.providers.impl.declarationProviders.CompositeKotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.providers.impl.packageProviders.CompositeKotlinPackageProvider
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.config.AnalysisFlags
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
@@ -63,7 +63,7 @@ internal class LLFirProviderHelper(
         )
     )
 
-    private val packageProvider = CompositeKotlinPackageProvider.create(
+    val packageProvider = CompositeKotlinPackageProvider.create(
         listOfNotNull(
             firSession.project.createPackageProvider(searchScope),
             extensionTool?.packageProvider,

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

@@ -78,7 +78,7 @@ internal open class StubBasedFirDeserializedSymbolProvider(
     private val functionCache = session.firCachesFactory.createCache(::loadFunctionsByCallableId)
     private val propertyCache = session.firCachesFactory.createCache(::loadPropertiesByCallableId)
 
-    private val packageProvider = project.createPackageProvider(scope)
+    final override val packageProvider = project.createPackageProvider(scope)
 
     private fun findAndDeserializeTypeAlias(
         classId: ClassId,