From 935f5c1180225c6ec2e2e96792c9369fc5ccf4b3 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Tue, 28 Jun 2022 12:03:30 +0300
Subject: [PATCH] [Parcelize] Fix registration of IR extension for K2 mode

---
 .../kotlin/parcelize/ParcelizeComponentRegistrar.kt   | 11 ++++++++---
 .../test/services/ParcelizeEnvironmentConfigurator.kt |  3 ++-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/plugins/parcelize/parcelize-compiler/parcelize.cli/src/org/jetbrains/kotlin/parcelize/ParcelizeComponentRegistrar.kt b/plugins/parcelize/parcelize-compiler/parcelize.cli/src/org/jetbrains/kotlin/parcelize/ParcelizeComponentRegistrar.kt
index eb48399ee43..026c9b8e9cb 100644
--- a/plugins/parcelize/parcelize-compiler/parcelize.cli/src/org/jetbrains/kotlin/parcelize/ParcelizeComponentRegistrar.kt
+++ b/plugins/parcelize/parcelize-compiler/parcelize.cli/src/org/jetbrains/kotlin/parcelize/ParcelizeComponentRegistrar.kt
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.codegen.extensions.ClassBuilderInterceptorExtension
 import org.jetbrains.kotlin.codegen.extensions.ExpressionCodegenExtension
 import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
+import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.container.StorageComponentContainer
 import org.jetbrains.kotlin.container.useInstance
@@ -22,9 +23,13 @@ import org.jetbrains.kotlin.resolve.extensions.SyntheticResolveExtension
 
 class ParcelizeComponentRegistrar : CompilerPluginRegistrar() {
     companion object {
-        fun registerParcelizeComponents(extensionStorage: ExtensionStorage) = with(extensionStorage) {
+        fun registerParcelizeComponents(extensionStorage: ExtensionStorage, useFir: Boolean) = with(extensionStorage) {
             ExpressionCodegenExtension.registerExtension(ParcelizeCodegenExtension())
-            IrGenerationExtension.registerExtension(ParcelizeIrGeneratorExtension())
+            if (useFir) {
+                IrGenerationExtension.registerExtension(ParcelizeFirIrGeneratorExtension())
+            } else {
+                IrGenerationExtension.registerExtension(ParcelizeIrGeneratorExtension())
+            }
             SyntheticResolveExtension.registerExtension(ParcelizeResolveExtension())
             ClassBuilderInterceptorExtension.registerExtension(ParcelizeClinitClassBuilderInterceptorExtension())
             StorageComponentContainerContributor.registerExtension(ParcelizeDeclarationCheckerComponentContainerContributor())
@@ -33,7 +38,7 @@ class ParcelizeComponentRegistrar : CompilerPluginRegistrar() {
     }
 
     override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
-        registerParcelizeComponents(this)
+        registerParcelizeComponents(this, configuration.getBoolean(CommonConfigurationKeys.USE_FIR))
     }
 
     override val supportsK2: Boolean
diff --git a/plugins/parcelize/parcelize-compiler/tests/org/jetbrains/kotlin/parcelize/test/services/ParcelizeEnvironmentConfigurator.kt b/plugins/parcelize/parcelize-compiler/tests/org/jetbrains/kotlin/parcelize/test/services/ParcelizeEnvironmentConfigurator.kt
index 6f45c3f17e1..7ccd3f35ee7 100644
--- a/plugins/parcelize/parcelize-compiler/tests/org/jetbrains/kotlin/parcelize/test/services/ParcelizeEnvironmentConfigurator.kt
+++ b/plugins/parcelize/parcelize-compiler/tests/org/jetbrains/kotlin/parcelize/test/services/ParcelizeEnvironmentConfigurator.kt
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
 import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.parcelize.ParcelizeComponentRegistrar
+import org.jetbrains.kotlin.test.model.FrontendKinds
 import org.jetbrains.kotlin.test.model.TestModule
 import org.jetbrains.kotlin.test.services.EnvironmentConfigurator
 import org.jetbrains.kotlin.test.services.TestServices
@@ -28,6 +29,6 @@ class ParcelizeEnvironmentConfigurator(testServices: TestServices) : Environment
         module: TestModule,
         configuration: CompilerConfiguration
     ) {
-        ParcelizeComponentRegistrar.registerParcelizeComponents(this)
+        ParcelizeComponentRegistrar.registerParcelizeComponents(this, useFir = module.frontendKind == FrontendKinds.FIR)
     }
 }