JVM_IR: Get rid of non-local related hacks in coroutines inlining

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.backend.jvm.intrinsics.JavaClassProperty
 import org.jetbrains.kotlin.backend.jvm.lower.MultifileFacadeFileEntry
 import org.jetbrains.kotlin.backend.jvm.lower.constantValue
 import org.jetbrains.kotlin.backend.jvm.lower.inlineclasses.unboxInlineClass
+import org.jetbrains.kotlin.backend.jvm.lower.suspendFunctionOriginal
 import org.jetbrains.kotlin.codegen.AsmUtil.*
 import org.jetbrains.kotlin.codegen.BaseExpressionCodegen
 import org.jetbrains.kotlin.codegen.CallGenerator
@@ -442,7 +443,7 @@ class ExpressionCodegen(
             // Check return type of non-lowered suspend call, in order to replace the result of the call with Unit,
             // otherwise, it would seem like the call returns non-unit upon resume.
             // See box/coroutines/tailCallOptimization/unit tests.
-            if (((expression.symbol.owner as? IrSimpleFunction)?.attributeOwnerId as? IrSimpleFunction)?.returnType?.isUnit() == true) {
+            if (expression.symbol.owner.suspendFunctionOriginal().returnType.isUnit()) {
                 addReturnsUnitMarker(mv)
             }
 
@@ -633,21 +634,10 @@ class ExpressionCodegen(
 
     override fun visitReturn(expression: IrReturn, data: BlockInfo): PromisedValue {
         val returnTarget = expression.returnTargetSymbol.owner
-        var owner =
+        val owner =
             returnTarget as? IrFunction
                 ?: (returnTarget as? IrReturnableBlock)?.inlineFunctionSymbol?.owner
                 ?: error("Unsupported IrReturnTarget: $returnTarget")
-        // $$forInline and $suspendImpl functions share the same attributes as originals, but the name is different
-        // fix the return target. TODO: attributes seem to be overloaded. We rely on their uniqueness across transformations to find views
-        // TODO: but in this case, they are not unique enough.
-        if ((irFunction.origin == FOR_INLINE_STATE_MACHINE_TEMPLATE ||
-                    irFunction.origin == SUSPEND_IMPL_STATIC_FUNCTION ||
-                    irFunction.origin == FOR_INLINE_STATE_MACHINE_TEMPLATE_CAPTURES_CROSSINLINE) &&
-            owner.isSuspend && irFunction is IrSimpleFunction && owner is IrSimpleFunction &&
-            irFunction.attributeOwnerId == owner.attributeOwnerId
-        ) {
-            owner = irFunction
-        }
         //TODO: should be owner != irFunction
         val isNonLocalReturn =
             methodSignatureMapper.mapFunctionName(owner) != methodSignatureMapper.mapFunctionName(irFunction)

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

@@ -152,11 +152,6 @@ class IrSourceCompilerForInline(
 
     override fun getContextLabels(): Set<String> {
         val name = codegen.irFunction.name.asString()
-        if (name == INVOKE_SUSPEND_METHOD_NAME) {
-            codegen.context.suspendLambdaToOriginalFunctionMap[codegen.irFunction.parentAsClass.attributeOwnerId]?.let {
-                return setOf(it.name.asString())
-            }
-        }
         return setOf(name)
     }
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt

@@ -56,10 +56,10 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         val suspendLambdas = findSuspendAndInlineLambdas(irFile)
         addContinuationObjectAndContinuationParameterToSuspendFunctions(irFile)
         transformSuspendLambdasIntoContinuations(irFile, suspendLambdas)
-        addContinuationParameterToSuspendCallsAndUpdateNonLocalReturns(irFile)
+        addContinuationParameterToSuspendCallsAndUpdateNonLocalReturns(irFile, suspendLambdas)
     }
 
-    private fun addContinuationParameterToSuspendCallsAndUpdateNonLocalReturns(irFile: IrFile) {
+    private fun addContinuationParameterToSuspendCallsAndUpdateNonLocalReturns(irFile: IrFile, suspendLambdas: List<SuspendLambdaInfo>) {
         irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
             val functionStack = mutableListOf<IrFunction>()
 
@@ -79,8 +79,8 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             override fun visitReturn(expression: IrReturn): IrExpression {
                 val ret = super.visitReturn(expression) as IrReturn
                 val irFunction = expression.returnTargetSymbol.owner as? IrFunction ?: return ret
-                val targetViewOrStub = irFunction.suspendFunctionViewOrStub(context)
-                return IrReturnImpl(ret.startOffset, ret.endOffset, ret.type, targetViewOrStub.symbol, ret.value)
+                val target = suspendLambdas.find { it.function == irFunction }?.invokeSuspend ?: irFunction
+                return IrReturnImpl(ret.startOffset, ret.endOffset, ret.type, target.symbol, ret.value)
             }
         })
     }
@@ -171,6 +171,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             if (info.capturesCrossinline) {
                 addInvokeSuspendForInlineForLambda(invokeSuspend, info.function, parametersFields, receiverField)
             }
+            info.invokeSuspend = invokeSuspend
             info.function.parentAsClass.declarations.remove(info.function)
             if (info.arity <= 1) {
                 val singleParameterField = receiverField ?: parametersWithoutArguments.singleOrNull()
@@ -726,6 +727,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         val capturesCrossinline: Boolean
     ) {
         lateinit var constructor: IrConstructor
+        lateinit var invokeSuspend: IrFunction
     }
 }