[LL FIR] reuse source sessions

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

@@ -9,9 +9,7 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 
-@Suppress("unused")
 internal class LLFirGlobalResolveComponents(
-    val useSiteKtModule: KtModule,
     val project: Project,
 ) {
     val phaseRunner: LLFirPhaseRunner = LLFirPhaseRunner()

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

@@ -10,16 +10,17 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootModificationTracker
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.*
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionProviderStorage
-import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirSourceResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirLibraryOrLibrarySourceResolvableResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirNotUnderContentRootResolvableResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirSourceResolveSession
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
 import org.jetbrains.kotlin.analysis.utils.caches.strongCachedValue
-import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.locks.ReentrantReadWriteLock
+import kotlin.concurrent.withLock
 
 internal class LLFirResolveSessionService(project: Project) {
     private val sessionProviderStorage = LLFirSessionProviderStorage(project)
@@ -28,11 +29,22 @@ internal class LLFirResolveSessionService(project: Project) {
         project.createProjectWideOutOfBlockModificationTracker(),
         ProjectRootModificationTracker.getInstance(project),
     ) {
-        ConcurrentHashMap<KtModule, LLFirResolvableResolveSession>()
+        mutableMapOf<KtModule, LLFirResolvableResolveSession>()
     }
 
-    fun getFirResolveSession(module: KtModule): LLFirResolvableResolveSession =
-        stateCache.computeIfAbsent(module) { createFirResolveSessionFor(module, sessionProviderStorage) }
+    private val cacheLock = ReentrantReadWriteLock()
+
+    fun getFirResolveSession(module: KtModule): LLFirResolvableResolveSession {
+        cacheLock.readLock().withLock {
+            stateCache[module]?.let { return it }
+        }
+        cacheLock.writeLock().withLock {
+            stateCache[module]?.let { return it }
+            val session = createFirResolveSessionFor(module, sessionProviderStorage)
+            stateCache[module] = session
+            return session
+        }
+    }
 
     companion object {
         fun getInstance(project: Project): LLFirResolveSessionService =
@@ -54,6 +66,7 @@ internal class LLFirResolveSessionService(project: Project) {
                         sessionProvider,
                     )
                 }
+
                 is KtLibraryModule, is KtLibrarySourceModule -> {
                     LLFirLibraryOrLibrarySourceResolvableResolveSession(
                         useSiteSession.moduleComponents.globalResolveComponents,
@@ -62,6 +75,7 @@ internal class LLFirResolveSessionService(project: Project) {
                         sessionProvider,
                     )
                 }
+
                 is KtNotUnderContentRootModule -> {
                     LLFirNotUnderContentRootResolvableResolveSession(
                         useSiteSession.moduleComponents.globalResolveComponents,
@@ -70,6 +84,7 @@ internal class LLFirResolveSessionService(project: Project) {
                         sessionProvider,
                     )
                 }
+
                 else -> {
                     error("Unexpected $useSiteKtModule")
                 }

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

@@ -54,7 +54,7 @@ internal class LLFirNonUnderContentRootSessionFactory(private val project: Proje
         val builtinsSession = LLFirBuiltinsSessionFactory.getInstance(project).getBuiltinsSession(JvmPlatforms.unspecifiedJvmPlatform)
         val languageVersionSettings = LanguageVersionSettingsImpl.DEFAULT
         val scopeProvider = FirKotlinScopeProvider(::wrapScopeWithJvmMapped)
-        val globalResolveComponents = LLFirGlobalResolveComponents(module, project)
+        val globalResolveComponents = LLFirGlobalResolveComponents(project)
         val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider)
         val contentScope = module.contentScope
         val session = LLFirNonUnderContentRootSession(module, project, components, builtinsSession.builtinTypes)

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

@@ -11,30 +11,29 @@ import kotlinx.collections.immutable.PersistentMap
 import kotlinx.collections.immutable.persistentMapOf
 import kotlinx.collections.immutable.toPersistentMap
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirGlobalResolveComponents
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirBuiltinsSessionFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirLibrarySessionFactory
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirNonUnderContentRootSessionFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.firKtModuleBasedModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.addValueFor
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.executeWithoutPCE
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.analysis.providers.createModuleWithoutDependenciesOutOfBlockModificationTracker
 import java.util.concurrent.ConcurrentHashMap
-import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirBuiltinsSessionFactory
-import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirNonUnderContentRootSessionFactory
 
 class LLFirSessionProviderStorage(val project: Project) {
-    private val sessionsCache = ConcurrentHashMap<KtModule, FromModuleViewSessionCache>()
+    private val sessionsCache = LLFirSessionsCache()
 
     private val librariesSessionFactory = LLFirLibrarySessionFactory.getInstance(project)
-    private val builtInsSessionFactopry = LLFirBuiltinsSessionFactory.getInstance(project)
+    private val builtInsSessionFactory = LLFirBuiltinsSessionFactory.getInstance(project)
+
+    private val globalComponents = LLFirGlobalResolveComponents(project)
 
     fun getSessionProvider(
         useSiteKtModule: KtModule,
         configureSession: (LLFirSession.() -> Unit)? = null
     ): LLFirSessionProvider {
-        val globalComponents = LLFirGlobalResolveComponents(useSiteKtModule, project)
-
-        val cache = sessionsCache.getOrPut(useSiteKtModule) { FromModuleViewSessionCache() }
-        val (sessions, session) = cache.withMappings(project) { mappings ->
+        val (sessions, session) = sessionsCache.withMappings(project) { mappings ->
             val sessions = mutableMapOf<KtModule, LLFirResolvableModuleSession>().apply { putAll(mappings) }
             val session = executeWithoutPCE {
                 when (useSiteKtModule) {
@@ -43,26 +42,29 @@ class LLFirSessionProviderStorage(val project: Project) {
                             project,
                             useSiteKtModule,
                             globalComponents,
-                            cache.sessionInvalidator,
+                            sessionsCache.sessionInvalidator,
                             sessions,
                             librariesSessionFactory,
                             configureSession = configureSession,
                         )
                     }
+
                     is KtLibraryModule, is KtLibrarySourceModule -> {
                         LLFirSessionFactory.createLibraryOrLibrarySourceResolvableSession(
                             project,
                             useSiteKtModule,
                             globalComponents,
-                            cache.sessionInvalidator,
-                            builtInsSessionFactopry.getBuiltinsSession(useSiteKtModule.platform),
+                            sessionsCache.sessionInvalidator,
+                            builtInsSessionFactory.getBuiltinsSession(useSiteKtModule.platform),
                             sessions,
                             configureSession = configureSession,
                         )
                     }
+
                     is KtNotUnderContentRootModule ->
                         LLFirNonUnderContentRootSessionFactory.getInstance(project)
                             .getNonUnderContentRootSession(useSiteKtModule, sessions)
+
                     else -> error("Unexpected ${useSiteKtModule::class.simpleName}")
                 }
 
@@ -74,7 +76,7 @@ class LLFirSessionProviderStorage(val project: Project) {
     }
 }
 
-private class FromModuleViewSessionCache {
+private class LLFirSessionsCache {
     @Volatile
     private var mappings: PersistentMap<KtModule, FirSessionWithModificationTracker> = persistentMapOf()