From efabd556465e9558830b03b77167df7475a096ae Mon Sep 17 00:00:00 2001
From: Marco Pennekamp <marco.pennekamp@jetbrains.com>
Date: Fri, 30 Jun 2023 18:50:05 +0200
Subject: [PATCH] [LL FIR] KT-58257 Don't invalidate dependent sessions when
 root session is absent

- We don't need to invalidate dependent sessions when the root session
  isn't in the session cache, or if the session is already marked as
  invalid, because any dependent sessions that might exist won't be
  referencing entities from the root session at that point.
---
 .../level/api/fir/sessions/LLFirSession.kt    |  2 ++
 .../api/fir/sessions/LLFirSessionCache.kt     | 23 +++++++++----------
 .../LLFirSessionInvalidationService.kt        |  9 +++++++-
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt
index 576907737ce..4f661dd3a39 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSession.kt
@@ -68,6 +68,8 @@ abstract class LLFirSession(
     }
 
     private fun invalidateInWriteAction() {
+        if (!isValid) return
+
         LLFirSessionInvalidationService.getInstance(project).invalidate(ktModule)
     }
 
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt
index af992e0caae..4c6e466ade0 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt
@@ -71,24 +71,23 @@ class LLFirSessionCache(private val project: Project) {
     }
 
     /**
-     * Removes the session(s) associated with [module] after it has been invalidated.
+     * Removes the session(s) associated with [module] after it has been invalidated. Must be called in a write action.
      *
-     * [removeSession] must be called in a write action.
+     * @return `true` if any sessions were removed.
      */
-    fun removeSession(module: KtModule) {
+    fun removeSession(module: KtModule): Boolean {
         ApplicationManager.getApplication().assertWriteAccessAllowed()
 
-        removeSessionFrom(module, sourceCache)
-        if (module is KtBinaryModule) {
-            removeSessionFrom(module, binaryCache)
-        }
+        val didSourceSessionExist = removeSessionFrom(module, sourceCache)
+        val didBinarySessionExist = if (module is KtBinaryModule) removeSessionFrom(module, binaryCache) else false
+
+        return didSourceSessionExist || didBinarySessionExist
     }
 
-    private fun removeSessionFrom(module: KtModule, storage: SessionStorage) {
-        val session = storage.remove(module)
-        if (session != null) {
-            session.isValid = false
-        }
+    private fun removeSessionFrom(module: KtModule, storage: SessionStorage): Boolean {
+        val session = storage.remove(module) ?: return false
+        session.isValid = false
+        return true
     }
 
     /**
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionInvalidationService.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionInvalidationService.kt
index 19076c450d3..993d3eb7fa1 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionInvalidationService.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionInvalidationService.kt
@@ -72,7 +72,14 @@ class LLFirSessionInvalidationService(private val project: Project) : Disposable
         ApplicationManager.getApplication().assertWriteAccessAllowed()
 
         val sessionCache = LLFirSessionCache.getInstance(project)
-        sessionCache.removeSession(module)
+        val didSessionExist = sessionCache.removeSession(module)
+
+        // We don't have to invalidate dependent sessions if the root session does not exist in the cache. It is true that sessions can be
+        // created without their dependency sessions being created, as session dependencies are lazy. So some of the root session's
+        // dependents might exist. But if the root session does not exist, its dependent sessions won't contain any elements resolved by the
+        // root session, so they effectively don't depend on the root session at that moment and don't need to be invalidated.
+        if (!didSessionExist) return
+
         KotlinModuleDependentsProvider.getInstance(project).getTransitiveDependents(module).forEach(sessionCache::removeSession)
     }