Check for COROUTINE_SUSPENDED inside callable reference

#KT-41429 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/ClosureCodegen.java

@@ -589,4 +589,8 @@ public class ClosureCodegen extends MemberCodegen<KtElement> {
         MemberScope scope = functionClass.getDefaultType().getMemberScope();
         return scope.getContributedFunctions(OperatorNameConventions.INVOKE, NoLookupLocation.FROM_BACKEND).iterator().next();
     }
+
+    public boolean isCallableReference() {
+        return functionReferenceTarget != null;
+    }
 }

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -2742,6 +2742,15 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             addInlineMarker(v, false);
         }
 
+        // If a suspend function returns unboxed inline class, we should check for COROUTINE_SUSPENDED and only then box the inline class.
+        if (insideCallableReference() && resolvedCall.getResultingDescriptor() instanceof FunctionDescriptor) {
+            KotlinType unboxedInlineClass = CoroutineCodegenUtilKt.originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass(
+                    (FunctionDescriptor) resolvedCall.getResultingDescriptor(), typeMapper);
+            if (unboxedInlineClass != null) {
+                CoroutineCodegenUtilKt.generateCoroutineSuspendedCheck(v, state.getLanguageVersionSettings());
+            }
+        }
+
         KotlinType returnType = resolvedCall.getResultingDescriptor().getReturnType();
         if (returnType != null && KotlinBuiltIns.isNothing(returnType)) {
             if (state.getUseKotlinNothingValueException()) {
@@ -2756,8 +2765,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
     }
 
     private boolean insideCallableReference() {
-        return (parentCodegen instanceof ClosureCodegen) &&
-               ((ClosureCodegen) parentCodegen).superClassAsmType.equals(FUNCTION_REFERENCE_IMPL);
+        return (parentCodegen instanceof ClosureCodegen) && ((ClosureCodegen) parentCodegen).isCallableReference();
     }
 
     private void putReceiverAndInlineMarkerIfNeeded(

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

@@ -726,12 +726,7 @@ class CoroutineCodegenForNamedFunction private constructor(
                         with(codegen.v) {
                             // We need to box the returned inline class in resume path.
                             // But first, check for COROUTINE_SUSPENDED, since the function can return it
-                            dup()
-                            loadCoroutineSuspendedMarker(languageVersionSettings)
-                            val elseLabel = Label()
-                            ifacmpne(elseLabel)
-                            areturn(AsmTypes.OBJECT_TYPE)
-                            mark(elseLabel)
+                            generateCoroutineSuspendedCheck(languageVersionSettings)
                             // Now we box the inline class
                             StackValue.coerce(AsmTypes.OBJECT_TYPE, typeMapper.mapType(inlineClassToBoxInInvokeSuspend), this)
                             StackValue.boxInlineClass(inlineClassToBoxInInvokeSuspend, this)

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

@@ -490,6 +490,15 @@ fun InstructionAdapter.loadCoroutineSuspendedMarker(languageVersionSettings: Lan
     )
 }
 
+internal fun InstructionAdapter.generateCoroutineSuspendedCheck(languageVersionSettings: LanguageVersionSettings) {
+    dup()
+    loadCoroutineSuspendedMarker(languageVersionSettings)
+    val elseLabel = Label()
+    ifacmpne(elseLabel)
+    areturn(OBJECT_TYPE)
+    mark(elseLabel)
+}
+
 fun InstructionAdapter.invokeDoResumeWithUnit(thisName: String) {
     // .doResume(Unit, null)
     StackValue.putUnitInstance(this)

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

@@ -7441,6 +7441,11 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReference.kt");
             }
 
+            @TestMetadata("boxReturnValueOfSuspendFunctionReferenceResume.kt")
+            public void testBoxReturnValueOfSuspendFunctionReferenceResume() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReferenceResume.kt");
+            }
+
             @TestMetadata("boxReturnValueOfSuspendLambda.kt")
             public void testBoxReturnValueOfSuspendLambda() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendLambda.kt");

compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReferenceResume.kt

@@ -0,0 +1,30 @@
+// IGNORE_BACKEND: JS_IR
+// IGNORE_BACKEND: JS_IR_ES6
+// WITH_RUNTIME
+// WITH_COROUTINES
+
+import helpers.*
+import kotlin.coroutines.*
+
+inline class R(val x: Any)
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+suspend fun <T> call(fn: suspend () -> T) = fn()
+
+fun useR(r: R) = if (r.x == "OK") "OK" else "fail: $r"
+
+var c: Continuation<R>? = null
+
+suspend fun ok() = suspendCoroutine<R> { c = it }
+
+fun box(): String {
+    var res: String = "fail"
+    builder {
+        res = useR(call(::ok))
+    }
+    c?.resume(R("OK"))
+    return res
+}

compiler/testData/codegen/box/coroutines/inlineClasses/returnResult.kt

@@ -1,5 +1,6 @@
 // WITH_RUNTIME
 // WITH_COROUTINES
+// IGNORE_BACKEND: JVM_IR
 
 import helpers.*
 import kotlin.coroutines.*

compiler/testData/codegen/bytecodeText/coroutines/inlineClasses/noInlineClassBoxingInSuspendFunReturn_Any.kt

@@ -23,4 +23,5 @@ suspend fun test() {
 
 // -- 1 in 'useICAny(suspendGeneric(ICAny("")))
 // -- 1 in 'equals-impl' for ICAny
-// 2 INVOKEVIRTUAL ICAny\.unbox-impl
+// -- 2 on resume path of suspendICAny
+// 4 INVOKEVIRTUAL ICAny\.unbox-impl

compiler/testData/codegen/bytecodeText/coroutines/inlineClasses/noInlineClassBoxingInSuspendFunReturn_InlineAny.kt

@@ -11,14 +11,14 @@ fun useIC(x: IC1) {}
 fun useAny(x: Any) {}
 
 suspend fun test() {
-    useIC(suspendIC())
+    useIC(suspendIC()) // IC1.unbox-impl of resume path
     useIC(suspendGeneric(IC1(IC0("")))) // IC1.box-impl, IC1.unbox-impl
     useAny(suspendAny())
-    useAny(suspendIC()) // IC1.box-impl
+    useAny(suspendIC()) // IC1.box-impl, IC1.unbox-impl of resume path
 }
 
 // 0 INVOKESTATIC IC0\.box-impl
 // 3 INVOKESTATIC IC1\.box-impl
 
 // 1 INVOKEVIRTUAL IC0\.unbox-impl
-// 2 INVOKEVIRTUAL IC1\.unbox-impl
+// 4 INVOKEVIRTUAL IC1\.unbox-impl

compiler/testData/codegen/bytecodeText/coroutines/inlineClasses/noInlineClassBoxingInSuspendFunReturn_String.kt

@@ -23,4 +23,5 @@ suspend fun test() {
 
 // -- 1 in 'useICString(suspendGeneric(ICString("")))
 // -- 1 in 'equals-impl' for ICString
-// 2 INVOKEVIRTUAL ICString\.unbox-impl
+// -- 2 in resume path of suspendICString
+// 4 INVOKEVIRTUAL ICString\.unbox-impl

compiler/testData/codegen/bytecodeText/coroutines/inlineClasses/returnResult.kt

@@ -7,5 +7,5 @@ inline class OurAny(val a: Any)
 
 suspend fun returnsUnboxed(): OurAny = OurAny("OK")
 
-// 1 INVOKESTATIC kotlin/Result.box-impl
+// 0 INVOKESTATIC kotlin/Result.box-impl
 // 0 INVOKESTATIC OurAny.box-impl

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

@@ -8151,6 +8151,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReference.kt");
             }
 
+            @TestMetadata("boxReturnValueOfSuspendFunctionReferenceResume.kt")
+            public void testBoxReturnValueOfSuspendFunctionReferenceResume() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReferenceResume.kt");
+            }
+
             @TestMetadata("boxReturnValueOfSuspendLambda.kt")
             public void testBoxReturnValueOfSuspendLambda() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendLambda.kt");

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

@@ -8151,6 +8151,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReference.kt");
             }
 
+            @TestMetadata("boxReturnValueOfSuspendFunctionReferenceResume.kt")
+            public void testBoxReturnValueOfSuspendFunctionReferenceResume() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReferenceResume.kt");
+            }
+
             @TestMetadata("boxReturnValueOfSuspendLambda.kt")
             public void testBoxReturnValueOfSuspendLambda() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendLambda.kt");

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

@@ -7441,6 +7441,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReference.kt");
             }
 
+            @TestMetadata("boxReturnValueOfSuspendFunctionReferenceResume.kt")
+            public void testBoxReturnValueOfSuspendFunctionReferenceResume() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReferenceResume.kt");
+            }
+
             @TestMetadata("boxReturnValueOfSuspendLambda.kt")
             public void testBoxReturnValueOfSuspendLambda() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendLambda.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/es6/semantics/IrJsCodegenBoxES6TestGenerated.java

@@ -6261,6 +6261,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReference.kt");
             }
 
+            @TestMetadata("boxReturnValueOfSuspendFunctionReferenceResume.kt")
+            public void testBoxReturnValueOfSuspendFunctionReferenceResume() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReferenceResume.kt");
+            }
+
             @TestMetadata("boxReturnValueOfSuspendLambda.kt")
             public void testBoxReturnValueOfSuspendLambda() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendLambda.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -6261,6 +6261,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReference.kt");
             }
 
+            @TestMetadata("boxReturnValueOfSuspendFunctionReferenceResume.kt")
+            public void testBoxReturnValueOfSuspendFunctionReferenceResume() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReferenceResume.kt");
+            }
+
             @TestMetadata("boxReturnValueOfSuspendLambda.kt")
             public void testBoxReturnValueOfSuspendLambda() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendLambda.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -6261,6 +6261,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReference.kt");
             }
 
+            @TestMetadata("boxReturnValueOfSuspendFunctionReferenceResume.kt")
+            public void testBoxReturnValueOfSuspendFunctionReferenceResume() throws Exception {
+                runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendFunctionReferenceResume.kt");
+            }
+
             @TestMetadata("boxReturnValueOfSuspendLambda.kt")
             public void testBoxReturnValueOfSuspendLambda() throws Exception {
                 runTest("compiler/testData/codegen/box/coroutines/inlineClasses/boxReturnValueOfSuspendLambda.kt");