[LL FIR] KT-57220 Flatten session dependency symbol providers

- Replace the "recurse one layer deep" approach of `LLFirDependentModuleProviders` with a flattened list of dependency symbol providers. - This refactoring is a necessary component of supporting combined symbol providers.
2023-03-09 14:01:47 +01:00
parent 1817213e1d
commit 79457ec074
4 changed files with 87 additions and 112 deletions
@@ -7,11 +7,10 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure

 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
-import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirDependentModuleProvidersByProviders
+import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirDependenciesSymbolProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirModuleWithDependenciesSymbolProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirJavaFacadeForBinaries
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
-import org.jetbrains.kotlin.analysis.project.structure.KtBinaryModule
 import org.jetbrains.kotlin.analysis.providers.createPackagePartProvider
 import org.jetbrains.kotlin.fir.BuiltinTypes
 import org.jetbrains.kotlin.fir.deserialization.SingleModuleDataProvider
@@ -36,7 +35,6 @@ internal object LLFirLibraryProviderFactory {
        val packagePartProvider = project.createPackagePartProvider(scope)
        return LLFirModuleWithDependenciesSymbolProvider(
            session,
-            LLFirDependentModuleProvidersByProviders(session, listOf(builtinSymbolProvider)),
            providers = listOf(
                JvmClassFileBasedSymbolProvider(
                    session,
@@ -48,6 +46,7 @@ internal object LLFirLibraryProviderFactory {
                ),
                OptionalAnnotationClassesProvider(session, moduleDataProvider, kotlinScopeProvider, packagePartProvider),
            ),
+            LLFirDependenciesSymbolProvider(session, listOf(builtinSymbolProvider)),
        )
    }

@@ -75,5 +74,4 @@ internal object LLFirLibraryProviderFactory {
            add(OptionalAnnotationClassesProvider(session, moduleDataProvider, kotlinScopeProvider, packagePartProvider))
        }
    }
-
 }
@@ -5,12 +5,10 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.providers

-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.utils.collections.buildSmartList
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
-import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
@@ -24,15 +22,13 @@ import org.jetbrains.kotlin.utils.SmartSet

 internal class LLFirModuleWithDependenciesSymbolProvider(
    session: FirSession,
-    val dependencyProvider: LLFirDependentModuleProviders,
-    private val providers: List<FirSymbolProvider>,
+    val providers: List<FirSymbolProvider>,
+    val dependencyProvider: LLFirDependenciesSymbolProvider,
 ) : FirSymbolProvider(session) {
-
    override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? =
        getClassLikeSymbolByFqNameWithoutDependencies(classId)
            ?: dependencyProvider.getClassLikeSymbolByClassId(classId)

-
    fun getClassLikeSymbolByFqNameWithoutDependencies(classId: ClassId): FirClassLikeSymbol<*>? =
        providers.firstNotNullOfOrNull { it.getClassLikeSymbolByClassId(classId) }

@@ -82,42 +78,34 @@ internal class LLFirModuleWithDependenciesSymbolProvider(
        getPackageWithoutDependencies(fqName)
            ?: dependencyProvider.getPackage(fqName)

-    override fun computePackageSetWithTopLevelCallables(): Set<String>? = null
-
-    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? = null
-
-    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? = null
-
-
    fun getPackageWithoutDependencies(fqName: FqName): FqName? =
        providers.firstNotNullOfOrNull { it.getPackage(fqName) }
+
+    override fun computePackageSetWithTopLevelCallables(): Set<String>? = null
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? = null
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? = null
 }

-internal abstract class LLFirDependentModuleProviders(
+internal class LLFirDependenciesSymbolProvider(
    session: FirSession,
+    val providers: List<FirSymbolProvider>,
 ) : FirSymbolProvider(session) {
-
-    abstract val dependentProviders: List<FirSymbolProvider>
-    abstract val dependentSessions: List<LLFirSession>
+    init {
+        require(providers.all { it !is LLFirModuleWithDependenciesSymbolProvider }) {
+            "${LLFirDependenciesSymbolProvider::class.simpleName} may not contain ${LLFirModuleWithDependenciesSymbolProvider::class.simpleName}:" +
+                    " dependency providers must be flattened during session creation."
+        }
+    }

    override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? =
-        dependentProviders.firstNotNullOfOrNull { provider ->
-            when (provider) {
-                is LLFirModuleWithDependenciesSymbolProvider -> provider.getClassLikeSymbolByFqNameWithoutDependencies(classId)
-                else -> provider.getClassLikeSymbolByClassId(classId)
-            }
-        }
-
+        providers.firstNotNullOfOrNull { it.getClassLikeSymbolByClassId(classId) }

    @FirSymbolProviderInternals
    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
        val facades = SmartSet.create<JvmClassName>()
-        for (provider in dependentProviders) {
+        for (provider in providers) {
            val newSymbols = buildSmartList {
-                when (provider) {
-                    is LLFirModuleWithDependenciesSymbolProvider -> provider.getTopLevelCallableSymbolsToWithoutDependencies(this, packageFqName, name)
-                    else -> provider.getTopLevelCallableSymbolsTo(this, packageFqName, name)
-                }
+                provider.getTopLevelCallableSymbolsTo(this, packageFqName, name)
            }
            addNewSymbolsConsideringJvmFacades(destination, newSymbols, facades)
        }
@@ -126,42 +114,27 @@ internal abstract class LLFirDependentModuleProviders(
    @FirSymbolProviderInternals
    override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
        val facades = SmartSet.create<JvmClassName>()
-        for (provider in dependentProviders) {
+        for (provider in providers) {
            val newSymbols = buildSmartList {
-                when (provider) {
-                    is LLFirModuleWithDependenciesSymbolProvider -> provider.getTopLevelFunctionSymbolsToWithoutDependencies(this, packageFqName, name)
-                    else -> provider.getTopLevelFunctionSymbolsTo(this, packageFqName, name)
-                }
+                provider.getTopLevelFunctionSymbolsTo(this, packageFqName, name)
            }
-
            addNewSymbolsConsideringJvmFacades(destination, newSymbols, facades)
        }
-
    }

    @FirSymbolProviderInternals
    override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
        val facades = SmartSet.create<JvmClassName>()
-        for (provider in dependentProviders) {
+        for (provider in providers) {
            val newSymbols = buildSmartList {
-                when (provider) {
-                    is LLFirModuleWithDependenciesSymbolProvider -> provider.getTopLevelPropertySymbolsToWithoutDependencies(this, packageFqName, name)
-                    else -> provider.getTopLevelPropertySymbolsTo(this, packageFqName, name)
-                }
+                provider.getTopLevelPropertySymbolsTo(this, packageFqName, name)
            }
            addNewSymbolsConsideringJvmFacades(destination, newSymbols, facades)
        }
    }

-    override fun getPackage(fqName: FqName): FqName? =
-        dependentProviders.firstNotNullOfOrNull { provider ->
-            when (provider) {
-                is LLFirModuleWithDependenciesSymbolProvider -> provider.getPackageWithoutDependencies(fqName)
-                else -> provider.getPackage(fqName)
-            }
-        }
+    override fun getPackage(fqName: FqName): FqName? = providers.firstNotNullOfOrNull { it.getPackage(fqName) }

-    // TODO: Consider having proper implementations for sake of optimizations
    override fun computePackageSetWithTopLevelCallables(): Set<String>? = null
    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? = null
    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? = null
@@ -191,31 +164,4 @@ internal abstract class LLFirDependentModuleProviders(
        val jvmPackagePartSource = fir.containerSource as? JvmPackagePartSource ?: return null
        return jvmPackagePartSource.facadeClassName ?: jvmPackagePartSource.className
    }
-
-}
-
-internal class LLFirDependentModuleProvidersBySessions(
-    session: FirSession,
-    override val dependentSessions: List<LLFirSession>
-) : LLFirDependentModuleProviders(session) {
-
-    override val dependentProviders: List<FirSymbolProvider> by lazy(LazyThreadSafetyMode.PUBLICATION) {
-        dependentSessions.map { it.symbolProvider }
-    }
-
-    constructor(session: FirSession, createSessions: MutableList<LLFirSession>.() -> Unit)
-            : this(session, buildList { createSessions() })
-}
-
-
-internal class LLFirDependentModuleProvidersByProviders(
-    session: FirSession,
-    override val dependentProviders: List<FirSymbolProvider>,
-) : LLFirDependentModuleProviders(session) {
-
-    override val dependentSessions: List<LLFirSession>
-        get() = dependentProviders.map { it.session as LLFirSession }
-
-    constructor(session: FirSession, createProviders: MutableList<FirSymbolProvider>.() -> Unit)
-            : this(session, buildList { createProviders() })
 }
@@ -38,6 +38,7 @@ import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.*
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirExtensionSyntheticFunctionInterfaceProvider
+import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
 import org.jetbrains.kotlin.fir.resolve.transformers.FirDummyCompilerLazyDeclarationResolver
 import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
@@ -157,10 +158,10 @@ internal class LLFirSessionCache(private val project: Project) {
                register(FirSwitchableExtensionDeclarationsSymbolProvider::class, it)
            }

-            val dependencyProvider = LLFirDependentModuleProvidersBySessions(this) {
-                addAll(dependencies)
-                add(builtinsSession)
-            }
+            val dependencyProvider = LLFirDependenciesSymbolProvider(this, buildList {
+                addDependencySymbolProvidersTo(session, dependencies, this)
+                add(builtinsSession.symbolProvider)
+            })

            val javaSymbolProvider = createJavaSymbolProvider(this, moduleData, project, contentScope)
            val syntheticFunctionalInterfaceProvider = FirExtensionSyntheticFunctionInterfaceProvider(this, moduleData, scopeProvider)
@@ -168,13 +169,13 @@ internal class LLFirSessionCache(private val project: Project) {
                FirSymbolProvider::class,
                LLFirModuleWithDependenciesSymbolProvider(
                    this,
-                    dependencyProvider,
                    providers = listOfNotNull(
                        provider.symbolProvider,
                        switchableExtensionDeclarationsSymbolProvider,
                        javaSymbolProvider,
                        syntheticFunctionalInterfaceProvider,
                    ),
+                    dependencyProvider,
                )
            )
            register(JavaSymbolProvider::class, javaSymbolProvider)
@@ -239,7 +240,7 @@ internal class LLFirSessionCache(private val project: Project) {
            register(FirRegisteredPluginAnnotations::class, LLFirIdeRegisteredPluginAnnotations(this, annotationsResolver))
            register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)

-            val dependencyProvider = LLFirDependentModuleProvidersByProviders(this) {
+            val dependencyProvider = LLFirDependenciesSymbolProvider(this, buildList {
                add(builtinsSession.symbolProvider)

                // Script dependencies are self-contained and should not depend on other libraries
@@ -255,18 +256,18 @@ internal class LLFirSessionCache(private val project: Project) {

                    addAll(restLibrariesProvider)
                }
-            }
+            })

            val javaSymbolProvider = createJavaSymbolProvider(this, moduleData, project, contentScope)
            register(
                FirSymbolProvider::class,
                LLFirModuleWithDependenciesSymbolProvider(
                    this,
-                    dependencyProvider,
                    providers = listOf(
                        provider.symbolProvider,
                        javaSymbolProvider,
                    ),
+                    dependencyProvider,
                )
            )
            register(JavaSymbolProvider::class, javaSymbolProvider)
@@ -358,10 +359,10 @@ internal class LLFirSessionCache(private val project: Project) {
            register(FirProvider::class, provider)
            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver())

-            val dependencyProvider = LLFirDependentModuleProvidersBySessions(this) {
-                addAll(dependencies)
-                add(builtinsSession)
-            }
+            val dependencyProvider = LLFirDependenciesSymbolProvider(this, buildList {
+                addDependencySymbolProvidersTo(session, dependencies, this)
+                add(builtinsSession.symbolProvider)
+            })

            val javaSymbolProvider = createJavaSymbolProvider(this, moduleData, project, contentScope)
            register(JavaSymbolProvider::class, javaSymbolProvider)
@@ -370,11 +371,11 @@ internal class LLFirSessionCache(private val project: Project) {
                FirSymbolProvider::class,
                LLFirModuleWithDependenciesSymbolProvider(
                    this,
-                    dependencyProvider,
                    providers = listOfNotNull(
                        javaSymbolProvider,
                        provider.symbolProvider,
-                    )
+                    ),
+                    dependencyProvider,
                )
            )

@@ -423,18 +424,16 @@ internal class LLFirSessionCache(private val project: Project) {
            register(FirProvider::class, provider)
            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver())

-            val dependencyProvider = LLFirDependentModuleProvidersBySessions(this) {
-                add(builtinsSession)
-            }
+            val dependencyProvider = LLFirDependenciesSymbolProvider(this, listOf(builtinsSession.symbolProvider))

            register(
                FirSymbolProvider::class,
                LLFirModuleWithDependenciesSymbolProvider(
                    this,
-                    dependencyProvider,
-                    providers = listOfNotNull(
+                    providers = listOf(
                        provider.symbolProvider,
-                    )
+                    ),
+                    dependencyProvider,
                )
            )

@@ -466,8 +465,6 @@ internal class LLFirSessionCache(private val project: Project) {
    }

    private fun collectSourceModuleDependencies(module: KtModule): List<LLFirSession> {
-        val dependencies = mutableListOf<LLFirSession>()
-
        fun getOrCreateSessionForDependency(dependency: KtModule): LLFirSession? = when (dependency) {
            is KtBuiltinsModule -> null // Built-ins are already added
            is KtBinaryModule -> getSession(dependency, preferBinary = true)
@@ -479,18 +476,15 @@ internal class LLFirSessionCache(private val project: Project) {
            is KtLibrarySourceModule -> error("Module $module cannot depend on ${dependency::class}: $dependency")
        }

-        module.directRegularDependencies.mapNotNullTo(dependencies, ::getOrCreateSessionForDependency)
+        val dependencyModules = buildSet {
+            addAll(module.directRegularDependencies)

-        // The dependency provider needs to have access to all direct and indirect `dependsOn` dependencies, as `dependsOn`
-        // dependencies are transitive.
-        val directRegularDependenciesSet = module.directRegularDependencies.toSet()
-        module.transitiveDependsOnDependencies.forEach { dependency ->
-            if (dependency !in directRegularDependenciesSet) {
-                getOrCreateSessionForDependency(dependency)?.let(dependencies::add)
-            }
+            // The dependency provider needs to have access to all direct and indirect `dependsOn` dependencies, as `dependsOn`
+            // dependencies are transitive.
+            addAll(module.transitiveDependsOnDependencies)
        }

-        return dependencies
+        return dependencyModules.mapNotNull(::getOrCreateSessionForDependency)
    }

    private fun createSourceModuleDependencyTracker(module: KtModule, exposedDependencies: List<LLFirSession>): ModificationTracker {
@@ -506,6 +500,43 @@ internal class LLFirSessionCache(private val project: Project) {
    private fun createModuleData(session: LLFirSession): LLFirModuleData {
        return LLFirModuleData(session.ktModule).apply { bindSession(session) }
    }
+
+    /**
+     * Adds dependency symbol providers from [dependencies] to [destination]. The function might combine, reorder, or exclude specific
+     * symbol providers for optimization.
+     */
+    private fun addDependencySymbolProvidersTo(
+        session: LLFirSession,
+        dependencies: List<LLFirSession>,
+        destination: MutableList<FirSymbolProvider>,
+    ) {
+        val dependencyProviders = buildList {
+            dependencies.forEach { session ->
+                when (val dependencyProvider = session.symbolProvider) {
+                    is LLFirModuleWithDependenciesSymbolProvider -> addAll(dependencyProvider.providers)
+                    else -> add(dependencyProvider)
+                }
+            }
+        }
+
+        dependencyProviders.mergeDependencySymbolProvidersInto(session, destination)
+    }
+
+    /**
+     * Merges dependency symbol providers of the same kind if possible. The merged symbol provider usually delegates to its subordinate
+     * symbol providers to preserve session semantics, but it will have some form of advantage over individual symbol providers (such as
+     * querying an index once instead of N times).
+     *
+     * [session] should be the session of the dependent module. Because all symbol providers are tied to a session, we need a session to
+     * create a combined symbol provider.
+     */
+    @Suppress("UNUSED_PARAMETER")
+    private fun List<FirSymbolProvider>.mergeDependencySymbolProvidersInto(
+        session: FirSession,
+        destination: MutableList<FirSymbolProvider>,
+    ) {
+        destination.addAll(this)
+    }
 }

 internal fun LLFirSessionConfigurator.Companion.configure(session: LLFirSession) {
@@ -125,7 +125,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
 private val LLFirModuleWithDependenciesSymbolProvider.friendBuiltinsProvider: FirSymbolProvider?
    get() {
        if (getPackageWithoutDependencies(StandardClassIds.BASE_KOTLIN_PACKAGE) != null) {
-            return dependencyProvider.dependentProviders.find { it.session is LLFirBuiltinsAndCloneableSession }
+            return dependencyProvider.providers.find { it.session is LLFirBuiltinsAndCloneableSession }
        }

        return null