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)