[LL API] Invalidate code fragment sessions on PSI modifications

Code fragments depend not only on their context module, but also on
body of the context declaration.

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

@@ -35,6 +35,7 @@ class LLFirSessionCache(private val project: Project) {
 
     private val sourceCache: SessionStorage = CollectionFactory.createConcurrentSoftValueMap()
     private val binaryCache: SessionStorage = CollectionFactory.createConcurrentSoftValueMap()
+    private val codeFragmentSessionCache: SessionStorage = CollectionFactory.createConcurrentSoftValueMap()
 
     /**
      * Returns the existing session if found, or creates a new session and caches it.
@@ -47,7 +48,12 @@ class LLFirSessionCache(private val project: Project) {
             }
         }
 
-        return getCachedSession(module, sourceCache, ::createSession)
+        val targetCache = when (module) {
+            is KtCodeFragmentModule -> codeFragmentSessionCache
+            else -> sourceCache
+        }
+
+        return getCachedSession(module, targetCache, ::createSession)
     }
 
     /**
@@ -79,9 +85,10 @@ class LLFirSessionCache(private val project: Project) {
         ApplicationManager.getApplication().assertWriteAccessAllowed()
 
         val didSourceSessionExist = removeSessionFrom(module, sourceCache)
-        val didBinarySessionExist = if (module is KtBinaryModule) removeSessionFrom(module, binaryCache) else false
+        val didBinarySessionExist = module is KtBinaryModule && removeSessionFrom(module, binaryCache)
+        val didCodeFragmentSessionExist = module is KtCodeFragmentModule && removeSessionFrom(module, codeFragmentSessionCache)
 
-        return didSourceSessionExist || didBinarySessionExist
+        return didSourceSessionExist || didBinarySessionExist || didCodeFragmentSessionExist
     }
 
     private fun removeSessionFrom(module: KtModule, storage: SessionStorage): Boolean {
@@ -106,6 +113,12 @@ class LLFirSessionCache(private val project: Project) {
             // `binaryCache` can only contain library modules, so we only need to remove sessions from `sourceCache`.
             removeAllMatchingSessionsFrom(sourceCache) { it !is KtBinaryModule && it !is KtLibrarySourceModule }
         }
+
+        removeAllCodeFragmentSessions()
+    }
+
+    fun removeAllCodeFragmentSessions() {
+        removeAllSessionsFrom(codeFragmentSessionCache)
     }
 
     // Removing script sessions is only needed temporarily until KTIJ-25620 has been implemented.

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.sessions
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
+import com.intellij.psi.util.PsiModificationTracker
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.analysis.providers.KotlinAnchorModuleProvider
 import org.jetbrains.kotlin.analysis.providers.analysisMessageBus
@@ -51,6 +52,10 @@ class LLFirSessionInvalidationService(private val project: Project) : Disposable
             KotlinTopics.GLOBAL_SOURCE_OUT_OF_BLOCK_MODIFICATION,
             KotlinGlobalSourceOutOfBlockModificationListener { invalidateAll(includeLibraryModules = false) },
         )
+        busConnection.subscribe(
+            PsiModificationTracker.TOPIC,
+            PsiModificationTracker.Listener { invalidateAllCodeFragments() }
+        )
     }
 
     /**
@@ -79,6 +84,8 @@ class LLFirSessionInvalidationService(private val project: Project) : Disposable
         if (module is KtScriptModule || module is KtScriptDependencyModule || module is KtLibraryModule) {
             sessionCache.removeAllScriptSessions()
         }
+
+        sessionCache.removeAllCodeFragmentSessions()
     }
 
     private fun invalidateAll(includeLibraryModules: Boolean) {
@@ -98,6 +105,10 @@ class LLFirSessionInvalidationService(private val project: Project) : Disposable
         LLFirSessionCache.getInstance(project).removeAllSessions(includeLibraryModules)
     }
 
+    private fun invalidateAllCodeFragments() {
+        LLFirSessionCache.getInstance(project).removeAllCodeFragmentSessions()
+    }
+
     override fun dispose() {
     }
 

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

@@ -15,6 +15,9 @@ import com.intellij.openapi.project.Project
  * should keep this limitation in mind and handle it separately. For example, a global modification event should be published for builtins
  * and SDK changes.
  *
+ * An empty set is also returned for [KtCodeFragmentModule]s, as no other module can depend on a code fragment.
+ * Additionally, [KtCodeFragmentModule] are not returned from [getDirectDependents] for their context modules.
+ *
  * Implementations of this provider should ensure that results are provided in reasonable time, for example by caching results, as its
  * functions may be called frequently.
  */