[Analysis API] Rework in-block modifications for dandling file modules

Before the change, all code fragment sessions were invalidated on any
PSI change. It was a simple solution, though not very effective.
Fragments were invalidated not only on a physical module change, but
also on change in another, unrelated code fragment.

As code fragment sessions were reworked to support also ordinary '.kt'
files, the old logic started to have even less sense, because ordinary
dangling files do not have a context element (they only have a context
module).

Currently, code fragment sessions are invalidated on any change in a
physical (non-dangling) module, and are also invalidated on any change
in their contextual dangling module, if any. With some work, this can be
improved further, so code fragments will be invalidated only on changes
in modules they depend on.

Relevant tests can be found in the IntelliJ project
(DanglingFileModuleInvalidationTest).

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

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2010-2023 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.providers.topics
+
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+
+public fun interface KotlinCodeFragmentContextModificationListener {
+    /**
+     * [onModification] is invoked in a write action before or after a context change for code fragments depending on the [module].
+     *
+     * All code fragments depending on [module], both directly or transitively, should be considered modified when this event is received.
+     *
+     * @see KotlinTopics
+     */
+    public fun onModification(module: KtModule)
+}

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

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.analysis.providers.analysisMessageBus
  *  - [KotlinGlobalModuleStateModificationListener]
  *  - [KotlinGlobalSourceModuleStateModificationListener]
  *  - [KotlinGlobalSourceOutOfBlockModificationListener]
+ *  - [KotlinCodeFragmentContextModificationListener]
  *
  * Care needs to be taken with the lack of interplay between different types of topics: Publishing a global modification event, for example,
  * does not imply the corresponding module-level event. Similarly, publishing a module state modification event does not imply out-of-block
@@ -60,4 +61,7 @@ public object KotlinTopics {
 
     public val GLOBAL_SOURCE_OUT_OF_BLOCK_MODIFICATION: Topic<KotlinGlobalSourceOutOfBlockModificationListener> =
         Topic(KotlinGlobalSourceOutOfBlockModificationListener::class.java, Topic.BroadcastDirection.TO_CHILDREN, true)
+
+    public val CODE_FRAGMENT_CONTEXT_MODIFICATION: Topic<KotlinCodeFragmentContextModificationListener> =
+        Topic(KotlinCodeFragmentContextModificationListener::class.java, Topic.BroadcastDirection.TO_CHILDREN, true)
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/LLFirDeclarationModificationService.kt

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecific
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.providers.analysisMessageBus
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics.CODE_FRAGMENT_CONTEXT_MODIFICATION
 import org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics.MODULE_OUT_OF_BLOCK_MODIFICATION
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.declarations.FirCodeFragment
@@ -166,26 +167,34 @@ class LLFirDeclarationModificationService(val project: Project) : Disposable {
         }
     }
 
-    private fun calculateChangeType(element: PsiElement, modificationType: ModificationType): ChangeType = when {
-        // If PSI is not valid, well something bad happened; OOBM won't hurt
-        !element.isValid -> ChangeType.OutOfBlock
+    private fun calculateChangeType(element: PsiElement, modificationType: ModificationType): ChangeType {
+        if (!element.isValid) {
+            // If PSI is not valid, well something bad happened; OOBM won't hurt
+            return ChangeType.OutOfBlock
+        }
 
-        element is PsiWhiteSpace || element is PsiComment -> ChangeType.Invisible
+        if (element is PsiWhiteSpace || element is PsiComment) {
+            return ChangeType.Invisible
+        }
 
-        // TODO improve for Java KTIJ-21684
-        element.language !is KotlinLanguage -> ChangeType.OutOfBlock
+        if (element.language !is KotlinLanguage) {
+            // TODO improve for Java KTIJ-21684
+            return ChangeType.OutOfBlock
+        }
 
-        else -> {
-            val inBlockModificationOwner = nonLocalDeclarationForLocalChange(element)
-            if (
-                inBlockModificationOwner != null
-                && !element.isNewDirectChildOf(inBlockModificationOwner, modificationType)
-                && !modificationType.isContractRemoval()
-            ) {
-                ChangeType.InBlock(inBlockModificationOwner, project)
-            } else {
-                ChangeType.OutOfBlock
-            }
+        val inBlockModificationOwner = nonLocalDeclarationForLocalChange(element) ?: return ChangeType.OutOfBlock
+
+        if (inBlockModificationOwner is KtCodeFragment) {
+            // All code fragment content is local
+            return ChangeType.InBlock(inBlockModificationOwner, project)
+        }
+
+        val isOutOfBlockChange = element.isNewDirectChildOf(inBlockModificationOwner, modificationType)
+                || modificationType.isContractRemoval()
+
+        return when {
+            !isOutOfBlockChange -> ChangeType.InBlock(inBlockModificationOwner, project)
+            else -> ChangeType.OutOfBlock
         }
     }
 
@@ -244,6 +253,8 @@ class LLFirDeclarationModificationService(val project: Project) : Disposable {
             ?: return // we do not have a cache for this file
 
         fileStructure.invalidateElement(declaration)
+
+        project.analysisMessageBus.syncPublisher(CODE_FRAGMENT_CONTEXT_MODIFICATION).onModification(ktModule)
     }
 
     private fun outOfBlockModification(element: PsiElement) {

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

@@ -117,6 +117,23 @@ class LLFirSessionCache(private val project: Project) {
         removeAllDanglingFileSessions()
     }
 
+    fun removeContextualDanglingFileSessions(contextModule: KtModule) {
+        if (contextModule is KtDanglingFileModule) {
+            removeAllMatchingSessionsFrom(danglingFileSessionCache) { it is KtDanglingFileModule && hasContextModule(it, contextModule) }
+        } else {
+            // Only code fragments can have a dangling file context
+            removeAllMatchingSessionsFrom(danglingFileSessionCache) { it is KtDanglingFileModule && it.isCodeFragment }
+        }
+    }
+
+    private tailrec fun hasContextModule(module: KtDanglingFileModule, contextModule: KtModule): Boolean {
+        return when (val candidate = module.contextModule) {
+            contextModule -> true
+            is KtDanglingFileModule -> hasContextModule(candidate, contextModule)
+            else -> false
+        }
+    }
+
     fun removeAllDanglingFileSessions() {
         removeAllSessionsFrom(danglingFileSessionCache)
     }
@@ -161,7 +178,6 @@ class LLFirSessionCache(private val project: Project) {
             is KtScriptModule -> sessionFactory.createScriptSession(module)
             is KtDanglingFileModule -> {
                 //  Dangling file context must have an analyzable session, so we can properly compile code against it.
-                // 'KtDanglingFileModule' is always a leaf module, so there might not be a circular reference.
                 val contextSession = getSession(module.contextModule, preferBinary = false)
                 sessionFactory.createDanglingFileSession(module, contextSession)
             }

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

@@ -8,7 +8,6 @@ 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
@@ -53,8 +52,8 @@ class LLFirSessionInvalidationService(private val project: Project) : Disposable
             KotlinGlobalSourceOutOfBlockModificationListener { invalidateAll(includeLibraryModules = false) },
         )
         busConnection.subscribe(
-            PsiModificationTracker.TOPIC,
-            PsiModificationTracker.Listener { invalidateAllDanglingFiles() }
+            KotlinTopics.CODE_FRAGMENT_CONTEXT_MODIFICATION,
+            KotlinCodeFragmentContextModificationListener { module -> invalidateCodeFragments(module) }
         )
     }
 
@@ -85,7 +84,11 @@ class LLFirSessionInvalidationService(private val project: Project) : Disposable
             sessionCache.removeAllScriptSessions()
         }
 
-        sessionCache.removeAllDanglingFileSessions()
+        if (module is KtDanglingFileModule) {
+            sessionCache.removeContextualDanglingFileSessions(module)
+        } else {
+            sessionCache.removeAllDanglingFileSessions()
+        }
     }
 
     private fun invalidateAll(includeLibraryModules: Boolean) {
@@ -105,8 +108,8 @@ class LLFirSessionInvalidationService(private val project: Project) : Disposable
         LLFirSessionCache.getInstance(project).removeAllSessions(includeLibraryModules)
     }
 
-    private fun invalidateAllDanglingFiles() {
-        LLFirSessionCache.getInstance(project).removeAllDanglingFileSessions()
+    private fun invalidateCodeFragments(contextModule: KtModule) {
+        LLFirSessionCache.getInstance(project).removeContextualDanglingFileSessions(contextModule)
     }
 
     override fun dispose() {