[LL FIR] incrementally update module caches

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/KotlinModificationTrackerFactory.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.ModificationTracker
 import org.jetbrains.annotations.TestOnly
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 
 
@@ -62,10 +63,40 @@ public abstract class KotlinModificationTrackerFactory {
      * See [KotlinModificationTrackerFactory] for the definition of **OOBM**.
      * @see ModificationTracker
      */
-    public abstract fun createLibrariesModificationTracker(): ModificationTracker
+    public abstract fun createLibrariesWideModificationTracker(): ModificationTracker
+
+    /**
+     * Creates [KtModuleStateTracker] which is incremented every time [module] roots changes.
+     *
+     * @see KtModuleStateTracker
+     */
+    public abstract fun createModuleStateTracker(module: KtModule): KtModuleStateTracker
+
 
     @TestOnly
     public abstract fun incrementModificationsCount()
+
+    public companion object {
+        public fun getService(project: Project): KotlinModificationTrackerFactory =
+            project.getService(KotlinModificationTrackerFactory::class.java)
+    }
+}
+
+/**
+ * Represents current state of [KtModule] validity, can be created via [org.jetbrains.kotlin.analysis.providers.KotlinModificationTrackerFactory.createModuleStateTracker]
+ */
+public interface KtModuleStateTracker {
+    /**
+     * If the module is still valid (i.e., it was not removed)
+     */
+    public val isValid: Boolean
+
+    /**
+     * Represents modification tracker of modified roots similar to the [ModificationTracker].
+     *
+     * If the [isValid] == `false`, when behaviour is unspecified
+     */
+    public val rootModificationCount: Long
 }
 
 /**
@@ -94,6 +125,6 @@ public fun KtSourceModule.createModuleWithoutDependenciesOutOfBlockModificationT
  * See [KotlinModificationTrackerFactory] for the definition of **OOBM**.
  * @see ModificationTracker
  */
-public fun Project.createLibrariesModificationTracker(): ModificationTracker =
+public fun Project.createAllLibrariesModificationTracker(): ModificationTracker =
     ServiceManager.getService(this, KotlinModificationTrackerFactory::class.java)
-        .createLibrariesModificationTracker()
+        .createLibrariesWideModificationTracker()

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticModificationTrackerFactory.kt

@@ -8,13 +8,17 @@ package org.jetbrains.kotlin.analysis.providers.impl
 import com.intellij.openapi.util.ModificationTracker
 import com.intellij.openapi.util.SimpleModificationTracker
 import org.jetbrains.annotations.TestOnly
+import org.jetbrains.kotlin.analysis.project.structure.KtBinaryModule
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.providers.KotlinModificationTrackerFactory
+import org.jetbrains.kotlin.analysis.providers.KtModuleStateTracker
 
 public class KotlinStaticModificationTrackerFactory : KotlinModificationTrackerFactory() {
     private val projectWide = SimpleModificationTracker()
-    private val library = SimpleModificationTracker()
+    private val librariesWide = SimpleModificationTracker()
-    private val forModule = mutableMapOf<KtSourceModule, SimpleModificationTracker>()
+    private val moduleOutOfBlock = mutableMapOf<KtSourceModule, SimpleModificationTracker>()
+    private val moduleState = mutableMapOf<KtModule, KtModuleStateTrackerImpl>()
 
     override fun createProjectWideOutOfBlockModificationTracker(): ModificationTracker {
         return projectWide
@@ -22,17 +26,35 @@ public class KotlinStaticModificationTrackerFactory : KotlinModificationTrackerF
 
 
     override fun createModuleWithoutDependenciesOutOfBlockModificationTracker(module: KtSourceModule): ModificationTracker {
-        return forModule.getOrPut(module) { SimpleModificationTracker() }
+        return moduleOutOfBlock.getOrPut(module) { SimpleModificationTracker() }
     }
 
-    override fun createLibrariesModificationTracker(): ModificationTracker {
+    override fun createLibrariesWideModificationTracker(): ModificationTracker {
-        return library
+        return librariesWide
+    }
+
+    override fun createModuleStateTracker(module: KtModule): KtModuleStateTracker {
+        return moduleState.getOrPut(module) { KtModuleStateTrackerImpl() }
     }
 
     @TestOnly
     override fun incrementModificationsCount() {
         projectWide.incModificationCount()
-        library.incModificationCount()
+        librariesWide.incModificationCount()
-        forModule.values.forEach { it.incModificationCount() }
+        moduleOutOfBlock.values.forEach { it.incModificationCount() }
+        moduleState.values.forEach { it.incModificationCount() }
+    }
+}
+
+private class KtModuleStateTrackerImpl: KtModuleStateTracker {
+    override val isValid: Boolean get() = true
+
+    private var _rootModificationCount = 0L
+
+    override val rootModificationCount: Long get() = _rootModificationCount
+
+    @TestOnly
+    fun incModificationCount() {
+        _rootModificationCount++
     }
 }

analysis/analysis-internal-utils/src/org/jetbrains/kotlin/analysis/utils/trackers/CompositeModificationTracker.kt

@@ -0,0 +1,20 @@
+/*
+ * 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.utils.trackers
+
+import com.intellij.openapi.util.ModificationTracker
+
+public class CompositeModificationTracker private constructor(private val trackers: List<ModificationTracker>) : ModificationTracker {
+    override fun getModificationCount(): Long = trackers.sumOf { it.modificationCount }
+
+    public companion object {
+        public fun create(trackers: List<ModificationTracker>): ModificationTracker = when (trackers.size) {
+            0 -> ModificationTracker.NEVER_CHANGED
+            1 -> trackers.single()
+            else -> CompositeModificationTracker(trackers)
+        }
+    }
+}

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

@@ -9,9 +9,9 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirFirClassByPsiClassProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirLibrarySessionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirLibrarySession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.KtBinaryModule
-import org.jetbrains.kotlin.analysis.providers.createLibrariesModificationTracker
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.utils.caches.softCachedValue
 import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor
 import org.jetbrains.kotlin.fir.SessionConfiguration
@@ -28,22 +28,16 @@ import org.jetbrains.kotlin.fir.session.registerJavaSpecificResolveComponents
 import org.jetbrains.kotlin.fir.session.registerModuleData
 import org.jetbrains.kotlin.fir.symbols.FirPhaseManager
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
-import java.util.concurrent.ConcurrentHashMap
-import org.jetbrains.kotlin.analysis.utils.caches.getValue
 
 @OptIn(PrivateSessionConstructor::class, SessionConfiguration::class)
 internal class LLFirLibrarySessionFactory(
     private val project: Project,
 ) {
-    private val librarySessionByModule by softCachedValue(project, project.createLibrariesModificationTracker()) {
-        ConcurrentHashMap<KtBinaryModule, LLFirLibrarySession>()
-    }
 
-    fun getLibrarySession(ktBinaryModule: KtBinaryModule): LLFirLibrarySession {
+    fun getLibrarySession(ktBinaryModule: KtBinaryModule, sessionsCache: MutableMap<KtModule, LLFirSession>): LLFirLibrarySession {
-        return librarySessionByModule.getOrPut(ktBinaryModule) { createModuleLibrariesSession(ktBinaryModule) }
+        return sessionsCache.getOrPut(ktBinaryModule) { createModuleLibrariesSession(ktBinaryModule) } as LLFirLibrarySession
     }
 
-
     private fun createModuleLibrariesSession(
         ktLibraryModule: KtBinaryModule,
     ): LLFirLibrarySession {

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

@@ -15,7 +15,6 @@ 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.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtNotUnderContentRootModule
 import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.createPackageProvider
@@ -41,9 +40,8 @@ internal class LLFirNonUnderContentRootSessionFactory(private val project: Proje
 
     fun getNonUnderContentRootSession(
         module: KtNotUnderContentRootModule,
-        sessions: MutableMap<KtModule, LLFirResolvableModuleSession>
     ): LLFirResolvableModuleSession {
-        return cache.getOrPut(module) { createSession(module, LLFirSessionInvalidator { cache.clear() }) }.also { sessions[module] = it }
+        return cache.getOrPut(module) { createSession(module, LLFirSessionInvalidator { cache.clear() }) }
     }
 
     @OptIn(PrivateSessionConstructor::class, SessionConfiguration::class)

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

@@ -49,7 +49,7 @@ internal object LLFirSessionFactory {
         module: KtSourceModule,
         globalResolveComponents: LLFirGlobalResolveComponents,
         sessionInvalidator: LLFirSessionInvalidator,
-        sessionsCache: MutableMap<KtModule, LLFirResolvableModuleSession>,
+        sessionsCache: MutableMap<KtModule, LLFirSession>,
         librariesSessionFactory: LLFirLibrarySessionFactory,
         configureSession: (LLFirSession.() -> Unit)? = null
     ): LLFirSourcesSession {
@@ -114,7 +114,7 @@ internal object LLFirSessionFactory {
                 module.directRegularDependencies.mapNotNullTo(this) { dependency ->
                     when (dependency) {
                         is KtBuiltinsModule -> null //  build in is already added
-                        is KtBinaryModule -> LLFirLibrarySessionFactory.getInstance(project).getLibrarySession(dependency)
+                        is KtBinaryModule -> LLFirLibrarySessionFactory.getInstance(project).getLibrarySession(dependency, sessionsCache)
                         is KtSourceModule -> {
                             createSourcesSession(
                                 project,
@@ -164,7 +164,7 @@ internal object LLFirSessionFactory {
         globalComponents: LLFirGlobalResolveComponents,
         sessionInvalidator: LLFirSessionInvalidator,
         builtinSession: LLFirBuiltinsAndCloneableSession,
-        sessionsCache: MutableMap<KtModule, LLFirResolvableModuleSession>,
+        sessionsCache: MutableMap<KtModule, LLFirSession>,
         languageVersionSettings: LanguageVersionSettings = LanguageVersionSettingsImpl.DEFAULT,
         configureSession: (LLFirSession.() -> Unit)? = null
     ): LLFirLibraryOrLibrarySourceResolvableModuleSession {

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

@@ -19,24 +19,19 @@ import org.jetbrains.kotlin.fir.resolve.providers.dependenciesSymbolProvider
 class LLFirSessionProvider internal constructor(
     val project: Project,
     internal val rootModuleSession: LLFirResolvableModuleSession,
-    private val moduleToResolvableSession: Map<KtModule, LLFirResolvableModuleSession>
+    private val ktModuleToSession: Map<KtModule, LLFirSession>
 ) : FirSessionProvider() {
-
-    private val moduleToSession = moduleToResolvableSession + moduleToResolvableSession.values.flatMap { module ->
-        (module.dependenciesSymbolProvider as LLFirDependentModuleProviders).dependentSessions
-    }.associateBy { it.ktModule }
-
     override fun getSession(moduleData: FirModuleData): LLFirSession {
         requireIsInstance<LLFirModuleData>(moduleData)
         return getResolvableSession(moduleData.ktModule)
     }
 
     fun getSession(module: KtModule): LLFirSession =
-        moduleToSession.getValue(module)
+        ktModuleToSession.getValue(module)
 
     fun getResolvableSession(module: KtModule): LLFirResolvableModuleSession =
-        moduleToResolvableSession.getValue(module)
+        ktModuleToSession.getValue(module) as LLFirResolvableModuleSession
 
-    val allSessions: Collection<LLFirModuleSession>
+    val allSessions: Collection<LLFirSession>
-        get() = moduleToResolvableSession.values
+        get() = ktModuleToSession.values
 }

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

@@ -18,10 +18,13 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirMo
 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 org.jetbrains.kotlin.analysis.providers.KotlinModificationTrackerFactory
+import org.jetbrains.kotlin.analysis.providers.KtModuleStateTracker
+import org.jetbrains.kotlin.analysis.utils.trackers.CompositeModificationTracker
 
 class LLFirSessionProviderStorage(val project: Project) {
-    private val sessionsCache = LLFirSessionsCache()
+    private val sourceAsUseSiteSessionCache = LLFirSessionsCache()
+    private val libraryAsUseSiteSessionCache = LLFirSessionsCache()
 
     private val librariesSessionFactory = LLFirLibrarySessionFactory.getInstance(project)
     private val builtInsSessionFactory = LLFirBuiltinsSessionFactory.getInstance(project)
@@ -31,46 +34,64 @@ class LLFirSessionProviderStorage(val project: Project) {
     fun getSessionProvider(
         useSiteKtModule: KtModule,
         configureSession: (LLFirSession.() -> Unit)? = null
-    ): LLFirSessionProvider {
+    ): LLFirSessionProvider = executeWithoutPCE {
-        val (sessions, session) = sessionsCache.withMappings(project) { mappings ->
+        when (useSiteKtModule) {
-            val sessions = mutableMapOf<KtModule, LLFirResolvableModuleSession>().apply { putAll(mappings) }
+            is KtSourceModule -> {
-            val session = executeWithoutPCE {
+                createSessionProviderForSourceSession(useSiteKtModule, configureSession)
-                when (useSiteKtModule) {
-                    is KtSourceModule -> {
-                        LLFirSessionFactory.createSourcesSession(
-                            project,
-                            useSiteKtModule,
-                            globalComponents,
-                            sessionsCache.sessionInvalidator,
-                            sessions,
-                            librariesSessionFactory,
-                            configureSession = configureSession,
-                        )
-                    }
-
-                    is KtLibraryModule, is KtLibrarySourceModule -> {
-                        LLFirSessionFactory.createLibraryOrLibrarySourceResolvableSession(
-                            project,
-                            useSiteKtModule,
-                            globalComponents,
-                            sessionsCache.sessionInvalidator,
-                            builtInsSessionFactory.getBuiltinsSession(useSiteKtModule.platform),
-                            sessions,
-                            configureSession = configureSession,
-                        )
-                    }
-
-                    is KtNotUnderContentRootModule ->
-                        LLFirNonUnderContentRootSessionFactory.getInstance(project)
-                            .getNonUnderContentRootSession(useSiteKtModule, sessions)
-
-                    else -> error("Unexpected ${useSiteKtModule::class.simpleName}")
-                }
-
             }
+
+            is KtLibraryModule, is KtLibrarySourceModule -> {
+                createSessionProviderForLibraryOrLibrarySource(useSiteKtModule, configureSession)
+            }
+
+            is KtNotUnderContentRootModule -> {
+                val session = LLFirNonUnderContentRootSessionFactory.getInstance(project)
+                    .getNonUnderContentRootSession(useSiteKtModule)
+                LLFirSessionProvider(project, session, mapOf(useSiteKtModule to session))
+            }
+
+            else -> error("Unexpected ${useSiteKtModule::class.simpleName}")
+        }
+    }
+
+    private fun createSessionProviderForSourceSession(
+        useSiteKtModule: KtSourceModule,
+        configureSession: (LLFirSession.() -> Unit)?
+    ): LLFirSessionProvider {
+        val (sessions, session) = sourceAsUseSiteSessionCache.withMappings(project) { mappings ->
+            val sessions = mutableMapOf<KtModule, LLFirSession>().apply { putAll(mappings) }
+            val session = LLFirSessionFactory.createSourcesSession(
+                project,
+                useSiteKtModule,
+                globalComponents,
+                libraryAsUseSiteSessionCache.sessionInvalidator,
+                sessions,
+                librariesSessionFactory,
+                configureSession = configureSession,
+            )
             sessions to session
         }
+        return LLFirSessionProvider(project, session, sessions)
+    }
 
+
+    private fun createSessionProviderForLibraryOrLibrarySource(
+        useSiteKtModule: KtModule,
+        configureSession: (LLFirSession.() -> Unit)?
+    ): LLFirSessionProvider {
+        val (sessions, session) = libraryAsUseSiteSessionCache.withMappings(project) { mappings ->
+            val sessions = mutableMapOf<KtModule, LLFirSession>().apply { putAll(mappings) }
+            val session = LLFirSessionFactory.createLibraryOrLibrarySourceResolvableSession(
+                project,
+                useSiteKtModule,
+                globalComponents,
+                libraryAsUseSiteSessionCache.sessionInvalidator,
+                builtInsSessionFactory.getBuiltinsSession(useSiteKtModule.platform),
+                sessions,
+                configureSession = configureSession,
+            )
+            sessions to session
+        }
         return LLFirSessionProvider(project, session, sessions)
     }
 }
@@ -86,20 +107,22 @@ private class LLFirSessionsCache {
 
     inline fun <R> withMappings(
         project: Project,
-        action: (Map<KtModule, LLFirResolvableModuleSession>) -> Pair<Map<KtModule, LLFirResolvableModuleSession>, R>
+        action: (Map<KtModule, LLFirSession>) -> Pair<Map<KtModule, LLFirSession>, R>
-    ): Pair<Map<KtModule, LLFirResolvableModuleSession>, R> {
+    ): Pair<Map<KtModule, LLFirSession>, R> {
         val (newMappings, result) = action(getSessions().mapValues { it.value })
         mappings = newMappings.mapValues { FirSessionWithModificationTracker(project, it.value) }.toPersistentMap()
         return newMappings to result
     }
 
     @OptIn(ExperimentalStdlibApi::class)
-    private fun getSessions(): Map<KtModule, LLFirResolvableModuleSession> = buildMap {
+    private fun getSessions(): Map<KtModule, LLFirSession> = buildMap {
         val sessions = mappings.values
         val wasSessionInvalidated = sessions.associateWithTo(hashMapOf()) { false }
 
         val reversedDependencies = sessions.reversedDependencies { session ->
-            session.firSession.ktModule.directRegularDependencies.mapNotNull { mappings[it] }
+            if (session.validityTracker.isValid) {
+                session.firSession.ktModule.directRegularDependencies.mapNotNull { mappings[it] }
+            } else emptyList()
         }
 
         fun markAsInvalidWithDfs(session: FirSessionWithModificationTracker) {
@@ -136,13 +159,31 @@ private class LLFirSessionsCache {
 
 private class FirSessionWithModificationTracker(
     project: Project,
-    val firSession: LLFirResolvableModuleSession,
+    val firSession: LLFirSession,
 ) {
-    private val modificationTracker =
+
-        when (val ktModule = firSession.llFirModuleData.ktModule) {
+    val validityTracker: KtModuleStateTracker
-            is KtSourceModule -> ktModule.createModuleWithoutDependenciesOutOfBlockModificationTracker(project)
+    private val modificationTracker: ModificationTracker
-            else -> ModificationTracker.NEVER_CHANGED
+
+    init {
+        val trackerFactory = KotlinModificationTrackerFactory.getService(project)
+        val ktModule = firSession.llFirModuleData.ktModule
+
+        validityTracker = trackerFactory.createModuleStateTracker(ktModule)
+
+        val outOfBlockTracker = when (ktModule) {
+            is KtSourceModule -> trackerFactory.createModuleWithoutDependenciesOutOfBlockModificationTracker(ktModule)
+            else -> null
         }
+        modificationTracker = CompositeModificationTracker.create(
+            listOfNotNull(
+                outOfBlockTracker,
+                object : ModificationTracker {
+                    override fun getModificationCount() = validityTracker.rootModificationCount
+                }
+            )
+        )
+    }
 
 
     private val timeStamp = modificationTracker.modificationCount
@@ -154,5 +195,8 @@ private class FirSessionWithModificationTracker(
         isInvalidated = true
     }
 
-    val isValid: Boolean get() = !isInvalidated && modificationTracker.modificationCount == timeStamp
+    val isValid: Boolean
+        get() = validityTracker.isValid
+                && !isInvalidated
+                && modificationTracker.modificationCount == timeStamp
 }