[Analysis API] Provide sessions in 'LLFirModuleData' for dependencies

Normally, 'bindSession()' is called on a 'FirModuleData' right after its
instantiation to make it fully initialized. For source sessions,
this is done directly during the session creation. However, module data
created for dependencies (see 'LLFirModuleData.dependencies' & friends)
stay uninitialized, breaking 'collectAllDependentSourceSessionsTo()'.

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

@@ -72,7 +72,7 @@ class LLFirBuiltinsSessionFactory(private val project: Project) {
         val builtinsModule = getBuiltinsModule(platform)
 
         val session = LLFirBuiltinsAndCloneableSession(builtinsModule, builtInTypes)
-        val moduleData = LLFirModuleData(builtinsModule).apply { bindSession(session) }
+        val moduleData = LLFirModuleData(session)
 
         return session.apply {
             registerIdeComponents(project)

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

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirModuleData
@@ -31,9 +33,11 @@ val FirBasedSymbol<*>.llFirModuleData: LLFirModuleData
     get() = fir.llFirModuleData
 
 
-class LLFirModuleData(
-    val ktModule: KtModule,
-) : FirModuleData() {
+class LLFirModuleData private constructor(val ktModule: KtModule) : FirModuleData() {
+    constructor(session: LLFirSession) : this(session.ktModule) {
+        bindSession(session)
+    }
+
     override val name: Name get() = Name.special("<${ktModule.moduleDescription}>")
 
     override val dependencies: List<FirModuleData> by lazy(LazyThreadSafetyMode.PUBLICATION) {
@@ -54,6 +58,9 @@ class LLFirModuleData(
 
     override val analyzerServices: PlatformDependentAnalyzerServices get() = ktModule.analyzerServices
 
+    override val session: FirSession
+        get() = boundSession ?: LLFirSessionCache.getInstance(ktModule.project).getSession(ktModule, preferBinary = true)
+
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
         if (javaClass != other?.javaClass) return false

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

@@ -616,7 +616,7 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
     }
 
     private fun createModuleData(session: LLFirSession): LLFirModuleData {
-        return LLFirModuleData(session.ktModule).apply { bindSession(session) }
+        return LLFirModuleData(session)
     }
 
     /**

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirModuleData.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir
 
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.platform.isCommon
@@ -58,16 +57,16 @@ abstract class FirModuleData : FirSessionComponent {
     open val capabilities: FirModuleCapabilities
         get() = FirModuleCapabilities.Empty
 
-    private var _session: FirSession? = null
-    val session: FirSession
-        get() = _session
-            ?: error("module data ${this::class.simpleName}:${name} not bound to session")
+    protected var boundSession: FirSession? = null
+        private set
+
+    abstract val session: FirSession
 
     fun bindSession(session: FirSession) {
-        if (_session != null) {
+        if (boundSession != null) {
             error("module data already bound to $this")
         }
-        _session = session
+        boundSession = session
     }
 
     override fun toString(): String {
@@ -84,7 +83,11 @@ class FirModuleDataImpl(
     override val analyzerServices: PlatformDependentAnalyzerServices,
     override val capabilities: FirModuleCapabilities = FirModuleCapabilities.Empty,
     override val isCommon: Boolean = platform.isCommon(),
-) : FirModuleData()
+) : FirModuleData() {
+    override val session: FirSession
+        get() = boundSession
+            ?: error("module data ${this::class.simpleName}:${name} not bound to session")
+}
 
 val FirSession.nullableModuleData: FirModuleData? by FirSession.nullableSessionComponentAccessor()
 val FirSession.moduleData: FirModuleData