JVM_IR: move state machine generation to ClassCodegen

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

@@ -354,7 +354,22 @@ open class ClassCodegen protected constructor(
         }
 
         val node = FunctionCodegen(method, this).generate()
-        node.accept(with(node) { visitor.newMethod(method.OtherOrigin, access, name, desc, signature, exceptions.toTypedArray()) })
+        val mv = with(node) { visitor.newMethod(method.OtherOrigin, access, name, desc, signature, exceptions.toTypedArray()) }
+        if (method.hasContinuation() || method.isInvokeSuspendOfLambda()) {
+            // Generate a state machine within this method. The continuation class for it should be generated
+            // lazily so that if tail call optimization kicks in, the unused class will not be written to the output.
+            val continuationClassCodegen = lazy {
+                createLocalClassCodegen(method.continuationClass()!!, method).also { it.generate() }
+            }
+            node.acceptWithStateMachine(method, this, mv) {
+                if (method.isSuspend) continuationClassCodegen.value.visitor else visitor
+            }
+            if (continuationClassCodegen.isInitialized() || method.alwaysNeedsContinuation()) {
+                continuationClassCodegen.value.done()
+            }
+        } else {
+            node.accept(mv)
+        }
         jvmSignatureClashDetector.trackMethod(method, RawSignature(node.name, node.desc, MemberKind.METHOD))
 
         val signature = Method(node.name, node.desc)

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

@@ -30,18 +30,16 @@ import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.isSuspend
 import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.org.objectweb.asm.MethodVisitor
+import org.jetbrains.org.objectweb.asm.tree.MethodNode
 
-internal fun generateStateMachine(
+internal fun MethodNode.acceptWithStateMachine(
     irFunction: IrFunction,
     classCodegen: ClassCodegen,
     methodVisitor: MethodVisitor,
-    access: Int,
-    signature: JvmMethodGenericSignature,
     obtainContinuationClassBuilder: () -> ClassBuilder
-): MethodVisitor {
+) {
     val state = classCodegen.context.state
     val languageVersionSettings = state.languageVersionSettings
     assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
@@ -49,8 +47,8 @@ internal fun generateStateMachine(
         irFunction.symbol.descriptor.psiElement ?: classCodegen.irClass.descriptor.psiElement
     else
         classCodegen.context.suspendLambdaToOriginalFunctionMap[classCodegen.irClass.attributeOwnerId]!!.symbol.descriptor.psiElement
-    return CoroutineTransformerMethodVisitor(
-        methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, signature.genericsSignature, null,
+    val visitor = CoroutineTransformerMethodVisitor(
+        methodVisitor, access, name, desc, signature, exceptions.toTypedArray(),
         obtainClassBuilderForCoroutineState = obtainContinuationClassBuilder,
         reportSuspensionPointInsideMonitor = { reportSuspensionPointInsideMonitor(element as KtElement, state, it) },
         lineNumber = element?.let { CodegenUtil.getLineNumberForElement(it, false) } ?: 0,
@@ -67,9 +65,10 @@ internal fun generateStateMachine(
             it.returnType.isUnit() && it.anyOfOverriddenFunctionsReturnsNonUnit()
         }
     )
+    accept(visitor)
 }
 
-internal fun IrFunction.anyOfOverriddenFunctionsReturnsNonUnit(): Boolean {
+private fun IrFunction.anyOfOverriddenFunctionsReturnsNonUnit(): Boolean {
     return (this as? IrSimpleFunction)?.allOverridden()?.toList()?.let { functions ->
         functions.isNotEmpty() && functions.any { !it.returnType.isUnit() }
     } == true

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

@@ -61,7 +61,7 @@ class FunctionCodegen(
             signature.genericsSignature.takeIf { flags.and(Opcodes.ACC_SYNTHETIC) == 0 },
             getThrownExceptions(irFunction)?.toTypedArray()
         )
-        var methodVisitor: MethodVisitor = wrapWithMaxLocalCalc(methodNode)
+        val methodVisitor: MethodVisitor = wrapWithMaxLocalCalc(methodNode)
 
         if (context.state.generateParametersMetadata && flags.and(Opcodes.ACC_SYNTHETIC) == 0) {
             generateParameterNames(irFunction, methodVisitor, signature, context.state)
@@ -86,19 +86,10 @@ class FunctionCodegen(
             }
         }
 
-        val continuationClassCodegen = lazy {
-            classCodegen.createLocalClassCodegen(irFunction.continuationClass()!!, irFunction).also { it.generate() }
-        }
         if (!context.state.classBuilderMode.generateBodies || flags.and(Opcodes.ACC_ABSTRACT) != 0 || irFunction.isExternal) {
             generateAnnotationDefaultValueIfNeeded(methodVisitor)
         } else {
             val frameMap = irFunction.createFrameMapWithReceivers()
-            if (irFunction.hasContinuation() || irFunction.isInvokeSuspendOfLambda()) {
-                methodVisitor = generateStateMachine(irFunction, classCodegen, methodVisitor, flags, signature) {
-                    // This has to be done lazily to avoid generating the class if tail call optimization makes it redundant.
-                    if (irFunction.isSuspend) continuationClassCodegen.value.visitor else classCodegen.visitor
-                }
-            }
             context.state.globalInlineContext.enterDeclaration(irFunction.suspendFunctionOriginal().descriptor)
             try {
                 val adapter = InstructionAdapter(methodVisitor)
@@ -109,9 +100,6 @@ class FunctionCodegen(
             methodVisitor.visitMaxs(-1, -1)
         }
         methodVisitor.visitEnd()
-        if (continuationClassCodegen.isInitialized() || irFunction.alwaysNeedsContinuation()) {
-            continuationClassCodegen.value.done()
-        }
         return methodNode
     }
 

compiler/testData/codegen/box/coroutines/varSpilling/refinedIntTypesAnalysis.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // WITH_RUNTIME
 
 class Test {