From b152600e35665ba96ac11549dd3632424822f332 Mon Sep 17 00:00:00 2001
From: Alexander Udalov <alexander.udalov@jetbrains.com>
Date: Thu, 16 Feb 2023 23:05:44 +0100
Subject: [PATCH] IR: add flag to IrGenerationExtension for executing in kapt
 mode

 #KT-56635 Fixed
---
 .../extensions/IrGenerationExtension.kt       | 13 +++++++++++-
 .../kotlin/backend/jvm/JvmIrCodegenFactory.kt | 21 ++++++++++++-------
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension.kt
index b9ab77d8554..4125203b33a 100644
--- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension.kt
+++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension.kt
@@ -19,10 +19,21 @@ interface IrGenerationExtension : IrDeserializer.IrLinkerExtension {
     fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext)
 
     fun getPlatformIntrinsicExtension(backendContext: BackendContext): IrIntrinsicExtension? = null
+
+    // Returns true if this extension should also be applied in the KAPT stub generation mode in Kotlin/JVM. This mode uses light analysis
+    // in the compiler frontend to produce an "API-only" class file which is then converted to a .java stub. Because of the light analysis,
+    // the resulting IR does not have function bodies and can contain references to error types. If this method returns true, the extension
+    // should be ready to handle such incomplete IR. Any modifications to the IR applied during the KAPT stub generation mode will only have
+    // effect on the .java stub, not the resulting .class files. Compilation to the .class files is done in a separate step after stub
+    // generation.
+    //
+    // K2 KAPT doesn't use stub generation, so this property has no effect on extensions applied when K2 is enabled.
+    @FirIncompatiblePluginAPI
+    val shouldAlsoBeAppliedInKaptStubGenerationMode: Boolean get() = false
 }
 
 /**
  * This interface for common IR is empty because intrinsics are done in a platform-specific way (because of inliner).
  * Currently, only JVM intrinsics are supported via JvmIrIntrinsicExtension interface.
  */
-interface IrIntrinsicExtension
\ No newline at end of file
+interface IrIntrinsicExtension
diff --git a/compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmIrCodegenFactory.kt b/compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmIrCodegenFactory.kt
index 9b210065e16..d39765dbcd8 100644
--- a/compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmIrCodegenFactory.kt
+++ b/compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmIrCodegenFactory.kt
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.backend.jvm
 
 import org.jetbrains.kotlin.analyzer.hasJdkCapability
+import org.jetbrains.kotlin.backend.common.extensions.FirIncompatiblePluginAPI
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContextImpl
@@ -209,15 +210,19 @@ open class JvmIrCodegenFactory(
             irLinker,
             messageLogger
         )
-        if (pluginExtensions.isNotEmpty() && psi2irContext.configuration.generateBodies) {
+        if (pluginExtensions.isNotEmpty()) {
             for (extension in pluginExtensions) {
-                psi2ir.addPostprocessingStep { module ->
-                    val old = stubGenerator.unboundSymbolGeneration
-                    try {
-                        stubGenerator.unboundSymbolGeneration = true
-                        extension.generate(module, pluginContext)
-                    } finally {
-                        stubGenerator.unboundSymbolGeneration = old
+                if (psi2irContext.configuration.generateBodies ||
+                    @OptIn(FirIncompatiblePluginAPI::class) extension.shouldAlsoBeAppliedInKaptStubGenerationMode
+                ) {
+                    psi2ir.addPostprocessingStep { module ->
+                        val old = stubGenerator.unboundSymbolGeneration
+                        try {
+                            stubGenerator.unboundSymbolGeneration = true
+                            extension.generate(module, pluginContext)
+                        } finally {
+                            stubGenerator.unboundSymbolGeneration = old
+                        }
                     }
                 }
             }