FIR IDE: fix module invalidation algorithm

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeSessionProviderStorage.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirTransformerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.trackers.KotlinFirOutOfBlockModificationTrackerFactory
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.addValueFor
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
 import java.util.concurrent.ConcurrentHashMap
 
@@ -71,36 +72,42 @@ private class FromModuleViewSessionCache(
     @OptIn(ExperimentalStdlibApi::class)
     private fun getSessions(): Map<ModuleSourceInfo, FirIdeSourcesSession> = buildMap {
         val sessions = mappings.values
-        val sessionToValidity = hashMapOf<FirSessionWithModificationTracker, Boolean>()
+        val wasSessionInvalidated = sessions.associateWithTo(hashMapOf()) { false }
 
-        var isValid = true
+        val reversedDependencies = sessions.reversedDependencies { session ->
-        fun dfs(session: FirSessionWithModificationTracker) {
+            session.firSession.dependencies.mapNotNull { mappings[it] }
-            sessionToValidity[session]?.let { valid ->
+        }
-                if (!valid) isValid = false
+
+        fun markAsInvalidWithDfs(session: FirSessionWithModificationTracker) {
+            if (wasSessionInvalidated.getValue(session)) {
+                // we already was in that branch
                 return
             }
-            sessionToValidity[session] = session.isValid
+            wasSessionInvalidated[session] = true
-            session.firSession.dependencies.forEach { dependency ->
+            reversedDependencies[session]?.forEach { dependsOn ->
-                mappings[dependency]?.let(::dfs)
+                markAsInvalidWithDfs(dependsOn)
-            }
-            if (!session.isValid) {
-                isValid = false
-            }
-            if (!isValid) {
-                sessionToValidity[session] = false
             }
         }
 
         for (session in sessions) {
-            if (session !in sessionToValidity) {
+            if (!session.isValid) {
-                isValid = true
+                markAsInvalidWithDfs(session)
-                dfs(session)
             }
         }
-        return sessionToValidity.entries
+        return wasSessionInvalidated.entries
-            .mapNotNull { (session, valid) -> session.takeIf { valid } }
+            .mapNotNull { (session, wasInvalidated) -> session.takeUnless { wasInvalidated } }
             .associate { session -> session.firSession.moduleInfo to session.firSession }
     }
+
+    private fun <T> Collection<T>.reversedDependencies(getDependencies: (T) -> List<T>): Map<T, List<T>> {
+        val result = hashMapOf<T, MutableList<T>>()
+        forEach { from ->
+            getDependencies(from).forEach { to ->
+                result.addValueFor(to, from)
+            }
+        }
+        return result
+    }
 }
 
 private class FirSessionWithModificationTracker(