[LL FIR] KT-58257 Keep library (source) sessions on global source invalidation

- `LLFirSessionCache.sourceCache` can contain library and library
  sources sessions. However, global source modification events should
  not remove such sessions from the cache, because they belong to
  "stable" modules.
- The commit introduces the term "stable module" as a combined term for
  binary and library source modules. Library sources are not binaries,
  but nevertheless they cannot cause nor be affected by out-of-block
  modification.
- The term "source modules" as used by global "source" modification is
  slightly imprecise, as it does *not* include library sources, but for
  the sake of conciseness, I don't plan to change that.

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

@@ -28,8 +28,8 @@ public abstract class KotlinGlobalModificationService {
     public abstract fun publishGlobalModuleStateModification()
 
     /**
-     * Publishes an event of global out-of-block modification of all [KtModule]s, including binary modules. The event does not invalidate
+     * Publishes an event of global out-of-block modification of all [KtModule]s. The event does not invalidate module state like
-     * module state like [publishGlobalModuleStateModification], so some module structure-specific caches might persist.
+     * [publishGlobalModuleStateModification], so some module structure-specific caches might persist.
      */
     @TestOnly
     public abstract fun publishGlobalOutOfBlockModification()

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/topics/KotlinGlobalModuleStateModificationListener.kt

@@ -9,8 +9,7 @@ import org.jetbrains.kotlin.analysis.project.structure.KtModule
 
 public fun interface KotlinGlobalModuleStateModificationListener {
     /**
-     * [onModification] is invoked in a write action before or after global modification of the module state of all [KtModule]s, including
+     * [onModification] is invoked in a write action before or after global modification of the module state of all [KtModule]s.
-     * binary modules.
      *
      * This event is published after SDK removal and to invalidate caches during/between tests.
      */

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/topics/KotlinGlobalOutOfBlockModificationListener.kt

@@ -9,8 +9,7 @@ import org.jetbrains.kotlin.analysis.project.structure.KtModule
 
 public fun interface KotlinGlobalOutOfBlockModificationListener {
     /**
-     * [onModification] is invoked in a write action before or after global out-of-block modification of all [KtModule]s, including binary
+     * [onModification] is invoked in a write action before or after global out-of-block modification of all [KtModule]s.
-     * modules.
      *
      * This event is published to invalidate caches during/between tests.
      */

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/topics/KotlinGlobalSourceModuleStateModificationListener.kt

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.analysis.project.structure.KtModule
 public fun interface KotlinGlobalSourceModuleStateModificationListener {
     /**
      * [onModification] is invoked in a write action before or after global modification of the module state of all source [KtModule]s,
-     * excluding binary modules.
+     * excluding [stable][org.jetbrains.kotlin.analysis.project.structure.isStableModule] modules.
      *
      * This event is published to invalidate caches during/between tests.
      */

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/topics/KotlinGlobalSourceOutOfBlockModificationListener.kt

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.analysis.project.structure.KtModule
 public fun interface KotlinGlobalSourceOutOfBlockModificationListener {
     /**
      * [onModification] is invoked in a write action before or after global out-of-block modification of all source [KtModule]s, excluding
-     * binary modules.
+     * [stable][org.jetbrains.kotlin.analysis.project.structure.isStableModule] modules.
      *
      * This event is published on global PSI changes.
      */

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

@@ -91,22 +91,20 @@ class LLFirSessionCache(private val project: Project) {
     }
 
     /**
-     * Removes all sessions after global invalidation. If [includeBinarySessions] is `false`, only source sessions will be removed.
+     * Removes all sessions after global invalidation. If [includeStableModules] is `false`, sessions of stable modules will not be removed.
      *
      * [removeAllSessions] must be called in a write action.
      */
-    fun removeAllSessions(includeBinarySessions: Boolean) {
+    fun removeAllSessions(includeStableModules: Boolean) {
         ApplicationManager.getApplication().assertWriteAccessAllowed()
 
-        removeAllSessionsFrom(sourceCache)
+        if (includeStableModules) {
-        if (includeBinarySessions) removeAllSessionsFrom(binaryCache)
+            removeAllSessionsFrom(sourceCache)
-    }
+            removeAllSessionsFrom(binaryCache)
-
+        } else {
-    private fun removeAllSessionsFrom(storage: SessionStorage) {
+            // `binaryCache` can only contain binary and thus stable modules, so we only need to remove sessions from `sourceCache`.
-        // Because `removeAllSessionsFrom` is executed in a write action, the order of setting `isValid` and clearing `storage` is not
+            removeAllMatchingSessionsFrom(sourceCache) { !it.isStableModule }
-        // important.
+        }
-        storage.values.forEach { it.isValid = false }
-        storage.clear()
     }
 
     // Removing script sessions is only needed temporarily until KTIJ-25620 has been implemented.
@@ -118,11 +116,22 @@ class LLFirSessionCache(private val project: Project) {
     }
 
     private fun removeAllScriptSessionsFrom(storage: SessionStorage) {
+        removeAllMatchingSessionsFrom(storage) { it is KtScriptModule || it is KtScriptDependencyModule }
+    }
+
+    private fun removeAllSessionsFrom(storage: SessionStorage) {
+        // Because `removeAllSessionsFrom` is executed in a write action, the order of setting `isValid` and clearing `storage` is not
+        // important.
+        storage.values.forEach { it.isValid = false }
+        storage.clear()
+    }
+
+    private inline fun removeAllMatchingSessionsFrom(storage: SessionStorage, shouldBeRemoved: (KtModule) -> Boolean) {
         // `ConcurrentSoftValueHashMap` (the implementation used by `storage`) does not back its entry set but rather creates a copy, which
         // is in violation of the contract of `Map.entrySet`, and thus changes to the entry set are not reflected in `storage`. Because this
         // function is executed in a write action, we do not need the weak consistency guarantees made by `ConcurrentMap`'s iterator, so a
         // "collect and remove" approach also works.
-        val scriptEntries = storage.entries.filter { (module, _) -> module is KtScriptModule || module is KtScriptDependencyModule }
+        val scriptEntries = storage.entries.filter { (module, _) -> shouldBeRemoved(module) }
         for ((module, session) in scriptEntries) {
             session.isValid = false
             storage.remove(module)

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

@@ -40,19 +40,19 @@ class LLFirSessionInvalidationService(private val project: Project) : Disposable
         )
         busConnection.subscribe(
             KotlinTopics.GLOBAL_MODULE_STATE_MODIFICATION,
-            KotlinGlobalModuleStateModificationListener { invalidateAll(includeBinaryModules = true) }
+            KotlinGlobalModuleStateModificationListener { invalidateAll(includeStableModules = true) }
         )
         busConnection.subscribe(
             KotlinTopics.GLOBAL_OUT_OF_BLOCK_MODIFICATION,
-            KotlinGlobalOutOfBlockModificationListener { invalidateAll(includeBinaryModules = true) }
+            KotlinGlobalOutOfBlockModificationListener { invalidateAll(includeStableModules = true) }
         )
         busConnection.subscribe(
             KotlinTopics.GLOBAL_SOURCE_MODULE_STATE_MODIFICATION,
-            KotlinGlobalSourceModuleStateModificationListener { invalidateAll(includeBinaryModules = false) },
+            KotlinGlobalSourceModuleStateModificationListener { invalidateAll(includeStableModules = false) },
         )
         busConnection.subscribe(
             KotlinTopics.GLOBAL_SOURCE_OUT_OF_BLOCK_MODIFICATION,
-            KotlinGlobalSourceOutOfBlockModificationListener { invalidateAll(includeBinaryModules = false) },
+            KotlinGlobalSourceOutOfBlockModificationListener { invalidateAll(includeStableModules = false) },
         )
     }
 
@@ -84,10 +84,10 @@ class LLFirSessionInvalidationService(private val project: Project) : Disposable
         }
     }
 
-    private fun invalidateAll(includeBinaryModules: Boolean) {
+    private fun invalidateAll(includeStableModules: Boolean) {
         ApplicationManager.getApplication().assertWriteAccessAllowed()
 
-        LLFirSessionCache.getInstance(project).removeAllSessions(includeBinaryModules)
+        LLFirSessionCache.getInstance(project).removeAllSessions(includeStableModules)
     }
 
     override fun dispose() {

analysis/project-structure/src/org/jetbrains/kotlin/analysis/project/structure/ktModuleUtils.kt

@@ -64,3 +64,11 @@ public inline fun <reified M : KtModule> KtModule.allDirectDependenciesOfType():
  */
 public fun computeTransitiveDependsOnDependencies(directDependsOnDependencies: List<KtModule>): List<KtModule> =
     topologicalSort(directDependsOnDependencies) { this.directDependsOnDependencies }
+
+/**
+ * A stable [KtModule] does not contain conventionally *editable* sources and cannot cause or be affected by out-of-block modification.
+ *
+ * Stable modules may still be subject to modification (e.g. when a JAR is added to a library), but in such cases, module state modification
+ * events will be raised.
+ */
+public val KtModule.isStableModule: Boolean get() = this is KtBinaryModule || this is KtLibrarySourceModule