JVM_IR: unify generateStateMachineFor{Lambda,NamedFunction}

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

@@ -33,7 +33,7 @@ import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 
-internal fun generateStateMachineForNamedFunction(
+internal fun generateStateMachine(
     irFunction: IrFunction,
     classCodegen: ClassCodegen,
     methodVisitor: MethodVisitor,
@@ -42,12 +42,11 @@ internal fun generateStateMachineForNamedFunction(
     obtainContinuationClassBuilder: () -> ClassBuilder,
     element: KtElement
 ): MethodVisitor {
-    assert(irFunction.isSuspend)
     val state = classCodegen.state
     val languageVersionSettings = state.languageVersionSettings
     assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
     return CoroutineTransformerMethodVisitor(
-        methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, null, null,
+        methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, signature.genericsSignature, null,
         obtainClassBuilderForCoroutineState = obtainContinuationClassBuilder,
         reportSuspensionPointInsideMonitor = { reportSuspensionPointInsideMonitor(element, state, it) },
         lineNumber = CodegenUtil.getLineNumberForElement(element, false) ?: 0,
@@ -55,12 +54,12 @@ internal fun generateStateMachineForNamedFunction(
         languageVersionSettings = languageVersionSettings,
         shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
         containingClassInternalName = classCodegen.visitor.thisName,
-        isForNamedFunction = true,
-        needDispatchReceiver = irFunction.dispatchReceiverParameter != null
-                || irFunction.origin == JvmLoweredDeclarationOrigin.SUSPEND_IMPL_STATIC_FUNCTION,
+        isForNamedFunction = irFunction.isSuspend, // SuspendLambda.invokeSuspend is not suspend
+        needDispatchReceiver = irFunction.isSuspend && (irFunction.dispatchReceiverParameter != null
+                || irFunction.origin == JvmLoweredDeclarationOrigin.SUSPEND_IMPL_STATIC_FUNCTION),
         internalNameForDispatchReceiver = classCodegen.visitor.thisName,
         putContinuationParameterToLvt = false,
-        disableTailCallOptimizationForFunctionReturningUnit = irFunction.suspendFunctionOriginal().let {
+        disableTailCallOptimizationForFunctionReturningUnit = irFunction.isSuspend && irFunction.suspendFunctionOriginal().let {
             it.returnType.isUnit() && it.anyOfOverriddenFunctionsReturnsNonUnit()
         }
     )
@@ -72,30 +71,6 @@ internal fun IrFunction.anyOfOverriddenFunctionsReturnsNonUnit(): Boolean {
     } == true
 }
 
-internal fun generateStateMachineForLambda(
-    classCodegen: ClassCodegen,
-    methodVisitor: MethodVisitor,
-    access: Int,
-    signature: JvmMethodGenericSignature,
-    element: KtElement
-): MethodVisitor {
-    val state = classCodegen.state
-    val languageVersionSettings = state.languageVersionSettings
-    assert(languageVersionSettings.isReleaseCoroutines()) { "Experimental coroutines are unsupported in JVM_IR backend" }
-    return CoroutineTransformerMethodVisitor(
-        methodVisitor, access, signature.asmMethod.name, signature.asmMethod.descriptor, signature.genericsSignature, null,
-        obtainClassBuilderForCoroutineState = { classCodegen.visitor },
-        reportSuspensionPointInsideMonitor = { reportSuspensionPointInsideMonitor(element, state, it) },
-        lineNumber = CodegenUtil.getLineNumberForElement(element, false) ?: 0,
-        sourceFile = classCodegen.irClass.file.name,
-        languageVersionSettings = languageVersionSettings,
-        shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization,
-        containingClassInternalName = classCodegen.visitor.thisName,
-        isForNamedFunction = false,
-        disableTailCallOptimizationForFunctionReturningUnit = false
-    )
-}
-
 internal fun IrFunction.continuationParameter(): IrValueParameter? = when {
     isInvokeSuspendOfLambda() || isInvokeSuspendForInlineOfLambda() -> dispatchReceiverParameter
     else -> valueParameters.singleOrNull { it.origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS }

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

@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.backend.common.lower.BOUND_VALUE_PARAMETER
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.codegen.AsmUtil
-import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
 import org.jetbrains.kotlin.codegen.mangleNameIfNeeded
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.codegen.visitAnnotableParameterCount
@@ -31,9 +30,7 @@ import org.jetbrains.kotlin.resolve.jvm.annotations.SYNCHRONIZED_ANNOTATION_FQ_N
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterKind
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature
-import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
-import org.jetbrains.kotlin.utils.sure
 import org.jetbrains.org.objectweb.asm.*
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 
@@ -93,29 +90,25 @@ open class FunctionCodegen(
             generateAnnotationDefaultValueIfNeeded(methodVisitor)
         } else {
             val frameMap = createFrameMapWithReceivers()
-            methodVisitor = when {
-                irFunction.hasContinuation() -> {
-                    if (irFunction.parentAsClass.declarations.any {
-                        (it as? IrAttributeContainer)?.attributeOwnerId == (irFunction as IrAttributeContainer).attributeOwnerId &&
-                                it.nameForIrSerialization.asString() == irFunction.nameForIrSerialization.asString() + FOR_INLINE_SUFFIX &&
+            if (irFunction.hasContinuation() || irFunction.isInvokeSuspendOfLambda()) {
+                if (irFunction is IrSimpleFunction && irFunction.parentAsClass.declarations.any {
+                        it is IrSimpleFunction && it.attributeOwnerId == irFunction.attributeOwnerId &&
                                 it.origin == JvmLoweredDeclarationOrigin.FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE
-                        }
-                    ) {
-                        // force generation of fake continuation for inliner.
-                        continuationClassCodegen.value
                     }
-                    generateStateMachineForNamedFunction(
-                        irFunction, classCodegen, methodVisitor,
-                        access = flags,
-                        signature = signature,
-                        obtainContinuationClassBuilder = { continuationClassCodegen.value.visitor },
-                        element = psiElement()
-                    )
+                ) {
+                    // Force generation of fake continuation for inliner.
+                    continuationClassCodegen.value
                 }
-                irFunction.isInvokeSuspendOfLambda() -> generateStateMachineForLambda(
-                    classCodegen, methodVisitor, flags, signature, psiElement()
+                // This has to be done lazily to avoid generating the class if tail call optimization makes it redundant.
+                val getContinuation = {
+                    if (irFunction.isSuspend)
+                        continuationClassCodegen.value.visitor
+                    else
+                        classCodegen.visitor
+                }
+                methodVisitor = generateStateMachine(
+                    irFunction, classCodegen, methodVisitor, flags, signature, getContinuation, psiElement()
                 )
-                else -> methodVisitor
             }
             ExpressionCodegen(irFunction, signature, frameMap, InstructionAdapter(methodVisitor), classCodegen, inlinedInto).generate()
             methodVisitor.visitMaxs(-1, -1)