[FIR] Run irGenerationExtensions after IR actualization

^KT-56173 Fixed

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/pipeline/convertToIr.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.pipeline
 import org.jetbrains.kotlin.backend.common.actualizer.IrActualizedResult
 import org.jetbrains.kotlin.backend.common.actualizer.IrActualizer
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.backend.jvm.JvmIrTypeSystemContext
 import org.jetbrains.kotlin.backend.jvm.serialization.JvmIdSignatureDescriptor
 import org.jetbrains.kotlin.builtins.DefaultBuiltIns
@@ -96,7 +97,6 @@ fun FirResult.convertToIrAndActualize(
             fir2IrResult = outputs.single().convertToIr(
                 fir2IrExtensions,
                 fir2IrConfiguration,
-                irGeneratorExtensions,
                 commonMemberStorage = commonMemberStorage,
                 irBuiltIns = null,
                 irMangler,
@@ -115,7 +115,6 @@ fun FirResult.convertToIrAndActualize(
                 it.convertToIr(
                     fir2IrExtensions,
                     fir2IrConfiguration,
-                    irGeneratorExtensions,
                     commonMemberStorage = commonMemberStorage,
                     irBuiltIns = irBuiltIns,
                     irMangler,
@@ -131,7 +130,6 @@ fun FirResult.convertToIrAndActualize(
             fir2IrResult = platformOutput.convertToIr(
                 fir2IrExtensions,
                 fir2IrConfiguration,
-                irGeneratorExtensions,
                 commonMemberStorage = commonMemberStorage,
                 irBuiltIns = irBuiltIns!!,
                 irMangler,
@@ -152,13 +150,25 @@ fun FirResult.convertToIrAndActualize(
     }
 
     val (irModuleFragment, components, pluginContext) = fir2IrResult
+    components.applyIrGenerationExtensions(irModuleFragment, irGeneratorExtensions)
     return Fir2IrActualizedResult(irModuleFragment, components, pluginContext, actualizationResult)
 }
 
+fun Fir2IrComponents.applyIrGenerationExtensions(irModuleFragment: IrModuleFragment, irGenerationExtensions: Collection<IrGenerationExtension>) {
+    if (irGenerationExtensions.isEmpty()) return
+    Fir2IrPluginContext(this, irModuleFragment.descriptor).applyIrGenerationExtensions(irModuleFragment, irGenerationExtensions)
+}
+
+fun IrPluginContext.applyIrGenerationExtensions(irModuleFragment: IrModuleFragment, irGenerationExtensions: Collection<IrGenerationExtension>) {
+    if (irGenerationExtensions.isEmpty()) return
+    for (extension in irGenerationExtensions) {
+        extension.generate(irModuleFragment, this)
+    }
+}
+
 private fun ModuleCompilerAnalyzedOutput.convertToIr(
     fir2IrExtensions: Fir2IrExtensions,
     fir2IrConfiguration: Fir2IrConfiguration,
-    irGeneratorExtensions: Collection<IrGenerationExtension>,
     commonMemberStorage: Fir2IrCommonMemberStorage,
     irBuiltIns: IrBuiltInsOverFir?,
     irMangler: KotlinMangler.IrMangler,
@@ -170,7 +180,6 @@ private fun ModuleCompilerAnalyzedOutput.convertToIr(
         fir2IrExtensions, fir2IrConfiguration,
         irMangler, IrFactoryImpl, visibilityConverter,
         Fir2IrJvmSpecialAnnotationSymbolProvider(), // TODO: replace with appropriate (probably empty) implementation for other backends.
-        irGeneratorExtensions,
         kotlinBuiltIns = kotlinBuiltIns,
         commonMemberStorage = commonMemberStorage,
         initializedIrBuiltIns = irBuiltIns

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt

@@ -53,7 +53,6 @@ class Fir2IrConverter(
     private fun runSourcesConversion(
         allFirFiles: List<FirFile>,
         irModuleFragment: IrModuleFragmentImpl,
-        irGenerationExtensions: Collection<IrGenerationExtension>,
         fir2irVisitor: Fir2IrVisitor,
         fir2IrExtensions: Fir2IrExtensions,
         runPreCacheBuiltinClasses: Boolean
@@ -100,13 +99,6 @@ class Fir2IrConverter(
 
         evaluateConstants(irModuleFragment, configuration)
 
-        if (irGenerationExtensions.isNotEmpty()) {
-            val pluginContext = Fir2IrPluginContext(components, irModuleFragment.descriptor)
-            for (extension in irGenerationExtensions) {
-                extension.generate(irModuleFragment, pluginContext)
-            }
-        }
-
         irModuleFragment.acceptVoid(ExternalPackageParentPatcher(components, fir2IrExtensions))
     }
 
@@ -461,7 +453,6 @@ class Fir2IrConverter(
             irFactory: IrFactory,
             visibilityConverter: Fir2IrVisibilityConverter,
             specialSymbolProvider: Fir2IrSpecialSymbolProvider,
-            irGenerationExtensions: Collection<IrGenerationExtension>,
             kotlinBuiltIns: KotlinBuiltIns,
             commonMemberStorage: Fir2IrCommonMemberStorage,
             initializedIrBuiltIns: IrBuiltInsOverFir?
@@ -508,7 +499,7 @@ class Fir2IrConverter(
             }
 
             converter.runSourcesConversion(
-                allFirFiles, irModuleFragment, irGenerationExtensions, fir2irVisitor, fir2IrExtensions,
+                allFirFiles, irModuleFragment, fir2irVisitor, fir2IrExtensions,
                 runPreCacheBuiltinClasses = initializedIrBuiltIns == null
             )
 

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/backend/ir/ActualizerOnlyFacade.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.backend.common.actualizer.IrActualizer
 import org.jetbrains.kotlin.backend.jvm.JvmIrTypeSystemContext
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.languageVersionSettings
+import org.jetbrains.kotlin.fir.pipeline.applyIrGenerationExtensions
 import org.jetbrains.kotlin.ir.types.IrTypeSystemContextImpl
 import org.jetbrains.kotlin.platform.jvm.isJvm
 import org.jetbrains.kotlin.test.model.*
@@ -32,6 +33,10 @@ class ActualizerOnlyFacade(
                 typeSystemContext,
                 testServices.compilerConfigurationProvider.getCompilerConfiguration(module).languageVersionSettings
             )
+            inputArtifact.irPluginContext.applyIrGenerationExtensions(
+                inputArtifact.irModuleFragment,
+                irGenerationExtensions = module.irGenerationExtensions(testServices)
+            )
         }
         return inputArtifact
     }

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/backend/ir/JvmIrBackendFacade.kt

@@ -8,14 +8,15 @@ package org.jetbrains.kotlin.test.backend.ir
 import org.jetbrains.kotlin.KtPsiSourceFile
 import org.jetbrains.kotlin.backend.common.BackendException
 import org.jetbrains.kotlin.backend.common.actualizer.IrActualizer
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.jvm.JvmIrTypeSystemContext
 import org.jetbrains.kotlin.backend.jvm.MultifileFacadeFileEntry
 import org.jetbrains.kotlin.backend.jvm.lower.getFileClassInfoFromIrFile
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.fileClasses.JvmFileClassUtil
+import org.jetbrains.kotlin.fir.pipeline.applyIrGenerationExtensions
 import org.jetbrains.kotlin.ir.PsiIrFileEntry
 import org.jetbrains.kotlin.ir.declarations.IrFile
-import org.jetbrains.kotlin.ir.types.IrTypeSystemContextImpl
 import org.jetbrains.kotlin.ir.util.NaiveSourceBasedFileEntryImpl
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.backend.classic.JavaCompilerFacade
@@ -45,6 +46,10 @@ class JvmIrBackendFacade(
                 JvmIrTypeSystemContext(inputArtifact.irModuleFragment.irBuiltins),
                 inputArtifact.state.languageVersionSettings
             )
+            inputArtifact.irPluginContext.applyIrGenerationExtensions(
+                inputArtifact.irModuleFragment,
+                module.irGenerationExtensions(testServices)
+            )
         }
 
         val state = inputArtifact.state
@@ -91,8 +96,12 @@ class JvmIrBackendFacade(
             }
         )
     }
-
-    private fun TestModule.useIrActualizer(): Boolean {
-        return frontendKind == FrontendKinds.FIR && languageVersionSettings.supportsFeature(LanguageFeature.MultiPlatformProjects)
-    }
+}
+
+internal fun TestModule.useIrActualizer(): Boolean {
+    return frontendKind == FrontendKinds.FIR && languageVersionSettings.supportsFeature(LanguageFeature.MultiPlatformProjects)
+}
+
+fun TestModule.irGenerationExtensions(testServices: TestServices): Collection<IrGenerationExtension> {
+    return IrGenerationExtension.getInstances(testServices.compilerConfigurationProvider.getProject(this))
 }

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/Fir2IrJsResultsConverter.kt

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.fir.backend.jvm.Fir2IrJvmSpecialAnnotationSymbolProv
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.descriptors.FirModuleDescriptor
+import org.jetbrains.kotlin.fir.pipeline.applyIrGenerationExtensions
 import org.jetbrains.kotlin.fir.serialization.FirKLibSerializerExtension
 import org.jetbrains.kotlin.fir.serialization.serializeSingleFirFile
 import org.jetbrains.kotlin.incremental.components.LookupTracker
@@ -41,6 +42,8 @@ import org.jetbrains.kotlin.library.metadata.resolver.KotlinResolvedLibrary
 import org.jetbrains.kotlin.library.unresolvedDependencies
 import org.jetbrains.kotlin.storage.LockBasedStorageManager
 import org.jetbrains.kotlin.test.backend.ir.IrBackendInput
+import org.jetbrains.kotlin.test.backend.ir.irGenerationExtensions
+import org.jetbrains.kotlin.test.backend.ir.useIrActualizer
 import org.jetbrains.kotlin.test.directives.CodegenTestDirectives
 import org.jetbrains.kotlin.test.model.BackendKinds
 import org.jetbrains.kotlin.test.model.Frontend2BackendConverter
@@ -120,7 +123,7 @@ class Fir2IrJsResultsConverter(
 
         var actualizedExpectDeclarations: Set<FirDeclaration>? = null
 
-        return IrBackendInput.JsIrBackendInput(
+        val result = IrBackendInput.JsIrBackendInput(
             mainIrPart,
             dependentIrParts,
             mainPluginContext,
@@ -152,6 +155,15 @@ class Fir2IrJsResultsConverter(
                 configuration.languageVersionSettings,
             )
         }
+
+        if (!module.useIrActualizer()) {
+            result.irPluginContext.applyIrGenerationExtensions(
+                result.irModuleFragment,
+                irGenerationExtensions = module.irGenerationExtensions(testServices)
+            )
+        }
+
+        return result
     }
 }
 
@@ -185,7 +197,6 @@ fun AbstractFirAnalyzerFacade.convertToJsIr(
         irMangler, IrFactoryImpl,
         Fir2IrVisibilityConverter.Default,
         Fir2IrJvmSpecialAnnotationSymbolProvider(), // TODO: replace with appropriate (probably empty) implementation
-        irGeneratorExtensions,
         kotlinBuiltIns = builtIns ?: DefaultBuiltIns.Instance, // TODO: consider passing externally,
         commonMemberStorage = commonMemberStorage,
         initializedIrBuiltIns = irBuiltIns

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/Fir2IrJvmResultsConverter.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.backend.jvm.FirJvmBackendClassResolver
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmBackendExtension
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmKotlinMangler
 import org.jetbrains.kotlin.fir.backend.jvm.JvmFir2IrExtensions
+import org.jetbrains.kotlin.fir.pipeline.applyIrGenerationExtensions
 import org.jetbrains.kotlin.fir.pipeline.signatureComposerForJvmFir2Ir
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.ir.backend.jvm.serialization.JvmIrMangler
@@ -29,6 +30,8 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.CompilerEnvironment
 import org.jetbrains.kotlin.resolve.lazy.declarations.FileBasedDeclarationProviderFactory
 import org.jetbrains.kotlin.test.backend.ir.IrBackendInput
+import org.jetbrains.kotlin.test.backend.ir.irGenerationExtensions
+import org.jetbrains.kotlin.test.backend.ir.useIrActualizer
 import org.jetbrains.kotlin.test.directives.CodegenTestDirectives
 import org.jetbrains.kotlin.test.model.BackendKinds
 import org.jetbrains.kotlin.test.model.Frontend2BackendConverter
@@ -126,7 +129,7 @@ class Fir2IrJvmResultsConverter(
             FirJvmBackendClassResolver(mainModuleComponents)
         ).build()
 
-        return IrBackendInput.JvmIrBackendInput(
+        val result = IrBackendInput.JvmIrBackendInput(
             generationState,
             codegenFactory,
             backendInput,
@@ -136,5 +139,14 @@ class Fir2IrJvmResultsConverter(
             irMangler = irMangler,
             firMangler = commonMemberStorage.firSignatureComposer.mangler,
         )
+
+        if (!module.useIrActualizer()) {
+            result.irPluginContext.applyIrGenerationExtensions(
+                backendInput.irModuleFragment,
+                irGenerationExtensions = module.irGenerationExtensions(testServices)
+            )
+        }
+
+        return result
     }
 }

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/Fir2IrWasmResultsConverter.kt

@@ -185,7 +185,6 @@ fun AbstractFirAnalyzerFacade.convertToWasmIr(
         irMangler, IrFactoryImpl,
         Fir2IrVisibilityConverter.Default,
         Fir2IrJvmSpecialAnnotationSymbolProvider(), // TODO: replace with appropriate (probably empty) implementation
-        irGeneratorExtensions,
         kotlinBuiltIns = builtIns ?: DefaultBuiltIns.Instance, // TODO: consider passing externally,
         commonMemberStorage = commonMemberStorage,
         initializedIrBuiltIns = irBuiltIns

compiler/tests-compiler-utils/tests/org/jetbrains/kotlin/fir/FirAnalyzerFacade.kt

@@ -119,7 +119,6 @@ class FirAnalyzerFacade(
             JvmIrMangler, IrFactoryImpl,
             FirJvmVisibilityConverter,
             Fir2IrJvmSpecialAnnotationSymbolProvider(),
-            irGeneratorExtensions,
             kotlinBuiltIns = DefaultBuiltIns.Instance, // TODO: consider passing externally,
             commonMemberStorage = commonMemberStorage,
             initializedIrBuiltIns = irBuiltIns

js/js.tests/test/org/jetbrains/kotlin/js/test/converters/FirJsKlibBackendFacade.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
+import org.jetbrains.kotlin.fir.pipeline.applyIrGenerationExtensions
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.ir.backend.js.JsFactories
 import org.jetbrains.kotlin.ir.backend.js.resolverLogger
@@ -24,6 +25,7 @@ import org.jetbrains.kotlin.library.KotlinAbiVersion
 import org.jetbrains.kotlin.storage.LockBasedStorageManager
 import org.jetbrains.kotlin.test.backend.ir.IrBackendFacade
 import org.jetbrains.kotlin.test.backend.ir.IrBackendInput
+import org.jetbrains.kotlin.test.backend.ir.irGenerationExtensions
 import org.jetbrains.kotlin.test.frontend.classic.ModuleDescriptorProvider
 import org.jetbrains.kotlin.test.frontend.classic.moduleDescriptorProvider
 import org.jetbrains.kotlin.test.frontend.fir.getAllJsDependenciesPaths
@@ -74,7 +76,12 @@ class FirJsKlibBackendFacade(
                         diagnosticReporter,
                         IrTypeSystemContextImpl(inputArtifact.irModuleFragment.irBuiltins),
                         configuration.languageVersionSettings
-                    )
+                    ).also {
+                        inputArtifact.irPluginContext.applyIrGenerationExtensions(
+                            inputArtifact.irModuleFragment,
+                            irGenerationExtensions = module.irGenerationExtensions(testServices)
+                        )
+                    }
                 } else {
                     null
                 }

plugins/fir-plugin-prototype/testData/box/replaceActualFunctionBodyWitExpectDefaultValue.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // LANGUAGE: +MultiPlatformProjects
 // ISSUE: KT-56173
 

wasm/wasm.tests/test/org/jetbrains/kotlin/wasm/test/converters/FirWasmKlibBackendFacade.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
+import org.jetbrains.kotlin.fir.pipeline.applyIrGenerationExtensions
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.ir.backend.js.JsFactories
 import org.jetbrains.kotlin.ir.backend.js.resolverLogger
@@ -22,6 +23,7 @@ import org.jetbrains.kotlin.library.KotlinAbiVersion
 import org.jetbrains.kotlin.storage.LockBasedStorageManager
 import org.jetbrains.kotlin.test.backend.ir.IrBackendFacade
 import org.jetbrains.kotlin.test.backend.ir.IrBackendInput
+import org.jetbrains.kotlin.test.backend.ir.irGenerationExtensions
 import org.jetbrains.kotlin.test.frontend.classic.ModuleDescriptorProvider
 import org.jetbrains.kotlin.test.frontend.classic.moduleDescriptorProvider
 import org.jetbrains.kotlin.test.frontend.fir.getAllWasmDependenciesPaths
@@ -71,7 +73,12 @@ class FirWasmKlibBackendFacade(
                         diagnosticReporter,
                         IrTypeSystemContextImpl(inputArtifact.irModuleFragment.irBuiltins),
                         configuration.languageVersionSettings
-                    )
+                    ).also {
+                        inputArtifact.irPluginContext.applyIrGenerationExtensions(
+                            inputArtifact.irModuleFragment,
+                            irGenerationExtensions = module.irGenerationExtensions(testServices)
+                        )
+                    }
                 } else {
                     null
                 }
@@ -120,4 +127,4 @@ class FirWasmKlibBackendFacade(
 
         return BinaryArtifacts.KLib(File(outputFile))
     }
-}
+}