[AA] Improve modification event test directives

- Rewrite `TestModule.publishModificationEventByDirective`: we can get
  the `KtModule` to publish an event for via the test module. The
  directive can now be made optional, which is necessary when we want to
  specify it in potentially multiple modules.
- Introduce the `WILDCARD_MODIFICATION_EVENT` directive, which allows
  specifying THAT some modification event should be raised for a module,
  but not WHICH one. This allows generating multiple tests which raise
  different modification events over the same test data.
- Add various convenience functions for publishing wildcard modification
  events. `publishWildcardModificationEventsByDirective` for the test
  module structure is quite opinionated, but takes a lot of work from
  test implementations.
- Because `ModificationEventDirectives` are part of the analysis test
  framework, we register them as a general Analysis API test directive.

^KT-56288

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/registerBaseTestServices.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.test.framework.base
 
 import com.intellij.openapi.Disposable
 import org.jetbrains.kotlin.analysis.test.framework.AnalysisApiTestDirectives
+import org.jetbrains.kotlin.analysis.test.framework.directives.ModificationEventDirectives
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.AnalysisApiKtModuleProvider
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.AnalysisApiKtModuleProviderImpl
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.AnalysisApiTestCodeFragmentDirectives
@@ -33,6 +34,5 @@ fun TestConfigurationBuilder.registerAnalysisApiBaseTestServices(
 
     useCustomCompilerConfigurationProvider(::AnalysisApiTestCompilerConfiguratorProvider)
     usePreAnalysisHandlers(::ProjectStructureInitialisationPreAnalysisHandler.bind(configurator))
-    useDirectives(AnalysisApiTestDirectives)
+    useDirectives(AnalysisApiTestDirectives, AnalysisApiTestCodeFragmentDirectives, ModificationEventDirectives)
-    useDirectives(AnalysisApiTestCodeFragmentDirectives)
 }

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/directives/ModificationEventDirectives.kt

@@ -10,14 +10,31 @@ import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.providers.analysisMessageBus
 import org.jetbrains.kotlin.analysis.providers.topics.KotlinModuleStateModificationKind
 import org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.test.framework.services.environmentManager
 import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
-import org.jetbrains.kotlin.test.directives.model.singleValue
 import org.jetbrains.kotlin.test.model.TestModule
+import org.jetbrains.kotlin.test.services.TestModuleStructure
+import org.jetbrains.kotlin.test.services.TestServices
 import org.jetbrains.kotlin.test.testFramework.runWriteAction
 
+/**
+ * Modification event directives allow publishing modification events in a specific module or globally.
+ *
+ * These directives are available by default in Analysis API tests and do not need to be registered separately. They are not applied
+ * automatically, however, and tests interested in publishing modification events must call one of the following publishing functions for
+ * the appropriate [TestModule]: [publishModificationEventByDirective] or [publishWildcardModificationEventByDirectiveIfPresent].
+ */
 object ModificationEventDirectives : SimpleDirectivesContainer() {
     val MODIFICATION_EVENT by enumDirective<ModificationEventKind>(
-        description = "The modification event to be raised by the test.",
+        description = "The modification event to be raised for the module or globally.",
+    )
+
+    val WILDCARD_MODIFICATION_EVENT by directive(
+        description = """
+            Specifies that a modification event should be raised for the module or globally, but the test has to specify the modification 
+            event kind itself. This allows generating multiple tests which raise different modification events over the same test data.
+        """.trimIndent()
     )
 }
 
@@ -29,29 +46,93 @@ enum class ModificationEventKind {
     GLOBAL_SOURCE_OUT_OF_BLOCK_MODIFICATION,
 }
 
+val ModificationEventKind.isModuleLevel: Boolean
+    get() = this == ModificationEventKind.MODULE_STATE_MODIFICATION || this == ModificationEventKind.MODULE_OUT_OF_BLOCK_MODIFICATION
+
+val ModificationEventKind.isGlobalLevel: Boolean
+    get() = !isModuleLevel
+
 /**
  * Publishes a modification event as defined in [KotlinTopics][org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics] based on the
- * [ModificationEventDirectives.MODIFICATION_EVENT] directive in a write action.
+ * [ModificationEventDirectives.MODIFICATION_EVENT] directive present in the test module in a write action.
  *
- * The function expects exactly one `MODIFICATION_EVENT` directive to be present in the registered directives.
+ * Module-level modification events will be published for the [TestModule]'s [KtModule].
  *
- * @param module The module for which module-level modification events should be published, or `null` if only global modification events
+ * The function expects exactly one `MODIFICATION_EVENT` directive to be present, unless [isOptional] is `true`.
- *  should be published. (Specifying a module-level modification event kind then raises an error.)
  */
-fun TestModule.publishModificationEventByDirective(project: Project, module: KtModule?) {
+fun TestModule.publishModificationEventByDirective(testServices: TestServices, isOptional: Boolean = false) {
-    val modificationEventKind = directives.singleValue(ModificationEventDirectives.MODIFICATION_EVENT)
+    val modificationEventKinds = directives[ModificationEventDirectives.MODIFICATION_EVENT]
-
+    val modificationEventKind = when (modificationEventKinds.size) {
-    publishModificationEventByKind(modificationEventKind, project, module)
+        0 -> {
+            if (isOptional) return
+            error("Expected a `${ModificationEventDirectives.MODIFICATION_EVENT.name}` to be present in the test module `$this`.")
+        }
+        1 -> modificationEventKinds.single()
+        else -> error("The test module `$this` must not specify multiple modification events.")
+    }
+    publishModificationEvent(modificationEventKind, getKtModule(testServices))
 }
 
-private fun publishModificationEventByKind(modificationEventKind: ModificationEventKind, project: Project, module: KtModule?) {
+/**
+ * If the given test module contains a [ModificationEventDirectives.WILDCARD_MODIFICATION_EVENT] directive, publishes a modification event
+ * as defined in [KotlinTopics][org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics] based on the given [modificationEventKind] in
+ * a write action.
+ *
+ * Module-level modification events will be published for the [TestModule]'s [KtModule].
+ */
+fun TestModule.publishWildcardModificationEventByDirectiveIfPresent(
+    modificationEventKind: ModificationEventKind,
+    testServices: TestServices,
+) {
+    if (ModificationEventDirectives.WILDCARD_MODIFICATION_EVENT !in directives) {
+        return
+    }
+    publishModificationEvent(modificationEventKind, getKtModule(testServices))
+}
+
+/**
+ * For each test module that contains a [ModificationEventDirectives.WILDCARD_MODIFICATION_EVENT] directive, publishes a modification event
+ * as defined in [KotlinTopics][org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics] based on the given [modificationEventKind] in
+ * a write action.
+ *
+ * Global-level modification events will only be published *once*, regardless of how many `WILDCARD_MODIFICATION_EVENT` directives the test
+ * modules contain, as long as at least one test module contains it (to support test cases which don't want to publish any modification
+ * events).
+ */
+fun TestModuleStructure.publishWildcardModificationEventsByDirective(
+    modificationEventKind: ModificationEventKind,
+    testServices: TestServices,
+) {
+    if (modificationEventKind.isModuleLevel) {
+        modules.forEach { testModule ->
+            testModule.publishWildcardModificationEventByDirectiveIfPresent(modificationEventKind, testServices)
+        }
+    } else {
+        if (!allDirectives.contains(ModificationEventDirectives.WILDCARD_MODIFICATION_EVENT)) {
+            return
+        }
+        publishGlobalModificationEvent(modificationEventKind, testServices.environmentManager.getProject())
+    }
+}
+
+fun publishModificationEvent(modificationEventKind: ModificationEventKind, ktModule: KtModule) {
+    publishModificationEventByKind(modificationEventKind, ktModule.project, ktModule)
+}
+
+fun publishGlobalModificationEvent(modificationEventKind: ModificationEventKind, project: Project) {
+    require(modificationEventKind.isGlobalLevel)
+
+    publishModificationEventByKind(modificationEventKind, project, ktModule = null)
+}
+
+private fun publishModificationEventByKind(modificationEventKind: ModificationEventKind, project: Project, ktModule: KtModule?) {
     runWriteAction {
         when (modificationEventKind) {
             ModificationEventKind.MODULE_STATE_MODIFICATION -> {
                 project.analysisMessageBus
                     .syncPublisher(KotlinTopics.MODULE_STATE_MODIFICATION)
                     .onModification(
-                        module ?: errorModuleRequired(),
+                        ktModule ?: errorModuleRequired(),
                         KotlinModuleStateModificationKind.UPDATE,
                     )
             }
@@ -59,7 +140,7 @@ private fun publishModificationEventByKind(modificationEventKind: ModificationEv
             ModificationEventKind.MODULE_OUT_OF_BLOCK_MODIFICATION -> {
                 project.analysisMessageBus
                     .syncPublisher(KotlinTopics.MODULE_OUT_OF_BLOCK_MODIFICATION)
-                    .onModification(module ?: errorModuleRequired())
+                    .onModification(ktModule ?: errorModuleRequired())
             }
 
             ModificationEventKind.GLOBAL_MODULE_STATE_MODIFICATION -> {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/resolve/extensions/AbstractResolveExtensionDisposalAfterModificationEventTest.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions
 
 import com.intellij.mock.MockApplication
 import com.intellij.mock.MockProject
-import com.intellij.openapi.Disposable
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtension
 import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionFile
@@ -17,13 +16,11 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.Analys
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedTest
-import org.jetbrains.kotlin.analysis.test.framework.directives.ModificationEventDirectives
 import org.jetbrains.kotlin.analysis.test.framework.directives.publishModificationEventByDirective
 import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfigurator
 import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestServiceRegistrar
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.model.TestModule
 import org.jetbrains.kotlin.test.services.TestServices
 import org.jetbrains.kotlin.test.services.assertions
@@ -46,7 +43,7 @@ abstract class AbstractResolveExtensionDisposalAfterModificationEventTest : Abst
             "The resolve extension should not be disposed before the modification event is published."
         }
 
-        mainModule.publishModificationEventByDirective(project, module)
+        mainModule.publishModificationEventByDirective(testServices)
 
         testServices.assertions.assertTrue(resolveExtension.isDisposed) {
             "The resolve extension should be disposed after the modification event has been published."
@@ -74,12 +71,6 @@ class KtResolveExtensionWithDisposalTrackerProvider() : KtResolveExtensionProvid
 }
 
 object ResolveExtensionDisposalTestConfigurator : AnalysisApiFirSourceTestConfigurator(analyseInDependentSession = false) {
-    override fun configureTest(builder: TestConfigurationBuilder, disposable: Disposable) {
-        super.configureTest(builder, disposable)
-
-        builder.useDirectives(ModificationEventDirectives)
-    }
-
     override val serviceRegistrars: List<AnalysisApiTestServiceRegistrar>
         get() = buildList {
             addAll(super.serviceRegistrars)