Add a key to enable spilling of all variables in a suspending context

This commit adds a new key that will allow users to enhance their
debugging experience in suspending contexts when using the IR backend.
After the key is enabled, the following things are changed:
1. All variables in a suspending context are spilled regardless their
liveness.
2. Their LVT records are not shrunk.
3. ACONST_NULL is not spilled to dead variables.

#KT-48678 In progress

(cherry picked from commit 38d97d0621)

compiler/backend/src/org/jetbrains/kotlin/codegen/coroutines/CoroutineTransformerMethodVisitor.kt

@@ -61,7 +61,8 @@ class CoroutineTransformerMethodVisitor(
     // JVM_IR backend generates $completion, while old backend does not
     private val putContinuationParameterToLvt: Boolean = true,
     // Parameters of suspend lambda are put to the same fields as spilled variables
-    private val initialVarsCountByType: Map<Type, Int> = emptyMap()
+    private val initialVarsCountByType: Map<Type, Int> = emptyMap(),
+    private val shouldOptimiseUnusedVariables: Boolean = true
 ) : TransformationMethodVisitor(delegate, access, name, desc, signature, exceptions) {
 
     private val classBuilderForCoroutineState: ClassBuilder by lazy(obtainClassBuilderForCoroutineState)
@@ -206,7 +207,9 @@ class CoroutineTransformerMethodVisitor(
         dropUnboxInlineClassMarkers(methodNode, suspensionPoints)
         methodNode.removeEmptyCatchBlocks()
 
-        updateLvtAccordingToLiveness(methodNode, isForNamedFunction, stateLabels)
+        if (shouldOptimiseUnusedVariables) {
+            updateLvtAccordingToLiveness(methodNode, isForNamedFunction, stateLabels)
+        }
 
         writeDebugMetadata(methodNode, suspensionPointLineNumbers, spilledToVariableMapping)
     }
@@ -667,7 +670,7 @@ class CoroutineTransformerMethodVisitor(
             for (slot in 0 until localsCount) {
                 if (slot == continuationIndex || slot == dataIndex) continue
                 val value = frame.getLocal(slot)
-                if (value.type == null || !livenessFrame.isAlive(slot)) continue
+                if (value.type == null || (shouldOptimiseUnusedVariables && !livenessFrame.isAlive(slot))) continue
 
                 if (value == StrictBasicValue.NULL_VALUE) {
                     referencesToSpill += slot to null
@@ -875,12 +878,16 @@ class CoroutineTransformerMethodVisitor(
             for ((slot, referenceToSpill) in referencesToSpillBySuspensionPointIndex[suspensionPointIndex]) {
                 generateSpillAndUnspill(suspension, slot, referenceToSpill)
             }
-            val (currentSpilledCount, predSpilledCount) = referencesToCleanBySuspensionPointIndex[suspensionPointIndex]
+
-            if (predSpilledCount > currentSpilledCount) {
+            if (shouldOptimiseUnusedVariables) {
-                for (fieldIndex in currentSpilledCount until predSpilledCount) {
+                val (currentSpilledCount, predSpilledCount) = referencesToCleanBySuspensionPointIndex[suspensionPointIndex]
-                    cleanUpField(suspension, fieldIndex)
+                if (predSpilledCount > currentSpilledCount) {
+                    for (fieldIndex in currentSpilledCount until predSpilledCount) {
+                        cleanUpField(suspension, fieldIndex)
+                    }
                 }
             }
+
             for ((slot, primitiveToSpill) in primitivesToSpillBySuspensionPointIndex[suspensionPointIndex]) {
                 generateSpillAndUnspill(suspension, slot, primitiveToSpill)
             }

compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JVMCompilerArguments.kt

@@ -511,6 +511,13 @@ Also sets `-jvm-target` value equal to the selected JDK version"""
     )
     var linkViaSignatures: Boolean by FreezableVar(false)
 
+    @Argument(
+        value = "-Xdebug",
+        description = "Enable debug mode for compilation.\n" +
+                "Currently this includes spilling all variables in a suspending context regardless their liveness."
+    )
+    var enableDebugMode: Boolean by FreezableVar(false)
+
     override fun configureAnalysisFlags(collector: MessageCollector, languageVersion: LanguageVersion): MutableMap<AnalysisFlag<*>, Any> {
         val result = super.configureAnalysisFlags(collector, languageVersion)
         result[JvmAnalysisFlags.strictMetadataVersionSemantics] = strictMetadataVersionSemantics

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/jvmArguments.kt

@@ -310,6 +310,8 @@ fun CompilerConfiguration.configureAdvancedJvmOptions(arguments: K2JVMCompilerAr
 
     put(JVMConfigurationKeys.LINK_VIA_SIGNATURES, arguments.linkViaSignatures)
 
+    put(JVMConfigurationKeys.ENABLE_DEBUG_MODE, arguments.enableDebugMode)
+
     val assertionsMode =
         JVMAssertionsMode.fromStringOrNull(arguments.assertionsMode)
     if (assertionsMode == null) {

compiler/config.jvm/src/org/jetbrains/kotlin/config/JVMConfigurationKeys.java

@@ -156,4 +156,7 @@ public class JVMConfigurationKeys {
 
     public static final CompilerConfigurationKey<Boolean> LINK_VIA_SIGNATURES =
             CompilerConfigurationKey.create("Link JVM IR symbols via signatures, instead of by descriptors");
+
+    public static final CompilerConfigurationKey<Boolean> ENABLE_DEBUG_MODE =
+            CompilerConfigurationKey.create("Enable debug mode");
 }

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.backend.jvm.unboxInlineClass
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.coroutines.CoroutineTransformerMethodVisitor
 import org.jetbrains.kotlin.codegen.coroutines.reportSuspensionPointInsideMonitor
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrBlockBody
@@ -36,13 +37,14 @@ internal fun MethodNode.acceptWithStateMachine(
     varsCountByType: Map<Type, Int>,
     obtainContinuationClassBuilder: () -> ClassBuilder,
 ) {
-    val state = classCodegen.context.state
+    val context = classCodegen.context
+    val state = context.state
     val languageVersionSettings = state.languageVersionSettings
     assert(languageVersionSettings.supportsFeature(LanguageFeature.ReleaseCoroutines)) { "Experimental coroutines are unsupported in JVM_IR backend" }
     val element = if (irFunction.isSuspend)
         irFunction.psiElement ?: classCodegen.irClass.psiElement
     else
-        classCodegen.context.suspendLambdaToOriginalFunctionMap[classCodegen.irClass.attributeOwnerId]!!.psiElement
+        context.suspendLambdaToOriginalFunctionMap[classCodegen.irClass.attributeOwnerId]!!.psiElement
 
     val lineNumber = if (irFunction.isSuspend) {
         val irFile = irFunction.file
@@ -74,6 +76,7 @@ internal fun MethodNode.acceptWithStateMachine(
         internalNameForDispatchReceiver = classCodegen.type.internalName,
         putContinuationParameterToLvt = false,
         initialVarsCountByType = varsCountByType,
+        shouldOptimiseUnusedVariables = !context.configuration.getBoolean(JVMConfigurationKeys.ENABLE_DEBUG_MODE)
     )
     accept(visitor)
 }