Support non-local returns to coroutine label from inline lambda

- Use proper class descriptor when retreiving continuation object (see previous commit)
- Use return type as VOID for such cases
- Load correct labels related to coroutine lambda from outer context

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

@@ -1665,10 +1665,20 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
                 tempVariables.put(argumentExpression, valueToReturn);
             }
 
+            // Currently only handleResult members are supported
+            ReceiverValue dispatchReceiver = resolvedCall.getDispatchReceiver();
+            assert dispatchReceiver != null : "Dispatch receiver is null for handleResult to " + resolvedCall.getResultingDescriptor();
+            assert dispatchReceiver instanceof ExtensionReceiver
+                    : "Argument for handleResult call to " + resolvedCall.getResultingDescriptor() +
+                      " should be a coroutine receiver parameter, but " + dispatchReceiver + " found";
+            ClassDescriptor coroutineClassDescriptor =
+                    bindingContext.get(CodegenBinding.CLASS_FOR_CALLABLE, ((ExtensionReceiver) dispatchReceiver).getDeclarationDescriptor());
+            assert coroutineClassDescriptor != null : "Coroutine class descriptor should not be null";
+
             // second argument for handleResult is always Continuation<T> ('this'-object in current implementation)
             tempVariables.put(
                     resolvedCall.getValueArgumentsByIndex().get(1).getArguments().get(0).getArgumentExpression(),
-                    StackValue.thisOrOuter(this, context.getThisDescriptor(), false, false));
+                    StackValue.thisOrOuter(this, coroutineClassDescriptor, false, false));
 
             return invokeFunction(resolvedCall, StackValue.none());
         }
@@ -1955,6 +1965,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
 
                 if (handleResultValue != null) {
                     handleResultValue.put(Type.VOID_TYPE, v);
+                    returnType = Type.VOID_TYPE;
                 }
                 else if (returnedExpression != null && valueToReturn != null) {
                     putStackValue(returnedExpression, returnType, valueToReturn);

compiler/backend/src/org/jetbrains/kotlin/codegen/context/ClosureContext.java

@@ -26,6 +26,24 @@ import static org.jetbrains.kotlin.codegen.binding.CodegenBinding.anonymousClass
 
 public class ClosureContext extends ClassContext {
     private final FunctionDescriptor functionDescriptor;
+    private final FunctionDescriptor coroutineDescriptor;
+
+    public ClosureContext(
+            @NotNull KotlinTypeMapper typeMapper,
+            @NotNull FunctionDescriptor functionDescriptor,
+            @Nullable CodegenContext parentContext,
+            @NotNull LocalLookup localLookup,
+            // original coroutine lambda descriptor having return type T (the type that can be returned within lambda)
+            @Nullable FunctionDescriptor coroutineDescriptor
+    ) {
+        super(typeMapper,
+              anonymousClassForCallable(
+                      typeMapper.getBindingContext(), coroutineDescriptor != null ? coroutineDescriptor : functionDescriptor),
+              OwnerKind.IMPLEMENTATION, parentContext, localLookup);
+
+        this.functionDescriptor = functionDescriptor;
+        this.coroutineDescriptor = coroutineDescriptor;
+    }
 
     public ClosureContext(
             @NotNull KotlinTypeMapper typeMapper,
@@ -33,10 +51,7 @@ public class ClosureContext extends ClassContext {
             @Nullable CodegenContext parentContext,
             @NotNull LocalLookup localLookup
     ) {
-        super(typeMapper, anonymousClassForCallable(typeMapper.getBindingContext(), functionDescriptor),
-              OwnerKind.IMPLEMENTATION, parentContext, localLookup);
-
-        this.functionDescriptor = functionDescriptor;
+        this(typeMapper, functionDescriptor, parentContext, localLookup, null);
     }
 
     @NotNull
@@ -48,4 +63,9 @@ public class ClosureContext extends ClassContext {
     public String toString() {
         return "Closure: " + getContextDescriptor();
     }
+
+    @Nullable
+    public FunctionDescriptor getCoroutineDescriptor() {
+        return coroutineDescriptor;
+    }
 }

compiler/backend/src/org/jetbrains/kotlin/codegen/context/CodegenContext.java

@@ -333,6 +333,18 @@ public abstract class CodegenContext<T extends DeclarationDescriptor> {
         return new ClosureContext(typeMapper, funDescriptor, this, localLookup);
     }
 
+    @NotNull
+    public ClosureContext intoCoroutineClosure(
+            // copy of lambda descriptor that have Continuation<Unit> return type (as it's ep)
+            @NotNull FunctionDescriptor funDescriptorReturningContinuation,
+            // original coroutine lambda descriptor having return type T (the type that can be returned within lambda)
+            @NotNull FunctionDescriptor coroutineLambdaDescriptor,
+            @NotNull LocalLookup localLookup,
+            @NotNull KotlinTypeMapper typeMapper
+    ) {
+        return new ClosureContext(typeMapper, funDescriptorReturningContinuation, this, localLookup, coroutineLambdaDescriptor);
+    }
+
     @Nullable
     public CodegenContext getParentContext() {
         return parentContext;

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

@@ -180,7 +180,8 @@ class CoroutineCodegen(
             return CoroutineCodegen(
                     state,
                     declaration,
-                    expressionCodegen.context.intoClosure(descriptorWithContinuationReturnType, expressionCodegen, state.typeMapper),
+                    expressionCodegen.context.intoCoroutineClosure(
+                            descriptorWithContinuationReturnType, originalCoroutineLambdaDescriptor, expressionCodegen, state.typeMapper),
                     FunctionGenerationStrategy.FunctionDefault(state, declaration), expressionCodegen.parentCodegen, classBuilder,
                     continuationSupertype,
                     originalCoroutineLambdaDescriptor.controllerTypeIfCoroutine!!)

compiler/backend/src/org/jetbrains/kotlin/codegen/inline/InlineCodegen.java

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.codegen.context.*;
 import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicArrayConstructorsKt;
 import org.jetbrains.kotlin.codegen.state.GenerationState;
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper;
+import org.jetbrains.kotlin.coroutines.CoroutineUtilKt;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCache;
 import org.jetbrains.kotlin.name.ClassId;
@@ -385,7 +386,7 @@ public class InlineCodegen extends CallGenerator {
         InlineResult result = inliner.doInline(adapter, remapper, true, LabelOwner.SKIP_ALL);
         result.getReifiedTypeParametersUsages().mergeAll(reificationResult);
 
-        CallableMemberDescriptor descriptor = codegen.getContext().getContextDescriptor();
+        CallableMemberDescriptor descriptor = getLabelOwnerDescriptor(codegen.getContext());
         final Set<String> labels = getDeclarationLabels(DescriptorToSourceUtils.descriptorToDeclaration(descriptor), descriptor);
         LabelOwner labelOwner = new LabelOwner() {
             @Override
@@ -410,6 +411,17 @@ public class InlineCodegen extends CallGenerator {
         return result;
     }
 
+    @NotNull
+    private static CallableMemberDescriptor getLabelOwnerDescriptor(@NotNull MethodContext context) {
+        if (context.getParentContext() instanceof ClosureContext &&
+            ((ClosureContext) context.getParentContext()).getCoroutineDescriptor() != null) {
+            //noinspection ConstantConditions
+            return ((ClosureContext) context.getParentContext()).getCoroutineDescriptor();
+        }
+
+        return context.getContextDescriptor();
+    }
+
     private static void removeStaticInitializationTrigger(@NotNull MethodNode methodNode) {
         InsnList insnList = methodNode.instructions;
         AbstractInsnNode insn = insnList.getFirst();

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

@@ -0,0 +1,43 @@
+class Controller {
+    var cResult = 0
+    suspend fun suspendHere(v: Int, x: Continuation<Int>) {
+        x.resume(v * 2)
+    }
+
+    operator fun handleResult(x: Int, y: Continuation<Nothing>) {
+        cResult = x
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>): Controller {
+    val controller = Controller()
+    c(controller).resume(Unit)
+
+    return controller
+}
+
+inline fun foo(x: (Int) -> Unit) {
+    for (i in 1..2) {
+        x(i)
+    }
+}
+
+fun box(): String {
+    var result = ""
+
+    val controllerResult = builder {
+        result += "-"
+        foo {
+            result += suspendHere(it).toString()
+            if (it == 2) return@builder 56
+        }
+        // Should be unreachable
+        result += "+"
+        1
+    }.cResult
+
+    if (result != "-24") return "fail 1: $result"
+    if (controllerResult != 56) return "fail 2: $controllerResult"
+
+    return "OK"
+}

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

@@ -0,0 +1,49 @@
+// WITH_RUNTIME
+
+class Controller {
+    var cResult = 0
+    suspend fun suspendHere(v: Int, x: Continuation<Int>) {
+        x.resume(v * 2)
+    }
+
+    operator fun handleResult(x: Int, y: Continuation<Nothing>) {
+        cResult = x
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>): Controller {
+    val controller = Controller()
+    c(controller).resume(Unit)
+
+    return controller
+}
+
+inline fun foo(x: (Int) -> Unit) {
+    for (i in 1..2) {
+        run {
+            x(i)
+        }
+    }
+}
+
+fun box(): String {
+    var result = ""
+
+    val controllerResult = builder {
+        result += "-"
+        foo {
+            run {
+                result += suspendHere(it).toString()
+                if (it == 2) return@builder 56
+            }
+        }
+        // Should be unreachable
+        result += "+"
+        1
+    }.cResult
+
+    if (result != "-24") return "fail 1: $result"
+    if (controllerResult != 56) return "fail 2: $controllerResult"
+
+    return "OK"
+}

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

@@ -4123,6 +4123,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("nonLocalReturnFromInlineLambda.kt")
+        public void testNonLocalReturnFromInlineLambda() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambda.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("nonLocalReturnFromInlineLambdaDeep.kt")
+        public void testNonLocalReturnFromInlineLambdaDeep() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/nonLocalReturnFromInlineLambdaDeep.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("returnByLabel.kt")
         public void testReturnByLabel() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/returnByLabel.kt");