[LL API] Unify cache for 'not-under content root' modules

Before, 'FirSession's for not-under content root modules were stored in
a project-global service, which essentially means that even if one
creates a non-cached 'LLFirResolveSession', 'FirSession's for such
modules will be cached and reused later. Such exact behavior happened in
sealed class inheritor pre-analysis handler.

analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/RegisterComponentService.java

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.analysis.api.standalone;
 import com.intellij.mock.MockProject;
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionService;
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirLibrarySessionFactory;
-import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirNonUnderContentRootSessionFactory;
 
 class RegisterComponentService {
     static void registerLLFirLibrarySessionFactory(MockProject project) {
@@ -18,13 +17,6 @@ class RegisterComponentService {
         );
     }
 
-    static void registerLLFirNonUnderContentRootSessionFactory(MockProject project) {
-        project.registerService(
-                LLFirNonUnderContentRootSessionFactory.class,
-                new LLFirNonUnderContentRootSessionFactory(project)
-        );
-    }
-
     static void registerLLFirResolveSessionService(MockProject project) {
         project.registerService(
                 LLFirResolveSessionService.class,

analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt

@@ -138,7 +138,6 @@ public class StandaloneAnalysisAPISessionBuilder(
             )
             registerService(LLFirBuiltinsSessionFactory::class.java, LLFirBuiltinsSessionFactory(this))
             RegisterComponentService.registerLLFirLibrarySessionFactory(this)
-            RegisterComponentService.registerLLFirNonUnderContentRootSessionFactory(this)
 
             registerService(KotlinReferenceProvidersService::class.java, HLApiReferenceProviderService::class.java)
             registerService(KotlinReferenceProviderContributor::class.java, KotlinFirReferenceContributor::class.java)

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

@@ -1,118 +0,0 @@
-/*
- * Copyright 2010-2022 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.analysis.low.level.api.fir.project.structure
-
-import com.intellij.openapi.project.Project
-import com.intellij.psi.util.PsiModificationTracker
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirGlobalResolveComponents
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirLazyDeclarationResolver
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
-import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirDependentModuleProvidersBySessions
-import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirModuleWithDependenciesSymbolProvider
-import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirProvider
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirNonUnderContentRootSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionInvalidator
-import org.jetbrains.kotlin.analysis.project.structure.KtNotUnderContentRootModule
-import org.jetbrains.kotlin.analysis.providers.createPackageProvider
-import org.jetbrains.kotlin.analysis.utils.caches.SoftCachedMap
-import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl
-import org.jetbrains.kotlin.fir.PrivateSessionConstructor
-import org.jetbrains.kotlin.fir.SessionConfiguration
-import org.jetbrains.kotlin.fir.backend.jvm.FirJvmTypeMapper
-import org.jetbrains.kotlin.fir.extensions.FirEmptyPredicateBasedProvider
-import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
-import org.jetbrains.kotlin.fir.extensions.FirRegisteredPluginAnnotations
-import org.jetbrains.kotlin.fir.resolve.providers.DEPENDENCIES_SYMBOL_PROVIDER_QUALIFIED_KEY
-import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
-import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
-import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
-import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
-import org.jetbrains.kotlin.fir.session.*
-import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
-import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
-
-internal class LLFirNonUnderContentRootSessionFactory(private val project: Project) {
-    private val cache = SoftCachedMap.create<KtNotUnderContentRootModule, LLFirResolvableModuleSession>(
-        project,
-        kind = SoftCachedMap.Kind.STRONG_KEYS_SOFT_VALUES,
-        trackers = listOf(PsiModificationTracker.MODIFICATION_COUNT),
-    )
-
-    fun getNonUnderContentRootSession(
-        module: KtNotUnderContentRootModule,
-    ): LLFirResolvableModuleSession {
-        return cache.getOrPut(module) { createSession(module, LLFirSessionInvalidator { cache.clear() }) }
-    }
-
-    @OptIn(PrivateSessionConstructor::class, SessionConfiguration::class)
-    private fun createSession(
-        module: KtNotUnderContentRootModule,
-        sessionInvalidator: LLFirSessionInvalidator
-    ): LLFirResolvableModuleSession {
-        val builtinsSession = LLFirBuiltinsSessionFactory.getInstance(project).getBuiltinsSession(JvmPlatforms.unspecifiedJvmPlatform)
-        val languageVersionSettings = LanguageVersionSettingsImpl.DEFAULT
-        val scopeProvider = FirKotlinScopeProvider(::wrapScopeWithJvmMapped)
-        val globalResolveComponents = LLFirGlobalResolveComponents(project)
-        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider, sessionInvalidator)
-        val contentScope = module.contentScope
-        val session = LLFirNonUnderContentRootSession(module, project, components, builtinsSession.builtinTypes)
-        components.session = session
-
-        return session.apply session@{
-            val moduleData = LLFirModuleData(module).apply { bindSession(this@session) }
-            registerModuleData(moduleData)
-            register(FirKotlinScopeProvider::class, scopeProvider)
-
-            registerIdeComponents(project)
-            registerCommonComponents(languageVersionSettings)
-            registerCommonJavaComponents(JavaModuleResolver.getInstance(project))
-            registerResolveComponents()
-            registerJavaSpecificResolveComponents()
-
-            val ktFile = module.file as? KtFile
-
-            val provider = LLFirProvider(
-                this,
-                components,
-                if (ktFile != null) FileBasedKotlinDeclarationProvider(ktFile) else EmptyKotlinDeclarationProvider,
-                project.createPackageProvider(contentScope),
-                canContainKotlinPackage = true,
-            )
-
-            register(FirProvider::class, provider)
-            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver())
-
-            val dependencyProvider = LLFirDependentModuleProvidersBySessions(this) {
-                add(builtinsSession)
-            }
-
-            register(
-                FirSymbolProvider::class,
-                LLFirModuleWithDependenciesSymbolProvider(
-                    this,
-                    dependencyProvider,
-                    providers = listOfNotNull(
-                        provider.symbolProvider,
-                    )
-                )
-            )
-
-            register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)
-            register(DEPENDENCIES_SYMBOL_PROVIDER_QUALIFIED_KEY, dependencyProvider)
-            register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
-            register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.Empty)
-        }
-    }
-
-    companion object {
-        fun getInstance(project: Project): LLFirNonUnderContentRootSessionFactory =
-            project.getService(LLFirNonUnderContentRootSessionFactory::class.java)
-    }
-}
-

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

@@ -7,13 +7,11 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions
 
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder
 import org.jetbrains.kotlin.analysis.project.structure.KtNotUnderContentRootModule
-import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.fir.BuiltinTypes
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor
 
-internal class LLFirNonUnderContentRootSession @PrivateSessionConstructor constructor(
+internal class LLFirNonUnderContentRootResolvableModuleSession @PrivateSessionConstructor constructor(
     override val ktModule: KtNotUnderContentRootModule,
     override val project: Project,
     override val moduleComponents: LLFirModuleResolveComponents,

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

@@ -37,6 +37,7 @@ import org.jetbrains.kotlin.fir.session.*
 import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
 import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
 
@@ -247,7 +248,75 @@ internal object LLFirSessionFactory {
 
             configureSession?.invoke(this)
         }
+    }
 
+    fun createNotUnderContentRootResolvableSession(
+        project: Project,
+        module: KtNotUnderContentRootModule,
+        sessionInvalidator: LLFirSessionInvalidator,
+        sessionsCache: MutableMap<KtModule, LLFirSession>,
+        configureSession: (LLFirSession.() -> Unit)? = null
+    ): LLFirNonUnderContentRootResolvableModuleSession {
+        sessionsCache[module]?.let { return it as LLFirNonUnderContentRootResolvableModuleSession }
+        checkCanceled()
+
+        val builtinsSession = LLFirBuiltinsSessionFactory.getInstance(project).getBuiltinsSession(JvmPlatforms.unspecifiedJvmPlatform)
+        val languageVersionSettings = LanguageVersionSettingsImpl.DEFAULT
+        val scopeProvider = FirKotlinScopeProvider(::wrapScopeWithJvmMapped)
+        val globalResolveComponents = LLFirGlobalResolveComponents(project)
+        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider, sessionInvalidator)
+        val contentScope = module.contentScope
+
+        val session = LLFirNonUnderContentRootResolvableModuleSession(module, project, components, builtinsSession.builtinTypes)
+        sessionsCache[module] = session
+        components.session = session
+
+        return session.apply session@{
+            val moduleData = LLFirModuleData(module).apply { bindSession(this@session) }
+            registerModuleData(moduleData)
+            register(FirKotlinScopeProvider::class, scopeProvider)
+
+            registerIdeComponents(project)
+            registerCommonComponents(languageVersionSettings)
+            registerCommonJavaComponents(JavaModuleResolver.getInstance(project))
+            registerResolveComponents()
+            registerJavaSpecificResolveComponents()
+
+            val ktFile = module.file as? KtFile
+
+            val provider = LLFirProvider(
+                this,
+                components,
+                if (ktFile != null) FileBasedKotlinDeclarationProvider(ktFile) else EmptyKotlinDeclarationProvider,
+                project.createPackageProvider(contentScope),
+                canContainKotlinPackage = true,
+            )
+
+            register(FirProvider::class, provider)
+            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver())
+
+            val dependencyProvider = LLFirDependentModuleProvidersBySessions(this) {
+                add(builtinsSession)
+            }
+
+            register(
+                FirSymbolProvider::class,
+                LLFirModuleWithDependenciesSymbolProvider(
+                    this,
+                    dependencyProvider,
+                    providers = listOfNotNull(
+                        provider.symbolProvider,
+                    )
+                )
+            )
+
+            register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)
+            register(DEPENDENCIES_SYMBOL_PROVIDER_QUALIFIED_KEY, dependencyProvider)
+            register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
+            register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.Empty)
+
+            configureSession?.invoke(this)
+        }
     }
 }
 

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

@@ -13,7 +13,6 @@ 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.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.addValueFor
 import org.jetbrains.kotlin.analysis.project.structure.*
@@ -25,6 +24,7 @@ import java.lang.ref.SoftReference
 class LLFirSessionProviderStorage(val project: Project) {
     private val sourceAsUseSiteSessionCache = LLFirSessionsCache()
     private val libraryAsUseSiteSessionCache = LLFirSessionsCache()
+    private val notUnderContentRootSessionCache = LLFirSessionsCache()
 
     private val librariesSessionFactory = LLFirLibrarySessionFactory.getInstance(project)
     private val builtInsSessionFactory = LLFirBuiltinsSessionFactory.getInstance(project)
@@ -44,9 +44,7 @@ class LLFirSessionProviderStorage(val project: Project) {
         }
 
         is KtNotUnderContentRootModule -> {
-            val session = LLFirNonUnderContentRootSessionFactory.getInstance(project)
-                .getNonUnderContentRootSession(useSiteKtModule)
-            LLFirSessionProvider(project, session, KtModuleToSessionMappingByMapImpl(mapOf(useSiteKtModule to session)))
+            createSessionProviderForNotUnderContentRootSession(useSiteKtModule, configureSession)
         }
 
         else -> error("Unexpected ${useSiteKtModule::class.simpleName}")
@@ -93,6 +91,24 @@ class LLFirSessionProviderStorage(val project: Project) {
         }
         return LLFirSessionProvider(project, session, KtModuleToSessionMappingByWeakValueMapImpl(sessions))
     }
+
+    private fun createSessionProviderForNotUnderContentRootSession(
+        useSiteKtModule: KtNotUnderContentRootModule,
+        configureSession: (LLFirSession.() -> Unit)?
+    ): LLFirSessionProvider {
+        val (sessions, session) = notUnderContentRootSessionCache.withMappings(project) { mappings ->
+            val sessions = mutableMapOf<KtModule, LLFirSession>().apply { putAll(mappings) }
+            val session = LLFirSessionFactory.createNotUnderContentRootResolvableSession(
+                project,
+                useSiteKtModule,
+                libraryAsUseSiteSessionCache.sessionInvalidator,
+                sessions,
+                configureSession = configureSession,
+            )
+            sessions to session
+        }
+        return LLFirSessionProvider(project, session, KtModuleToSessionMappingByWeakValueMapImpl(sessions))
+    }
 }
 
 private class LLFirSessionsCache {
@@ -175,6 +191,7 @@ private class FirSessionWithModificationTracker(
 
         val outOfBlockTracker = when (ktModule) {
             is KtSourceModule -> trackerFactory.createModuleWithoutDependenciesOutOfBlockModificationTracker(ktModule)
+            is KtNotUnderContentRootModule -> trackerFactory.createProjectWideOutOfBlockModificationTracker()
             else -> null
         }
         modificationTracker = CompositeModificationTracker.create(

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/test/base/AnalysisApiFirTestServiceRegistrar.kt

@@ -17,7 +17,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionServic
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.services.FirSealedClassInheritorsProcessorFactory
 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.services.KtCompilerPluginsProviderForTests
 import org.jetbrains.kotlin.analysis.low.level.api.fir.services.LLFirSealedClassInheritorsProcessorFactoryForTests
 import org.jetbrains.kotlin.analysis.low.level.api.fir.services.PackagePartProviderTestImpl
@@ -50,7 +49,6 @@ object AnalysisApiFirTestServiceRegistrar : AnalysisApiTestServiceRegistrar() {
             registerService(LLFirResolveSessionService::class.java)
             registerService(LLFirLibrarySessionFactory::class.java)
             registerService(LLFirBuiltinsSessionFactory::class.java)
-            registerService(LLFirNonUnderContentRootSessionFactory::class.java)
             registerService(PackagePartProviderFactory::class.java, PackagePartProviderTestImpl(testServices))
 
             registerService(KotlinAsJavaSupport::class.java, SymbolKotlinAsJavaSupport(project))