[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
2023-10-10 13:17:40 +02:00
parent c63dde4f7e
commit b9e3d848a6
12 changed files with 75 additions and 45 deletions
@@ -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
@@ -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)
        }
    }
@@ -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)
@@ -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
-    project: Project,
-    protected val searchScope: GlobalSearchScope
+    protected val project: Project,
+    public val searchScope: GlobalSearchScope,
 ) : KotlinPackageProvider() {

    override fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {
@@ -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."
+                }
+            }
+        }
+}
@@ -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
-private constructor(
-    private val providers: List<KotlinPackageProvider>
-) : KotlinPackageProvider() {
-
+public class CompositeKotlinPackageProvider private constructor(
+    override val providers: List<KotlinPackageProvider>,
+) : KotlinPackageProvider(), KotlinCompositeProvider<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 {
-        fun create(providers: List<KotlinPackageProvider>): KotlinPackageProvider {
-            return when (providers.size) {
-                0 -> EmptyKotlinPackageProvider
-                1 -> providers.single()
-                else -> CompositeKotlinPackageProvider(providers)
-            }
-        }
+    public companion object {
+        public val factory: KotlinCompositeProviderFactory<KotlinPackageProvider> = KotlinCompositeProviderFactory(
+            EmptyKotlinPackageProvider,
+            ::CompositeKotlinPackageProvider,
+        )
+
+        public fun create(providers: List<KotlinPackageProvider>): KotlinPackageProvider = factory.create(providers)
    }
 }
@@ -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 {
-        return false
-    }
+    override fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean = false

-    override fun doesKotlinOnlyPackageExist(packageFqName: FqName): Boolean {
-        return false
-    }
+    override fun doesKotlinOnlyPackageExist(packageFqName: FqName): Boolean = false

-    override fun doesPlatformSpecificPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {
-        return false
-    }
+    override fun doesPlatformSpecificPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean = false

-    override fun getSubPackageFqNames(packageFqName: FqName, platform: TargetPlatform, nameFilter: (Name) -> Boolean): Set<Name> {
-        return emptySet()
-    }
+    override fun getSubPackageFqNames(packageFqName: FqName, platform: TargetPlatform, nameFilter: (Name) -> Boolean): Set<Name> =
+        emptySet()

-    override fun getKotlinOnlySubPackagesFqNames(packageFqName: FqName, nameFilter: (Name) -> Boolean): Set<Name> {
-        return emptySet()
-    }
+    override fun getKotlinOnlySubPackagesFqNames(packageFqName: FqName, nameFilter: (Name) -> Boolean): Set<Name> = emptySet()

    override fun getPlatformSpecificSubPackagesFqNames(
        packageFqName: FqName,
        platform: TargetPlatform,
-        nameFilter: (Name) -> Boolean
-    ): Set<Name> {
-        return emptySet()
-    }
+        nameFilter: (Name) -> Boolean,
+    ): Set<Name> = emptySet()
 }
@@ -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,
@@ -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`.
@@ -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
@@ -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,
@@ -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,