JVM_IR: Support interface delegation of suspend functions

The issue was, that built IR function does not have a PSI element,
which is required to report error on suspend functions inside monitors.
In this case, use PSI element of the class, containing the function,
which is consistent with old BE.

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

@@ -14,9 +14,6 @@ import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_IMPL_NAME_SUFFIX
 import org.jetbrains.kotlin.codegen.mangleNameIfNeeded
 import org.jetbrains.kotlin.codegen.state.GenerationState
-import org.jetbrains.kotlin.codegen.ClassBuilderMode
-import org.jetbrains.kotlin.codegen.coroutines.SUSPEND_IMPL_NAME_SUFFIX
-import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
 import org.jetbrains.kotlin.codegen.visitAnnotableParameterCount
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.Modality
@@ -122,8 +119,11 @@ open class FunctionCodegen(
         isAnonymousObject || origin == JvmLoweredDeclarationOrigin.CONTINUATION_CLASS || origin == JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA
 
     private fun psiElement(): KtElement =
-        if (irFunction.isSuspend) irFunction.symbol.descriptor.psiElement as KtElement
-        else context.suspendLambdaToOriginalFunctionMap[irFunction.parentAsClass.attributeOwnerId]!!.symbol.descriptor.psiElement as KtElement
+        (if (irFunction.isSuspend)
+            irFunction.symbol.descriptor.psiElement ?: irFunction.parentAsClass.descriptor.psiElement
+        else
+            context.suspendLambdaToOriginalFunctionMap[irFunction.parentAsClass.attributeOwnerId]!!.symbol.descriptor.psiElement)
+                as KtElement
 
     private fun IrFunction.hasContinuation(): Boolean = isSuspend &&
             // We do not generate continuation and state-machine for synthetic accessors, bridges, and delegated members,

compiler/testData/codegen/box/coroutines/tailCallOptimizations/interfaceDelegation.kt

@@ -0,0 +1,56 @@
+// IGNORE_BACKEND_FIR: JVM_IR
+// TARGET_BACKEND: JVM
+// FULL_JDK
+// WITH_RUNTIME
+// WITH_COROUTINES
+// CHECK_TAIL_CALL_OPTIMIZATION
+import helpers.*
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.*
+
+suspend fun suspendThere(v: String): String = suspendCoroutineUninterceptedOrReturn { x ->
+    TailCallOptimizationChecker.saveStackTrace(x)
+    x.resume(v)
+    COROUTINE_SUSPENDED
+}
+
+interface I {
+    suspend fun suspendHere(): String
+
+    suspend fun suspendHereNoTailCall(): String
+}
+
+class A : I {
+    override suspend fun suspendHere(): String = suspendThere("OK")
+
+    override suspend fun suspendHereNoTailCall(): String {
+        suspendThere("FAIL 2")
+        return "OK"
+    }
+}
+
+open class B(val x: I) : I by x // open override suspend fun suspendHere() = x.suspendHere()
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+fun box(): String {
+    var result = ""
+
+    builder {
+        result = B(A()).suspendHere()
+    }
+    TailCallOptimizationChecker.checkNoStateMachineIn("suspendHere")
+    TailCallOptimizationChecker.checkNoStateMachineIn("suspendHere\$suspendImpl")
+
+    if (result != "OK") return "FAIL 1"
+
+    builder {
+        result = B(A()).suspendHereNoTailCall()
+    }
+    TailCallOptimizationChecker.checkStateMachineIn("suspendHereNoTailCall")
+    TailCallOptimizationChecker.checkNoStateMachineIn("suspendHereNoTailCall\$suspendImpl")
+
+    return result
+}

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -8612,6 +8612,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/tailCallOptimizations/innerObjectRetransformation.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/interfaceDelegation.kt");
+            }
+
             @TestMetadata("simple.kt")
             public void testSimple() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/simple.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -8612,6 +8612,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/tailCallOptimizations/innerObjectRetransformation.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/interfaceDelegation.kt");
+            }
+
             @TestMetadata("simple.kt")
             public void testSimple() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/simple.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/FirBlackBoxCodegenTestGenerated.java

@@ -7517,6 +7517,11 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/tailCallOptimizations/innerObjectRetransformation.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/interfaceDelegation.kt");
+            }
+
             @TestMetadata("simple.kt")
             public void testSimple() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/simple.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -7517,6 +7517,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/tailCallOptimizations/innerObjectRetransformation.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("interfaceDelegation.kt")
+            public void testInterfaceDelegation() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/interfaceDelegation.kt");
+            }
+
             @TestMetadata("simple.kt")
             public void testSimple() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/tailCallOptimizations/simple.kt");