[Analysis API] Support compiler plugins in code fragments

Compiler plugins are currently only supported in source modules.
This commit makes them also available in code fragments with source
context.

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/compilerFacility/FirIdeNormalAnalysisSourceModuleMultiModuleCompilerFacilityTestGenerated.java

@@ -63,4 +63,20 @@ public class FirIdeNormalAnalysisSourceModuleMultiModuleCompilerFacilityTestGene
     public void testInternalUsage() throws Exception {
         runTest("analysis/analysis-api/testData/components/compilerFacility/compilationMultiModule/internalUsage.kt");
     }
+
+    @Nested
+    @TestMetadata("analysis/analysis-api/testData/components/compilerFacility/compilationMultiModule/codeFragments")
+    @TestDataPath("$PROJECT_ROOT")
+    public class CodeFragments {
+        @Test
+        public void testAllFilesPresentInCodeFragments() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/compilerFacility/compilationMultiModule/codeFragments"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @Test
+        @TestMetadata("compilerPluginUsage.kt")
+        public void testCompilerPluginUsage() throws Exception {
+            runTest("analysis/analysis-api/testData/components/compilerFacility/compilationMultiModule/codeFragments/compilerPluginUsage.kt");
+        }
+    }
 }

analysis/analysis-api/testData/components/compilerFacility/compilationMultiModule/codeFragments/compilerPluginUsage.fragment.kt

@@ -0,0 +1,8 @@
+@org.jetbrains.kotlin.fir.plugin.AllOpen
+class Foo {
+    fun call() {}
+}
+
+class Bar : Foo() {}
+
+Bar().call()

analysis/analysis-api/testData/components/compilerFacility/compilationMultiModule/codeFragments/compilerPluginUsage.ir.txt

@@ -0,0 +1,59 @@
+MODULE_FRAGMENT
+  FILE fqName:<root> fileName:fragment.kt
+    CLASS CLASS name:CodeFragment modality:FINAL visibility:public superTypes:[kotlin.Any]
+      $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.CodeFragment
+      CONSTRUCTOR visibility:public <> () returnType:<root>.CodeFragment [primary]
+        BLOCK_BODY
+          DELEGATING_CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in kotlin.Any'
+      FUN name:run visibility:public modality:FINAL <> () returnType:kotlin.Unit
+        EXPRESSION_BODY
+          BLOCK type=kotlin.Unit origin=null
+            CLASS CLASS name:Foo modality:OPEN visibility:local superTypes:[kotlin.Any]
+              annotations:
+                AllOpen
+              $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.CodeFragment.run.Foo
+              CONSTRUCTOR visibility:public <> () returnType:<root>.CodeFragment.run.Foo [primary]
+                BLOCK_BODY
+                  DELEGATING_CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in kotlin.Any'
+                  INSTANCE_INITIALIZER_CALL classDescriptor='CLASS CLASS name:Foo modality:OPEN visibility:local superTypes:[kotlin.Any]'
+              FUN FAKE_OVERRIDE name:equals visibility:public modality:OPEN <> ($this:kotlin.Any, other:kotlin.Any?) returnType:kotlin.Boolean [fake_override,operator]
+                overridden:
+                  public open fun equals (other: kotlin.Any?): kotlin.Boolean [operator] declared in kotlin.Any
+                $this: VALUE_PARAMETER name:<this> type:kotlin.Any
+                VALUE_PARAMETER name:other index:0 type:kotlin.Any?
+              FUN FAKE_OVERRIDE name:hashCode visibility:public modality:OPEN <> ($this:kotlin.Any) returnType:kotlin.Int [fake_override]
+                overridden:
+                  public open fun hashCode (): kotlin.Int declared in kotlin.Any
+                $this: VALUE_PARAMETER name:<this> type:kotlin.Any
+              FUN FAKE_OVERRIDE name:toString visibility:public modality:OPEN <> ($this:kotlin.Any) returnType:kotlin.String [fake_override]
+                overridden:
+                  public open fun toString (): kotlin.String declared in kotlin.Any
+                $this: VALUE_PARAMETER name:<this> type:kotlin.Any
+              FUN name:call visibility:public modality:OPEN <> ($this:<root>.CodeFragment.run.Foo) returnType:kotlin.Unit
+                $this: VALUE_PARAMETER name:<this> type:<root>.CodeFragment.run.Foo
+                BLOCK_BODY
+            CLASS CLASS name:Bar modality:FINAL visibility:local superTypes:[<root>.CodeFragment.run.Foo]
+              $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.CodeFragment.run.Bar
+              CONSTRUCTOR visibility:public <> () returnType:<root>.CodeFragment.run.Bar [primary]
+                BLOCK_BODY
+                  DELEGATING_CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in <root>.CodeFragment.run.Foo'
+                  INSTANCE_INITIALIZER_CALL classDescriptor='CLASS CLASS name:Bar modality:FINAL visibility:local superTypes:[<root>.CodeFragment.run.Foo]'
+              FUN FAKE_OVERRIDE name:call visibility:public modality:OPEN <> ($this:<root>.CodeFragment.run.Foo) returnType:kotlin.Unit [fake_override]
+                overridden:
+                  public open fun call (): kotlin.Unit declared in <root>.CodeFragment.run.Foo
+                $this: VALUE_PARAMETER name:<this> type:<root>.CodeFragment.run.Foo
+              FUN FAKE_OVERRIDE name:equals visibility:public modality:OPEN <> ($this:kotlin.Any, other:kotlin.Any?) returnType:kotlin.Boolean [fake_override,operator]
+                overridden:
+                  public open fun equals (other: kotlin.Any?): kotlin.Boolean [fake_override,operator] declared in <root>.CodeFragment.run.Foo
+                $this: VALUE_PARAMETER name:<this> type:kotlin.Any
+                VALUE_PARAMETER name:other index:0 type:kotlin.Any?
+              FUN FAKE_OVERRIDE name:hashCode visibility:public modality:OPEN <> ($this:kotlin.Any) returnType:kotlin.Int [fake_override]
+                overridden:
+                  public open fun hashCode (): kotlin.Int [fake_override] declared in <root>.CodeFragment.run.Foo
+                $this: VALUE_PARAMETER name:<this> type:kotlin.Any
+              FUN FAKE_OVERRIDE name:toString visibility:public modality:OPEN <> ($this:kotlin.Any) returnType:kotlin.String [fake_override]
+                overridden:
+                  public open fun toString (): kotlin.String [fake_override] declared in <root>.CodeFragment.run.Foo
+                $this: VALUE_PARAMETER name:<this> type:kotlin.Any
+            CALL 'public open fun call (): kotlin.Unit [fake_override] declared in <root>.CodeFragment.run.Bar' type=kotlin.Unit origin=null
+              $this: CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in <root>.CodeFragment.run.Bar' type=<root>.CodeFragment.run.Bar origin=null

analysis/analysis-api/testData/components/compilerFacility/compilationMultiModule/codeFragments/compilerPluginUsage.kt

@@ -0,0 +1,5 @@
+// WITH_FIR_TEST_COMPILER_PLUGIN
+
+fun test() {
+    <caret>val x = 0
+}

analysis/analysis-api/testData/components/compilerFacility/compilationMultiModule/codeFragments/compilerPluginUsage.txt

@@ -0,0 +1,23 @@
+public final class CodeFragment$run$Bar {
+    // source: 'fragment.kt'
+    enclosing method CodeFragment.run()V
+    inner (local) class CodeFragment$run$Bar Bar
+    inner (local) class CodeFragment$run$Foo Foo
+    public method <init>(): void
+}
+
+public class CodeFragment$run$Foo {
+    // source: 'fragment.kt'
+    enclosing method CodeFragment.run()V
+    inner (local) class CodeFragment$run$Foo Foo
+    public method <init>(): void
+    public method call(): void
+}
+
+public final class CodeFragment {
+    // source: 'fragment.kt'
+    inner (local) class CodeFragment$run$Bar Bar
+    inner (local) class CodeFragment$run$Foo Foo
+    public method <init>(): void
+    public final static method run(): void
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/sessionFactoryHelpers.kt

@@ -17,13 +17,11 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirIdeRegiste
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirNonEmptyResolveExtensionTool
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.extensions.LLFirResolveExtensionTool
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSourcesSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirExceptionHandler
 import org.jetbrains.kotlin.analysis.project.structure.KtCompilerPluginsProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.providers.KotlinResolutionScopeProvider
 import org.jetbrains.kotlin.analysis.providers.createAnnotationResolver
-import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
 import org.jetbrains.kotlin.fir.FirExceptionHandler
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.SessionConfiguration
@@ -75,7 +73,7 @@ internal fun FirSession.registerCompilerPluginExtensions(project: Project, modul
 }
 
 @SessionConfiguration
-internal fun LLFirSourcesSession.registerCompilerPluginServices(
+internal fun LLFirSession.registerCompilerPluginServices(
     project: Project,
     module: KtSourceModule
 ) {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdePredicateBasedProvider.kt

@@ -9,6 +9,7 @@ import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSourcesSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
@@ -40,7 +41,7 @@ import org.jetbrains.kotlin.psi.*
  * PSI index based implementation of [FirPredicateBasedProvider].
  */
 internal class LLFirIdePredicateBasedProvider(
-    private val session: LLFirSourcesSession,
+    private val session: LLFirSession,
     private val annotationsResolver: KotlinAnnotationsResolver,
 ) : FirPredicateBasedProvider() {
     private val projectStructureProvider by lazy { ProjectStructureProvider.getInstance(session.project) }

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

@@ -499,6 +499,15 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
             register(FirProvider::class, firProvider)
             register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver())
 
+            val contextModule = module.contextModule
+            if (contextModule is KtSourceModule) {
+                registerCompilerPluginServices(project, contextModule)
+                registerCompilerPluginExtensions(project, contextModule)
+            } else {
+                register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotationsImpl(this))
+                register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)
+            }
+
             registerCommonComponentsAfterExtensionsAreConfigured()
 
             val dependencyProvider = LLFirDependenciesSymbolProvider(this) {
@@ -519,20 +528,19 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
 
             extensionService.additionalCheckers.forEach(session.checkersComponent::register)
 
-            //TODO: remove, replace with plugin loading
-            run {
-                register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)
-                register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotationsImpl(this))
-            }
-
             val syntheticFunctionInterfaceProvider = FirExtensionSyntheticFunctionInterfaceProvider
                 .createIfNeeded(this, moduleData, scopeProvider)
 
+            val switchableExtensionDeclarationsSymbolProvider = FirSwitchableExtensionDeclarationsSymbolProvider
+                .createIfNeeded(this)
+                ?.also { register(FirSwitchableExtensionDeclarationsSymbolProvider::class, it) }
+
             val context = CodeFragmentSessionCreationContext(
                 moduleData,
                 firProvider,
                 dependencyProvider,
-                syntheticFunctionInterfaceProvider
+                syntheticFunctionInterfaceProvider,
+                switchableExtensionDeclarationsSymbolProvider
             )
 
             additionalSessionConfiguration(context, this)
@@ -543,7 +551,8 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
         val moduleData: LLFirModuleData,
         val firProvider: LLFirProvider,
         val dependencyProvider: LLFirDependenciesSymbolProvider,
-        val syntheticFunctionInterfaceProvider: FirExtensionSyntheticFunctionInterfaceProvider?
+        val syntheticFunctionInterfaceProvider: FirExtensionSyntheticFunctionInterfaceProvider?,
+        val switchableExtensionDeclarationsSymbolProvider: FirSwitchableExtensionDeclarationsSymbolProvider?,
     )
 
     private fun wrapLanguageVersionSettings(original: LanguageVersionSettings): LanguageVersionSettings {

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

@@ -86,6 +86,7 @@ internal class LLFirJsSessionFactory(project: Project) : LLFirAbstractSessionFac
                     this,
                     providers = listOfNotNull(
                         firProvider.symbolProvider,
+                        switchableExtensionDeclarationsSymbolProvider,
                         syntheticFunctionInterfaceProvider,
                     ),
                     dependencyProvider,

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

@@ -91,6 +91,7 @@ internal class LLFirJvmSessionFactory(project: Project) : LLFirAbstractSessionFa
                     this,
                     providers = listOfNotNull(
                         firProvider.symbolProvider,
+                        switchableExtensionDeclarationsSymbolProvider,
                         syntheticFunctionInterfaceProvider,
                         javaSymbolProvider
                     ),

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

@@ -82,6 +82,7 @@ internal class LLFirNativeSessionFactory(project: Project) : LLFirAbstractSessio
                     this,
                     providers = listOfNotNull(
                         firProvider.symbolProvider,
+                        switchableExtensionDeclarationsSymbolProvider,
                         syntheticFunctionInterfaceProvider,
                     ),
                     dependencyProvider,