[FIR] Add late initialization of session in FirModuleData

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/FirModuleInfoBasedModuleData.kt

@@ -7,27 +7,23 @@ package org.jetbrains.kotlin.fir.session
 
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.fir.FirModuleData
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
 
-class FirModuleInfoBasedModuleData(
+class FirModuleInfoBasedModuleData(val moduleInfo: ModuleInfo) : FirModuleData() {
-    val moduleInfo: ModuleInfo,
-    override val session: FirSession
-) : FirModuleData() {
     override val name: Name
         get() = moduleInfo.name
     override val dependencies: List<FirModuleData> = moduleInfo.dependencies()
         .filterNot { it == moduleInfo }
-        .map { FirModuleInfoBasedModuleData(it, session) }
+        .map { FirModuleInfoBasedModuleData(it) }
 
     override val dependsOnDependencies: List<FirModuleData> = moduleInfo.expectedBy
         .filterNot { it == moduleInfo }
-        .map { FirModuleInfoBasedModuleData(it, session) }
+        .map { FirModuleInfoBasedModuleData(it) }
 
     override val friendDependencies: List<FirModuleData> = moduleInfo.modulesWhoseInternalsAreVisible()
         .filterNot { it == moduleInfo }
-        .map { FirModuleInfoBasedModuleData(it, session) }
+        .map { FirModuleInfoBasedModuleData(it) }
     override val platform: TargetPlatform
         get() = moduleInfo.platform
 

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/FirSessionFactory.kt

@@ -78,8 +78,8 @@ object FirSessionFactory {
         lookupTracker: LookupTracker? = null,
         init: FirSessionConfigurator.() -> Unit = {}
     ): FirJavaModuleBasedSession {
-        return FirJavaModuleBasedSession(moduleInfo, sessionProvider).apply {
+        return FirJavaModuleBasedSession(moduleInfo, sessionProvider).apply session@{
-            val moduleData = FirModuleInfoBasedModuleData(moduleInfo, this)
+            val moduleData = FirModuleInfoBasedModuleData(moduleInfo).apply { bindSession(this@session) }
             registerModuleData(moduleData)
             registerCliCompilerOnlyComponents()
             registerCommonComponents(languageVersionSettings)
@@ -137,8 +137,8 @@ object FirSessionFactory {
         packagePartProvider: PackagePartProvider,
         languageVersionSettings: LanguageVersionSettings = LanguageVersionSettingsImpl.DEFAULT,
     ): FirLibrarySession {
-        return FirLibrarySession(moduleInfo, sessionProvider).apply {
+        return FirLibrarySession(moduleInfo, sessionProvider).apply session@{
-            val moduleData = FirModuleInfoBasedModuleData(moduleInfo, this)
+            val moduleData = FirModuleInfoBasedModuleData(moduleInfo).apply { bindSession(this@session) }
             registerCliCompilerOnlyComponents()
             registerCommonComponents(languageVersionSettings)
 

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

@@ -14,7 +14,17 @@ abstract class FirModuleData : FirSessionComponent {
     abstract val dependsOnDependencies: List<FirModuleData>
     abstract val friendDependencies: List<FirModuleData>
     abstract val platform: TargetPlatform
-    abstract val session: FirSession
+
+    private var _session: FirSession? = null
+    val session: FirSession
+        get() = _session!!
+
+    fun bindSession(session: FirSession) {
+        if (_session != null) {
+            error("session already bound to $this")
+        }
+        _session = session
+    }
 }
 
 class FirModuleDataImpl(
@@ -22,8 +32,7 @@ class FirModuleDataImpl(
     override val dependencies: List<FirModuleData>,
     override val dependsOnDependencies: List<FirModuleData>,
     override val friendDependencies: List<FirModuleData>,
-    override val platform: TargetPlatform,
+    override val platform: TargetPlatform
-    override val session: FirSession,
 ) : FirModuleData()
 
 val FirSession.moduleData: FirModuleData by FirSession.sessionComponentAccessor()

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeSessionFactory.kt

@@ -74,8 +74,8 @@ internal object FirIdeSessionFactory {
         sessionsCache[moduleInfo] = session
 
 
-        return session.apply {
+        return session.apply session@{
-            val moduleData = FirModuleInfoBasedModuleData(moduleInfo, this)
+            val moduleData = FirModuleInfoBasedModuleData(moduleInfo).apply { bindSession(this@session) }
             registerModuleData(moduleData)
 
             val cache = ModuleFileCacheImpl(this)
@@ -107,7 +107,7 @@ internal object FirIdeSessionFactory {
                     this,
                     providers = listOf(
                         provider.symbolProvider,
-                        JavaSymbolProvider(this@apply, moduleData, project, searchScope),
+                        JavaSymbolProvider(this@session, moduleData, project, searchScope),
                     ),
                     dependentProviders = buildList {
                         add(
@@ -171,8 +171,8 @@ internal object FirIdeSessionFactory {
         val packagePartProvider = IDEPackagePartProvider(searchScope)
 
         val kotlinClassFinder = VirtualFileFinderFactory.getInstance(project).create(searchScope)
-        FirIdeLibrariesSession(moduleInfo, project, searchScope, builtinTypes).apply {
+        FirIdeLibrariesSession(moduleInfo, project, searchScope, builtinTypes).apply session@{
-            val moduleData = FirModuleInfoBasedModuleData(moduleInfo, this)
+            val moduleData = FirModuleInfoBasedModuleData(moduleInfo).apply { bindSession(this@session) }
 
             registerIdeComponents()
             register(FirPhaseManager::class, FirPhaseCheckingPhaseManager)
@@ -216,15 +216,16 @@ internal object FirIdeSessionFactory {
         languageVersionSettings: LanguageVersionSettings = LanguageVersionSettingsImpl.DEFAULT,
         configureSession: (FirIdeSession.() -> Unit)? = null,
     ): FirIdeBuiltinsAndCloneableSession {
-        return FirIdeBuiltinsAndCloneableSession(project, builtinTypes).apply {
+        return FirIdeBuiltinsAndCloneableSession(project, builtinTypes).apply session@{
             val moduleData = FirModuleDataImpl(
                 Name.special("<builtins module>"),
                 emptyList(),
                 emptyList(),
                 emptyList(),
-                JvmPlatforms.unspecifiedJvmPlatform,
+                JvmPlatforms.unspecifiedJvmPlatform
-                this
+            ).apply {
-            )
+                bindSession(this@session)
+            }
             registerIdeComponents()
             register(FirPhaseManager::class, FirPhaseCheckingPhaseManager)
             registerCommonComponents(languageVersionSettings)