FIR IDE: Add a way to use test compiler plugin in FIR IDE tests

Compiler plugins can be enabled
by WITH_FIR_TEST_COMPILER_PLUGIN directive

analysis/analysis-api-fir/build.gradle.kts

@@ -37,6 +37,10 @@ dependencies {
     testApi(toolsJar())
     testApiJUnit5()
     testApi(project(":analysis:symbol-light-classes"))
+
+    // We use 'api' instead of 'implementation' because other modules might be using these jars indirectly
+    testApi(project(":plugins:fir-plugin-prototype"))
+    testApi(projectTests(":plugins:fir-plugin-prototype"))
 }
 
 sourceSets {
@@ -48,6 +52,9 @@ projectTest(jUnitMode = JUnitMode.JUnit5) {
     dependsOn(":dist")
     workingDir = rootDir
     useJUnitPlatform()
+
+    // PluginAnnotationsProvider needs this jar during tests
+    dependsOn(":plugins:fir-plugin-prototype:plugin-annotations:jar")
 }
 
 testsJar()

analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/FirFrontendApiTestConfiguratorService.kt

@@ -9,9 +9,9 @@ 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.fir.utils.configureOptionalTestCompilerPlugin
 import org.jetbrains.kotlin.analysis.api.impl.barebone.test.FrontendApiTestConfiguratorService
 import org.jetbrains.kotlin.analysis.api.impl.base.references.HLApiReferenceProviderService
-import org.jetbrains.kotlin.analysis.low.level.api.fir.compiler.based.registerTestServices
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.FirLowLevelFrontendApiTestConfiguratorService
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.config.CompilerConfiguration
@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 object FirFrontendApiTestConfiguratorService : FrontendApiTestConfiguratorService {
     override fun TestConfigurationBuilder.configureTest(disposable: Disposable) {
         with(FirLowLevelFrontendApiTestConfiguratorService) { configureTest(disposable) }
+        configureOptionalTestCompilerPlugin()
     }
 
     override fun processTestFiles(files: List<KtFile>): List<KtFile> {

analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/utils/compilerPluginConfiguration.kt

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2010-2022 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.api.fir.utils
+
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.config.AnalysisFlag
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.LanguageVersion
+import org.jetbrains.kotlin.fir.plugin.services.ExtensionRegistrarConfigurator
+import org.jetbrains.kotlin.fir.plugin.services.PluginAnnotationsProvider
+import org.jetbrains.kotlin.test.Constructor
+import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
+import org.jetbrains.kotlin.test.directives.model.Directive
+import org.jetbrains.kotlin.test.directives.model.DirectiveApplicability
+import org.jetbrains.kotlin.test.directives.model.RegisteredDirectives
+import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
+import org.jetbrains.kotlin.test.model.TestModule
+import org.jetbrains.kotlin.test.services.AbstractEnvironmentConfigurator
+
+private object CompilerPluginsDirectives : SimpleDirectivesContainer() {
+    val WITH_FIR_TEST_COMPILER_PLUGIN by directive(
+        description = "Configure test compiler plugin from :plugins:fir-plugin-prototype module",
+        applicability = DirectiveApplicability.Global
+    )
+}
+
+internal fun TestConfigurationBuilder.configureOptionalTestCompilerPlugin() {
+    useDirectives(CompilerPluginsDirectives)
+
+    useConfigurators(
+        ::PluginAnnotationsProvider.enabledByDirective(CompilerPluginsDirectives.WITH_FIR_TEST_COMPILER_PLUGIN),
+        ::ExtensionRegistrarConfigurator.enabledByDirective(CompilerPluginsDirectives.WITH_FIR_TEST_COMPILER_PLUGIN),
+    )
+}
+
+private fun Constructor<AbstractEnvironmentConfigurator>.enabledByDirective(
+    directive: Directive
+): Constructor<AbstractEnvironmentConfigurator> {
+    val originalConstructor = this
+    return { testServices -> EnabledByDirectiveConfiguratorDecorator(originalConstructor(testServices), directive) }
+}
+
+private class EnabledByDirectiveConfiguratorDecorator(
+    private val original: AbstractEnvironmentConfigurator,
+    private val directive: Directive
+) : AbstractEnvironmentConfigurator() {
+    override fun configureCompileConfigurationWithAdditionalConfigurationKeys(configuration: CompilerConfiguration, module: TestModule) {
+        if (directive !in module.directives) return
+
+        original.configureCompileConfigurationWithAdditionalConfigurationKeys(configuration, module)
+    }
+
+    override fun provideAdditionalAnalysisFlags(
+        directives: RegisteredDirectives,
+        languageVersion: LanguageVersion
+    ): Map<AnalysisFlag<*>, Any?> = original.provideAdditionalAnalysisFlags(directives, languageVersion)
+
+    override fun registerCompilerExtensions(project: Project, module: TestModule) {
+        if (directive !in module.directives) return
+
+        original.registerCompilerExtensions(project, module)
+    }
+}