FIR IDE: fix session provider implementation as it only provides only 2 sessions

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.idea.util.psiModificationTrackerBasedCachedValue
 import java.util.concurrent.ConcurrentHashMap
 import org.jetbrains.kotlin.idea.caches.project.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeModuleLibraryDependenciesSession
+import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
 
 internal class FirIdeResolveStateService(private val project: Project) {
@@ -19,19 +20,22 @@ internal class FirIdeResolveStateService(private val project: Project) {
     }
 
     private fun createResolveStateFor(moduleInfo: IdeaModuleInfo): FirModuleResolveStateImpl {
-        val sessionProvider = FirIdeSessionProvider(project)
         val firPhaseRunner = FirPhaseRunner()
+        val sessionProvider = FirIdeSessionProvider(project)
 
         val librariesSession = FirIdeModuleLibraryDependenciesSession.create(moduleInfo as ModuleSourceInfo, sessionProvider, project)
-
         val sourcesSession = FirIdeSourcesSession.create(
             project,
             moduleInfo,
             firPhaseRunner,
             sessionProvider,
             librariesSession,
-        ).apply { sessionProvider.registerSession(moduleInfo, this) }
+        )
 
+        sessionProvider.apply {
+            setSourcesSession(sourcesSession)
+            setLibrariesSession(librariesSession)
+        }
 
         return FirModuleResolveStateImpl(
             moduleInfo,

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

@@ -1,30 +0,0 @@
-/*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * 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.idea.fir.low.level.api
-
-import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.analyzer.ModuleInfo
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
-import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
-import org.jetbrains.kotlin.psi.KtElement
-import java.util.concurrent.ConcurrentHashMap
-import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
-
-@ThreadSafe
-internal class FirIdeSessionProvider(project: Project) : FirProjectSessionProvider(project) {
-    override val sessionCache = ConcurrentHashMap<ModuleInfo, FirSession>()
-
-    fun getSession(psi: KtElement): FirSession {
-        val moduleInfo = psi.getModuleInfo()
-        return getSession(moduleInfo)
-    }
-
-    override fun getSession(moduleInfo: ModuleInfo): FirSession {
-        return sessionCache[moduleInfo]
-            ?: error("$moduleInfo")
-    }
-}

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.PsiToFirCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/FirFileBuilder.kt

@@ -9,9 +9,7 @@ import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.createTransformerBasedProcessorByPhase
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeSessionProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
@@ -22,7 +20,6 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
  */
 @ThreadSafe
 internal class FirFileBuilder(
-    private val sessionProvider: FirIdeSessionProvider,
     private val scopeProvider: FirScopeProvider,
     private val firPhaseRunner: FirPhaseRunner
 ) {

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

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.sessions
 
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
-import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
 import org.jetbrains.kotlin.fir.java.deserialization.KotlinDeserializedJvmSymbolsProvider
 import org.jetbrains.kotlin.fir.registerCommonComponents
@@ -19,7 +18,6 @@ import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
 import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.resolve.IDEPackagePartProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeSessionProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.IdeSessionComponents
 import org.jetbrains.kotlin.load.java.JavaClassFinderImpl
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinderFactory

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

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.sessions
 
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeSessionProvider
 
 internal abstract class FirIdeSession(sessionProvider: FirIdeSessionProvider) : FirSession(sessionProvider) {
     abstract val scope: GlobalSearchScope

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.idea.fir.low.level.api.sessions
+
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionProvider
+import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
+import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
+import org.jetbrains.kotlin.idea.caches.project.isLibraryClasses
+import org.jetbrains.kotlin.psi.KtElement
+import java.util.concurrent.ConcurrentHashMap
+import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
+
+@ThreadSafe
+internal class FirIdeSessionProvider(
+    override val project: Project,
+) : FirSessionProvider {
+    private lateinit var sourcesSession: FirIdeSourcesSession
+    private lateinit var librariesSession: FirIdeModuleLibraryDependenciesSession
+
+
+    fun setSourcesSession(sourcesSession: FirIdeSourcesSession) {
+        check(!this::sourcesSession.isInitialized)
+        this.sourcesSession = sourcesSession
+    }
+
+    fun setLibrariesSession(librariesSession: FirIdeModuleLibraryDependenciesSession) {
+        check(!this::librariesSession.isInitialized)
+        this.librariesSession = librariesSession
+    }
+
+    override fun getSession(moduleInfo: ModuleInfo): FirSession = when {
+        moduleInfo is IdeaModuleInfo && moduleInfo.isLibraryClasses() -> {
+            librariesSession /* TODO check if library is in libraries session scope */
+        }
+        moduleInfo is ModuleSourceInfo -> {
+            sourcesSession /* TODO check if source is in sources session scope */
+        }
+        else -> error("Invalid module info $moduleInfo")
+    }
+}

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

@@ -24,7 +24,6 @@ import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
 import org.jetbrains.kotlin.fir.resolve.transformers.PhasedFirFileResolver
 import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
-import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeSessionProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.IdePhasedFirFileResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
@@ -49,11 +48,11 @@ internal class FirIdeSourcesSession private constructor(
             project: Project,
             moduleInfo: ModuleSourceInfo,
             firPhaseRunner: FirPhaseRunner,
-            sessionProvider: FirSessionProvider,
+            sessionProvider: FirIdeSessionProvider,
             librariesSession: FirIdeModuleLibraryDependenciesSession,
         ): FirIdeSourcesSession {
             val scopeProvider = KotlinScopeProvider(::wrapScopeWithJvmMapped)
-            val firBuilder = FirFileBuilder(sessionProvider as FirIdeSessionProvider, scopeProvider, firPhaseRunner)
+            val firBuilder = FirFileBuilder(scopeProvider, firPhaseRunner)
             val dependentModules = moduleInfo.collectTransitiveDependenciesWithSelf().filterIsInstance<ModuleSourceInfo>()
             val searchScope = ModuleWithDependentsScope(project, dependentModules.map { it.module })
             return FirIdeSourcesSession(sessionProvider, searchScope, firBuilder).apply {