Use non-local return target instead of inline site in suspend function

return type coercion.
 #KT-43226 Fixed

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

@@ -46,6 +46,7 @@ import org.jetbrains.kotlin.config.ApiVersion;
 import org.jetbrains.kotlin.config.JVMAssertionsMode;
 import org.jetbrains.kotlin.config.LanguageFeature;
 import org.jetbrains.kotlin.descriptors.*;
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.SyntheticFieldDescriptor;
 import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor;
@@ -1705,26 +1706,38 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             Type returnType;
             KotlinType returnKotlinType;
             if (isNonLocalReturn) {
-                // This is inline lambda. Find inline-site and check, whether it is suspend functions returning unboxed inline class
+                FunctionDescriptor returnTarget =
-                CodegenContext<?> inlineSiteContext = this.context.getFirstCrossInlineOrNonInlineContext();
+                        nonLocalReturn.descriptor instanceof FunctionDescriptor
-                KotlinType originalInlineClass = null;
+                        ? (FunctionDescriptor) nonLocalReturn.descriptor
-                boolean invokeSuspendOfLambda = false;
+                        : null;
-                FunctionDescriptor inlineSiteDescriptor = null;
+                if (returnTarget == null || !returnTarget.isSuspend()) {
-                if (inlineSiteContext instanceof MethodContext) {
+                    JvmKotlinType jvmKotlinType = nonLocalReturn.getJvmKotlinType(typeMapper);
-                    inlineSiteDescriptor = ((MethodContext) inlineSiteContext).getFunctionDescriptor();
+                    returnType = jvmKotlinType.getType();
-                    originalInlineClass = CoroutineCodegenUtilKt
+                    returnKotlinType = jvmKotlinType.getKotlinType();
-                            .originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass(inlineSiteDescriptor, typeMapper);
+                } else if (returnTarget instanceof AnonymousFunctionDescriptor) {
-                    invokeSuspendOfLambda = CoroutineCodegenUtilKt.isInvokeSuspendOfLambda(inlineSiteDescriptor);
+                    // Suspend lambdas always return Any?
-                }
-                if (originalInlineClass != null) {
-                    returnType = typeMapper.mapType(originalInlineClass);
-                    returnKotlinType = originalInlineClass;
-                } else if (!invokeSuspendOfLambda) {
-                    returnType = nonLocalReturn.returnType.getType();
-                    returnKotlinType = nonLocalReturn.returnType.getKotlinType();
-                } else {
                     returnType = OBJECT_TYPE;
-                    returnKotlinType = inlineSiteDescriptor.getReturnType();
+                    returnKotlinType = state.getModule().getBuiltIns().getNullableAnyType();
+                } else {
+                    // This is inline lambda, but return target is ordinary, yet suspend, function.
+                    // Find inline-site and check, whether it is suspend functions returning unboxed inline class
+                    CodegenContext<?> inlineSiteContext = this.context.getFirstCrossInlineOrNonInlineContext();
+                    KotlinType originalInlineClass = null;
+                    if (inlineSiteContext instanceof MethodContext) {
+                        FunctionDescriptor view = CoroutineCodegenUtilKt.getOrCreateJvmSuspendFunctionView(returnTarget, state);
+                        originalInlineClass =
+                                CoroutineCodegenUtilKt.originalReturnTypeOfSuspendFunctionReturningUnboxedInlineClass(view, typeMapper);
+                    }
+                    if (originalInlineClass != null) {
+                        // As an optimization, suspend functions, returning inline classes with reference underlying
+                        // type return unboxed inline class. Save the type so the coercer will not box it.
+                        returnType = typeMapper.mapType(originalInlineClass);
+                        returnKotlinType = originalInlineClass;
+                    } else {
+                        JvmKotlinType jvmKotlinType = nonLocalReturn.getJvmKotlinType(typeMapper);
+                        returnType = jvmKotlinType.getType();
+                        returnKotlinType = jvmKotlinType.getKotlinType();
+                    }
                 }
             }
             else {
@@ -1788,13 +1801,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
                     FunctionDescriptor containingFunction =
                             BindingContextUtils.getContainingFunctionSkipFunctionLiterals(descriptor, true).getFirst();
                     //FIRST_FUN_LABEL to prevent clashing with existing labels
-                    return new NonLocalReturnInfo(
+                    return new NonLocalReturnInfo(containingFunction, FIRST_FUN_LABEL);
-                            new JvmKotlinType(
-                                    typeMapper.mapReturnType(containingFunction),
-                                    containingFunction.getReturnType()
-                            ),
-                            FIRST_FUN_LABEL
-                    );
                 } else {
                     //local
                     return null;
@@ -1807,10 +1814,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
                 assert element != null : "Expression should be not null " + expression.getText();
                 assert elementDescriptor != null : "Descriptor should be not null: " + element.getText();
                 CallableDescriptor function = (CallableDescriptor) elementDescriptor;
-                return new NonLocalReturnInfo(
+                return new NonLocalReturnInfo(function, expression.getLabelName());
-                        new JvmKotlinType(typeMapper.mapReturnType(function), function.getReturnType()),
-                        expression.getLabelName()
-                );
             }
         }
         return null;
@@ -5427,14 +5431,18 @@ The "returned" value of try expression with no finally is either the last expres
 
     private static class NonLocalReturnInfo {
 
-        private final JvmKotlinType returnType;
+        private final CallableDescriptor descriptor;
 
         private final String labelName;
 
-        private NonLocalReturnInfo(@NotNull JvmKotlinType type, @NotNull String name) {
+        private NonLocalReturnInfo(@NotNull CallableDescriptor descriptor, @NotNull String name) {
-            returnType = type;
+            this.descriptor = descriptor;
             labelName = name;
         }
+
+        private JvmKotlinType getJvmKotlinType(@NotNull KotlinTypeMapper typeMapper) {
+            return new JvmKotlinType(typeMapper.mapReturnType(descriptor), descriptor.getReturnType());
+        }
     }
 
     @NotNull

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

@@ -6883,6 +6883,11 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/noSuspensionPoints.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("nonLocalReturn.kt")
+        public void testNonLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/box/coroutines/nonLocalReturn.kt");
+        }
+
         @TestMetadata("nonLocalReturnFromInlineLambdaDeep.kt")
         public void testNonLocalReturnFromInlineLambdaDeep_1_3() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambdaDeep.kt", "kotlin.coroutines");

compiler/testData/codegen/box/coroutines/inlineClasses/direct/boxUnboxInsideCoroutine_nonLocalReturn.kt

@@ -60,6 +60,18 @@ class Test3 {
     suspend fun test() = bar().s
 }
 
+class Test4 {
+    suspend fun <T> foo(value: T): T = value
+
+    suspend fun bar(): IC? {
+        run {
+            return foo(IC("OK"))
+        }
+    }
+
+    suspend fun test() = bar()!!.s
+}
+
 fun box(): String {
 
     var result = "FAIL"
@@ -83,5 +95,13 @@ fun box(): String {
         result = Test3().test()
     }
 
+    if (result != "OK") return "FAIL 3 $result"
+
+    result = "FAIL 4"
+
+    builder {
+        result = Test4().test()
+    }
+
     return result
 }

compiler/testData/codegen/box/coroutines/inlineClasses/resume/boxUnboxInsideCoroutine_nonLocalReturn.kt

@@ -67,6 +67,18 @@ class Test3 {
     suspend fun test() = bar().s
 }
 
+class Test4 {
+    suspend fun <T> foo(value: T): T = value
+
+    suspend fun bar(): IC? {
+        run {
+            return foo(suspendMe())
+        }
+    }
+
+    suspend fun test() = bar()!!.s
+}
+
 fun box(): String {
 
     var result = "FAIL"
@@ -93,5 +105,14 @@ fun box(): String {
     }
     c?.resume(IC("OK"))
 
+    if (result != "OK") return "FAIL 3 $result"
+
+    result = "FAIL 4"
+
+    builder {
+        result = Test4().test()
+    }
+    c?.resume(IC("OK"))
+
     return result
 }

compiler/testData/codegen/box/coroutines/inlineClasses/resumeWithException/boxUnboxInsideCoroutine_nonLocalReturn.kt

@@ -71,6 +71,18 @@ class Test3 {
     suspend fun test() = bar().s
 }
 
+class Test4 {
+    suspend fun <T> foo(value: T): T = value
+
+    suspend fun bar(): IC? {
+        run {
+            return foo(suspendMe())
+        }
+    }
+
+    suspend fun test() = bar()!!.s
+}
+
 fun box(): String {
     builder {
         Test1().test()
@@ -95,5 +107,12 @@ fun box(): String {
     }
     c?.resumeWithException(IllegalStateException("OK"))
 
+    if (result != "OK") return "FAIL 3 $result"
+
+    builder {
+        Test4().test()
+    }
+    c?.resumeWithException(IllegalStateException("OK"))
+
     return result
 }

compiler/testData/codegen/box/coroutines/nonLocalReturn.kt

@@ -0,0 +1,29 @@
+// KJS_WITH_FULL_RUNTIME
+// WITH_RUNTIME
+import kotlin.coroutines.*
+
+suspend fun coroutineScope(c: suspend () -> Unit) {
+    c()
+}
+
+var counter = 0
+
+suspend fun whatever() = coroutineScope {
+    repeat(10) { // repeat hides a loop, that plays a part in the compiler crash
+        run {
+            counter++
+            return@repeat // required to reproduce the crash
+        }
+    }
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(Continuation(EmptyCoroutineContext) {})
+}
+
+fun box(): String {
+    builder {
+        whatever()
+    }
+    return if (counter != 10) "FAIL $counter" else "OK"
+}

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

@@ -7208,6 +7208,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/noSuspensionPoints.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("nonLocalReturn.kt")
+        public void testNonLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/box/coroutines/nonLocalReturn.kt");
+        }
+
         @TestMetadata("nonLocalReturnFromInlineLambdaDeep.kt")
         public void testNonLocalReturnFromInlineLambdaDeep_1_2() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambdaDeep.kt", "kotlin.coroutines.experimental");

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

@@ -7208,6 +7208,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/noSuspensionPoints.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("nonLocalReturn.kt")
+        public void testNonLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/box/coroutines/nonLocalReturn.kt");
+        }
+
         @TestMetadata("nonLocalReturnFromInlineLambdaDeep.kt")
         public void testNonLocalReturnFromInlineLambdaDeep_1_2() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambdaDeep.kt", "kotlin.coroutines.experimental");

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

@@ -6883,6 +6883,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/noSuspensionPoints.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("nonLocalReturn.kt")
+        public void testNonLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/box/coroutines/nonLocalReturn.kt");
+        }
+
         @TestMetadata("nonLocalReturnFromInlineLambdaDeep.kt")
         public void testNonLocalReturnFromInlineLambdaDeep_1_3() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambdaDeep.kt", "kotlin.coroutines");

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

@@ -5723,6 +5723,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/noSuspensionPoints.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("nonLocalReturn.kt")
+        public void testNonLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/box/coroutines/nonLocalReturn.kt");
+        }
+
         @TestMetadata("nonLocalReturnFromInlineLambdaDeep.kt")
         public void testNonLocalReturnFromInlineLambdaDeep_1_3() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambdaDeep.kt", "kotlin.coroutines");

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

@@ -5723,6 +5723,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/noSuspensionPoints.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("nonLocalReturn.kt")
+        public void testNonLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/box/coroutines/nonLocalReturn.kt");
+        }
+
         @TestMetadata("nonLocalReturnFromInlineLambdaDeep.kt")
         public void testNonLocalReturnFromInlineLambdaDeep_1_3() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambdaDeep.kt", "kotlin.coroutines");

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

@@ -5723,6 +5723,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/noSuspensionPoints.kt", "kotlin.coroutines");
         }
 
+        @TestMetadata("nonLocalReturn.kt")
+        public void testNonLocalReturn() throws Exception {
+            runTest("compiler/testData/codegen/box/coroutines/nonLocalReturn.kt");
+        }
+
         @TestMetadata("nonLocalReturnFromInlineLambdaDeep.kt")
         public void testNonLocalReturnFromInlineLambdaDeep_1_3() throws Exception {
             runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambdaDeep.kt", "kotlin.coroutines");