[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.

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

@@ -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))
         }
     }
-
 }

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

@@ -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,
+    val providers: List<FirSymbolProvider>,
-    private 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) {
-
+    init {
-    abstract val dependentProviders: List<FirSymbolProvider>
+        require(providers.all { it !is LLFirModuleWithDependenciesSymbolProvider }) {
-    abstract val dependentSessions: List<LLFirSession>
+            "${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 ->
+        providers.firstNotNullOfOrNull { it.getClassLikeSymbolByClassId(classId) }
-            when (provider) {
-                is LLFirModuleWithDependenciesSymbolProvider -> provider.getClassLikeSymbolByFqNameWithoutDependencies(classId)
-                else -> provider.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) {
+                provider.getTopLevelCallableSymbolsTo(this, packageFqName, name)
-                    is LLFirModuleWithDependenciesSymbolProvider -> provider.getTopLevelCallableSymbolsToWithoutDependencies(this, packageFqName, name)
-                    else -> 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) {
+                provider.getTopLevelFunctionSymbolsTo(this, packageFqName, name)
-                    is LLFirModuleWithDependenciesSymbolProvider -> provider.getTopLevelFunctionSymbolsToWithoutDependencies(this, packageFqName, name)
-                    else -> 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) {
+                provider.getTopLevelPropertySymbolsTo(this, packageFqName, name)
-                    is LLFirModuleWithDependenciesSymbolProvider -> provider.getTopLevelPropertySymbolsToWithoutDependencies(this, packageFqName, name)
-                    else -> provider.getTopLevelPropertySymbolsTo(this, packageFqName, name)
-                }
             }
             addNewSymbolsConsideringJvmFacades(destination, newSymbols, facades)
         }
     }
 
-    override fun getPackage(fqName: FqName): FqName? =
+    override fun getPackage(fqName: FqName): FqName? = providers.firstNotNullOfOrNull { it.getPackage(fqName) }
-        dependentProviders.firstNotNullOfOrNull { provider ->
-            when (provider) {
-                is LLFirModuleWithDependenciesSymbolProvider -> provider.getPackageWithoutDependencies(fqName)
-                else -> provider.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() })
 }

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

@@ -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) {
+            val dependencyProvider = LLFirDependenciesSymbolProvider(this, buildList {
-                addAll(dependencies)
+                addDependencySymbolProvidersTo(session, dependencies, this)
-                add(builtinsSession)
+                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) {
+            val dependencyProvider = LLFirDependenciesSymbolProvider(this, buildList {
-                addAll(dependencies)
+                addDependencySymbolProvidersTo(session, dependencies, this)
-                add(builtinsSession)
+                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) {
+            val dependencyProvider = LLFirDependenciesSymbolProvider(this, listOf(builtinsSession.symbolProvider))
-                add(builtinsSession)
-            }
 
             register(
                 FirSymbolProvider::class,
                 LLFirModuleWithDependenciesSymbolProvider(
                     this,
-                    dependencyProvider,
+                    providers = listOf(
-                    providers = listOfNotNull(
                         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`
+            // The dependency provider needs to have access to all direct and indirect `dependsOn` dependencies, as `dependsOn`
-        // dependencies are transitive.
+            // dependencies are transitive.
-        val directRegularDependenciesSet = module.directRegularDependencies.toSet()
+            addAll(module.transitiveDependsOnDependencies)
-        module.transitiveDependsOnDependencies.forEach { dependency ->
-            if (dependency !in directRegularDependenciesSet) {
-                getOrCreateSessionForDependency(dependency)?.let(dependencies::add)
-            }
         }
 
-        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) {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirDeclarationForCompiledElementSearcher.kt

@@ -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