Fix codegen issue on suspend functions with receiver

When the extension receiver of a named suspend function was marked as used
from the inner lambda, codegen used to throw a "Don't know how to generate outer expression for class" exception.

It may seem quite tricky, but currently for suspend lambda body
its extension receiver is treated as it's defined inside the relevant "doResume"
(there is an actual bytecode part that fills the relevant local variable)

The problem was that inside ExpressionCodegen for "doResume" of named
suspend function we couldn't determine that original function has
an extension receiver.

 #KT-15821 Fixed
 #KT-15820 Fixed

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

@@ -1742,7 +1742,7 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
         if (captureReceiver != null) {
             StackValue capturedReceiver =
                     functionReferenceReceiver != null ? functionReferenceReceiver :
-                    generateExtensionReceiver(unwrapOriginalDescriptorForSuspendLambda(context));
+                    generateExtensionReceiver(unwrapOriginalReceiverOwnerForSuspendFunction(context));
             callGenerator.putCapturedValueOnStack(capturedReceiver, capturedReceiver.type, paramIndex++);
         }
 
@@ -1781,9 +1781,13 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
     }
 
     @NotNull
-    private static CallableDescriptor unwrapOriginalDescriptorForSuspendLambda(@NotNull MethodContext context) {
-        FunctionDescriptor originalSuspendLambdaDescriptor = getOriginalSuspendLambdaDescriptorFromContext(context);
-        if (originalSuspendLambdaDescriptor != null) return originalSuspendLambdaDescriptor;
+    private static CallableDescriptor unwrapOriginalReceiverOwnerForSuspendFunction(@NotNull MethodContext context) {
+        FunctionDescriptor originalForDoResume =
+                context.getFunctionDescriptor().getUserData(CoroutineCodegenUtilKt.INITIAL_SUSPEND_DESCRIPTOR_FOR_DO_RESUME);
+
+        if (originalForDoResume != null) {
+            return originalForDoResume;
+        }
 
         if (context.getFunctionDescriptor().isSuspend()) {
             return CoroutineCodegenUtilKt.unwrapInitialDescriptorForSuspendFunction(context.getFunctionDescriptor());

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

@@ -44,7 +44,6 @@ import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.types.typeUtil.makeNullable
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.kotlin.utils.singletonOrEmptyList
-import org.jetbrains.org.objectweb.asm.AnnotationVisitor
 import org.jetbrains.org.objectweb.asm.MethodVisitor
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
@@ -52,19 +51,19 @@ import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 import org.jetbrains.org.objectweb.asm.commons.Method
 
 
-class CoroutineCodegen(
+class CoroutineCodegen private constructor(
         outerExpressionCodegen: ExpressionCodegen,
         element: KtElement,
         private val closureContext: ClosureContext,
         classBuilder: ClassBuilder,
-        private val originalSuspendLambdaDescriptor: FunctionDescriptor?
+        private val originalSuspendFunctionDescriptor: FunctionDescriptor,
+        private val isSuspendLambda: Boolean
 ) : ClosureCodegen(
         outerExpressionCodegen.state,
         element, null, closureContext, null,
         FailingFunctionGenerationStrategy,
         outerExpressionCodegen.parentCodegen, classBuilder
 ) {
-
     private val classDescriptor = closureContext.contextDescriptor
     private val builtIns = funDescriptor.builtIns
 
@@ -98,7 +97,8 @@ class CoroutineCodegen(
                         ),
                         builtIns.nullableAnyType,
                         Modality.FINAL,
-                        Visibilities.PUBLIC
+                        Visibilities.PUBLIC,
+                        mapOf(INITIAL_SUSPEND_DESCRIPTOR_FOR_DO_RESUME to originalSuspendFunctionDescriptor)
                 )
             }
 
@@ -130,14 +130,14 @@ class CoroutineCodegen(
     }
 
     override fun generateBridges() {
-        if (originalSuspendLambdaDescriptor == null) return
+        if (!isSuspendLambda) return
         super.generateBridges()
     }
 
     override fun generateBody() {
         super.generateBody()
 
-        if (originalSuspendLambdaDescriptor == null) return
+        if (!isSuspendLambda) return
 
         // create() = ...
         functionCodegen.generateMethod(JvmDeclarationOrigin.NO_ORIGIN, createCoroutineDescriptor,
@@ -224,7 +224,7 @@ class CoroutineCodegen(
     }
 
     private fun generateCreateCoroutineMethod(codegen: ExpressionCodegen) {
-        assert(originalSuspendLambdaDescriptor != null) { "create method should only be generated for suspend lambdas" }
+        assert(isSuspendLambda) { "create method should only be generated for suspend lambdas" }
         val classDescriptor = closureContext.contextDescriptor
         val owner = typeMapper.mapClass(classDescriptor)
 
@@ -273,7 +273,10 @@ class CoroutineCodegen(
     }
 
     private fun allLambdaParameters() =
-            originalSuspendLambdaDescriptor?.extensionReceiverParameter.singletonOrEmptyList() + originalSuspendLambdaDescriptor?.valueParameters.orEmpty()
+            if (isSuspendLambda)
+                originalSuspendFunctionDescriptor.extensionReceiverParameter.singletonOrEmptyList() + originalSuspendFunctionDescriptor.valueParameters.orEmpty()
+            else
+                emptyList()
 
     private fun ExpressionCodegen.generateLoadField(fieldInfo: FieldInfo) {
         StackValue.field(fieldInfo, generateThisOrOuter(context.thisDescriptor, false)).put(fieldInfo.fieldType, v)
@@ -304,7 +307,7 @@ class CoroutineCodegen(
     }
 
     override fun generateKotlinMetadataAnnotation() {
-        if (originalSuspendLambdaDescriptor != null) {
+        if (isSuspendLambda) {
             super.generateKotlinMetadataAnnotation()
         }
         else {
@@ -334,7 +337,8 @@ class CoroutineCodegen(
                             originalSuspendLambdaDescriptor, expressionCodegen, expressionCodegen.state.typeMapper
                     ),
                     classBuilder,
-                    originalSuspendLambdaDescriptor
+                    originalSuspendLambdaDescriptor,
+                    isSuspendLambda = true
             )
         }
 
@@ -356,7 +360,8 @@ class CoroutineCodegen(
                             originalSuspendDescriptor, expressionCodegen, expressionCodegen.state.typeMapper
                     ),
                     cv,
-                    originalSuspendLambdaDescriptor = null
+                    originalSuspendDescriptor,
+                    isSuspendLambda = false
             )
         }
     }

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

@@ -67,6 +67,10 @@ data class ResolvedCallWithRealDescriptor(val resolvedCall: ResolvedCall<*>, val
 
 @JvmField
 val INITIAL_DESCRIPTOR_FOR_SUSPEND_FUNCTION = object : FunctionDescriptor.UserDataKey<FunctionDescriptor> {}
+
+@JvmField
+val INITIAL_SUSPEND_DESCRIPTOR_FOR_DO_RESUME = object : FunctionDescriptor.UserDataKey<FunctionDescriptor> {}
+
 // Resolved calls to suspension function contain descriptors as they visible within coroutines:
 // E.g. `fun <V> await(f: CompletableFuture<V>): V` instead of `fun <V> await(f: CompletableFuture<V>, machine: Continuation<V>): Unit`
 // See `createJvmSuspendFunctionView` and it's usages for clarification

compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/inlineTwoReceivers.kt

@@ -0,0 +1,36 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// TARGET_BACKEND: JVM
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.SUSPENDED_MARKER
+import kotlin.coroutines.intrinsics.suspendCoroutineOrReturn
+
+class MyTest {
+    suspend fun act(value: String): String = suspendCoroutineOrReturn {
+        it.resume(value)
+        SUSPENDED_MARKER
+    }
+}
+
+inline suspend fun <T> testAsync(routine: suspend MyTest.() -> T): T = routine(MyTest())
+
+inline suspend fun Iterable<String>.test() = testAsync {
+    var sum = ""
+    for (v in this@test) {
+        sum += act(v)
+    }
+    sum
+}
+
+fun builder(x: suspend () -> Unit) {
+    x.startCoroutine(EmptyContinuation)
+}
+
+fun box(): String {
+    var res = ""
+    builder {
+        res = listOf("O", "K").test()
+    }
+
+    return res
+}

compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/noinlineTwoReceivers.kt

@@ -0,0 +1,35 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+import kotlin.coroutines.*
+import kotlin.coroutines.intrinsics.SUSPENDED_MARKER
+import kotlin.coroutines.intrinsics.suspendCoroutineOrReturn
+
+class MyTest {
+    suspend fun act(value: String): String = suspendCoroutineOrReturn {
+        it.resume(value)
+        SUSPENDED_MARKER
+    }
+}
+
+suspend fun <T> testAsync(routine: suspend MyTest.() -> T): T = routine(MyTest())
+
+suspend fun Iterable<String>.test() = testAsync {
+    var sum = ""
+    for (v in this@test) {
+        sum += act(v)
+    }
+    sum
+}
+
+fun builder(x: suspend () -> Unit) {
+    x.startCoroutine(EmptyContinuation)
+}
+
+fun box(): String {
+    var res = ""
+    builder {
+        res = listOf("O", "K").test()
+    }
+
+    return res
+}

compiler/testData/codegen/light-analysis/coroutines/suspendFunctionAsCoroutine/inlineTwoReceivers.txt

@@ -0,0 +1,38 @@
+@kotlin.Metadata
+public final class CoroutineUtilKt {
+    public final static @org.jetbrains.annotations.NotNull method handleExceptionContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+    public final static @org.jetbrains.annotations.NotNull method handleResultContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+}
+
+@kotlin.Metadata
+public class EmptyContinuation {
+    public final static field Companion: EmptyContinuation.Companion
+    private final @org.jetbrains.annotations.NotNull field context: kotlin.coroutines.CoroutineContext
+    inner class EmptyContinuation/Companion
+    public @synthetic.kotlin.jvm.GeneratedByJvmOverloads method <init>(): void
+    public method <init>(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.CoroutineContext): void
+    public synthetic method <init>(p0: kotlin.coroutines.CoroutineContext, p1: int, p2: kotlin.jvm.internal.DefaultConstructorMarker): void
+    public @org.jetbrains.annotations.NotNull method getContext(): kotlin.coroutines.CoroutineContext
+    public method resume(@org.jetbrains.annotations.Nullable p0: java.lang.Object): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+}
+
+@kotlin.Metadata
+public final static class EmptyContinuation/Companion {
+    inner class EmptyContinuation/Companion
+    private method <init>(): void
+}
+
+@kotlin.Metadata
+public final class InlineTwoReceiversKt {
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void
+    public final static @org.jetbrains.annotations.Nullable method test(@org.jetbrains.annotations.NotNull p0: java.lang.Iterable, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
+    public final static @org.jetbrains.annotations.Nullable method testAsync(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function2, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
+}
+
+@kotlin.Metadata
+public final class MyTest {
+    public method <init>(): void
+    public final @org.jetbrains.annotations.Nullable method act(@org.jetbrains.annotations.NotNull p0: java.lang.String, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
+}

compiler/testData/codegen/light-analysis/coroutines/suspendFunctionAsCoroutine/noinlineTwoReceivers.txt

@@ -0,0 +1,38 @@
+@kotlin.Metadata
+public final class CoroutineUtilKt {
+    public final static @org.jetbrains.annotations.NotNull method handleExceptionContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+    public final static @org.jetbrains.annotations.NotNull method handleResultContinuation(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): kotlin.coroutines.Continuation
+}
+
+@kotlin.Metadata
+public class EmptyContinuation {
+    public final static field Companion: EmptyContinuation.Companion
+    private final @org.jetbrains.annotations.NotNull field context: kotlin.coroutines.CoroutineContext
+    inner class EmptyContinuation/Companion
+    public @synthetic.kotlin.jvm.GeneratedByJvmOverloads method <init>(): void
+    public method <init>(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.CoroutineContext): void
+    public synthetic method <init>(p0: kotlin.coroutines.CoroutineContext, p1: int, p2: kotlin.jvm.internal.DefaultConstructorMarker): void
+    public @org.jetbrains.annotations.NotNull method getContext(): kotlin.coroutines.CoroutineContext
+    public method resume(@org.jetbrains.annotations.Nullable p0: java.lang.Object): void
+    public method resumeWithException(@org.jetbrains.annotations.NotNull p0: java.lang.Throwable): void
+}
+
+@kotlin.Metadata
+public final static class EmptyContinuation/Companion {
+    inner class EmptyContinuation/Companion
+    private method <init>(): void
+}
+
+@kotlin.Metadata
+public final class MyTest {
+    public method <init>(): void
+    public final @org.jetbrains.annotations.Nullable method act(@org.jetbrains.annotations.NotNull p0: java.lang.String, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
+}
+
+@kotlin.Metadata
+public final class NoinlineTwoReceiversKt {
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void
+    public final static @org.jetbrains.annotations.Nullable method test(@org.jetbrains.annotations.NotNull p0: java.lang.Iterable, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
+    public final static @org.jetbrains.annotations.Nullable method testAsync(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function2, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
+}

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

@@ -5161,12 +5161,24 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 doTest(fileName);
             }
 
+            @TestMetadata("inlineTwoReceivers.kt")
+            public void testInlineTwoReceivers() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/inlineTwoReceivers.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("member.kt")
             public void testMember() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/member.kt");
                 doTest(fileName);
             }
 
+            @TestMetadata("noinlineTwoReceivers.kt")
+            public void testNoinlineTwoReceivers() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/noinlineTwoReceivers.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("returnNoSuspend.kt")
             public void testReturnNoSuspend() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/returnNoSuspend.kt");

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

@@ -5161,12 +5161,24 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("inlineTwoReceivers.kt")
+            public void testInlineTwoReceivers() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/inlineTwoReceivers.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("member.kt")
             public void testMember() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/member.kt");
                 doTest(fileName);
             }
 
+            @TestMetadata("noinlineTwoReceivers.kt")
+            public void testNoinlineTwoReceivers() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/noinlineTwoReceivers.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("returnNoSuspend.kt")
             public void testReturnNoSuspend() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/returnNoSuspend.kt");

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

@@ -5161,12 +5161,24 @@ public class LightAnalysisModeCodegenTestGenerated extends AbstractLightAnalysis
                 doTest(fileName);
             }
 
+            @TestMetadata("inlineTwoReceivers.kt")
+            public void testInlineTwoReceivers() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/inlineTwoReceivers.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("member.kt")
             public void testMember() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/member.kt");
                 doTest(fileName);
             }
 
+            @TestMetadata("noinlineTwoReceivers.kt")
+            public void testNoinlineTwoReceivers() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/noinlineTwoReceivers.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("returnNoSuspend.kt")
             public void testReturnNoSuspend() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/returnNoSuspend.kt");

core/descriptors/src/org/jetbrains/kotlin/descriptors/impl/SimpleFunctionDescriptorImpl.java

@@ -23,7 +23,9 @@ import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.types.KotlinType;
 
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 
 public class SimpleFunctionDescriptorImpl extends FunctionDescriptorImpl implements SimpleFunctionDescriptor {
     protected SimpleFunctionDescriptorImpl(
@@ -58,9 +60,29 @@ public class SimpleFunctionDescriptorImpl extends FunctionDescriptorImpl impleme
             @Nullable KotlinType unsubstitutedReturnType,
             @Nullable Modality modality,
             @NotNull Visibility visibility
+    ) {
+        return initialize(receiverParameterType, dispatchReceiverParameter, typeParameters, unsubstitutedValueParameters,
+                          unsubstitutedReturnType, modality, visibility, null);
+    }
+
+    @NotNull
+    public SimpleFunctionDescriptorImpl initialize(
+            @Nullable KotlinType receiverParameterType,
+            @Nullable ReceiverParameterDescriptor dispatchReceiverParameter,
+            @NotNull List<? extends TypeParameterDescriptor> typeParameters,
+            @NotNull List<ValueParameterDescriptor> unsubstitutedValueParameters,
+            @Nullable KotlinType unsubstitutedReturnType,
+            @Nullable Modality modality,
+            @NotNull Visibility visibility,
+            @Nullable Map<? extends UserDataKey<?>, ?> userData
     ) {
         super.initialize(receiverParameterType, dispatchReceiverParameter, typeParameters, unsubstitutedValueParameters,
                          unsubstitutedReturnType, modality, visibility);
+
+        if (userData != null) {
+            userDataMap = new LinkedHashMap<UserDataKey<?>, Object>(userData);
+        }
+
         return this;
     }
 

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

@@ -5846,6 +5846,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("noinlineTwoReceivers.kt")
+            public void testNoinlineTwoReceivers() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/noinlineTwoReceivers.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("returnNoSuspend.kt")
             public void testReturnNoSuspend() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/returnNoSuspend.kt");