Support suspending extension functions

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

@@ -1869,21 +1869,30 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
     }
 
     @NotNull
-    public StackValue genCoroutineInstanceValueFromResolvedCall(ResolvedCall<?> resolvedCall) {
-        // Currently only handleResult/suspend members are supported
-        ReceiverValue dispatchReceiver = resolvedCall.getDispatchReceiver();
-        assert dispatchReceiver != null : "Dispatch receiver is null for handleResult/suspend to " + resolvedCall.getResultingDescriptor();
-        assert dispatchReceiver instanceof ExtensionReceiver
-                : "Argument for handleResult call to " + resolvedCall.getResultingDescriptor() +
-                  " should be a coroutine receiver parameter, but " + dispatchReceiver + " found";
+    public StackValue genCoroutineInstanceValueFromResolvedCall(@NotNull ResolvedCall<?> resolvedCall) {
+        ExtensionReceiver controllerReceiver = getControllerReceiverFromResolvedCall(resolvedCall);
         ClassDescriptor coroutineClassDescriptor =
-                bindingContext.get(CodegenBinding.CLASS_FOR_CALLABLE, ((ExtensionReceiver) dispatchReceiver).getDeclarationDescriptor());
+                bindingContext.get(CodegenBinding.CLASS_FOR_CALLABLE, controllerReceiver.getDeclarationDescriptor());
         assert coroutineClassDescriptor != null : "Coroutine class descriptor should not be null";
 
         // second argument for handleResult is always Continuation<T> ('this'-object in current implementation)
         return StackValue.thisOrOuter(this, coroutineClassDescriptor, false, false);
     }
 
+    private static ExtensionReceiver getControllerReceiverFromResolvedCall(@NotNull ResolvedCall<?> resolvedCall) {
+        ReceiverValue controllerReceiver =
+                resolvedCall.getDispatchReceiver() != null
+                ? resolvedCall.getDispatchReceiver()
+                : resolvedCall.getExtensionReceiver();
+
+        assert controllerReceiver != null : "Both dispatch and extension receivers are null for handleResult/suspend to " + resolvedCall.getResultingDescriptor();
+        assert controllerReceiver instanceof ExtensionReceiver
+                : "Argument for handleResult call to " + resolvedCall.getResultingDescriptor() +
+                  " should be a coroutine receiver parameter, but " + controllerReceiver + " found";
+
+        return (ExtensionReceiver) controllerReceiver;
+    }
+
     @NotNull
     private Type getVariableType(@NotNull VariableDescriptor variableDescriptor) {
         Type sharedVarType = typeMapper.getSharedVarType(variableDescriptor);

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/tower/ScopeTower.kt

@@ -69,6 +69,8 @@ interface CandidateWithBoundDispatchReceiver<out D : CallableDescriptor> {
     val diagnostics: List<ResolutionDiagnostic>
 
     val dispatchReceiver: ReceiverValue?
+
+    fun copy(newDescriptor: @UnsafeVariance D): CandidateWithBoundDispatchReceiver<D>
 }
 
 data class ResolutionCandidateStatus(val diagnostics: List<ResolutionDiagnostic>) {

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/tower/ScopeTowerProcessors.kt

@@ -21,6 +21,8 @@ import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
+import org.jetbrains.kotlin.resolve.coroutine.CoroutineReceiverValue
+import org.jetbrains.kotlin.resolve.coroutine.createCoroutineSuspensionFunctionView
 import org.jetbrains.kotlin.resolve.scopes.receivers.QualifierReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.Receiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
@@ -102,9 +104,23 @@ private class NoExplicitReceiverScopeTowerProcessor<D : CallableDescriptor, C: C
                     }
                 }
                 is TowerData.BothTowerLevelAndImplicitReceiver -> {
-                    data.level.collectCandidates(name, data.implicitReceiver).filter { it.requiresExtensionReceiver }.map {
-                        context.createCandidate(it, ExplicitReceiverKind.NO_EXPLICIT_RECEIVER, extensionReceiver = data.implicitReceiver)
+                    val result = mutableListOf<C>()
+
+                    data.level.collectCandidates(name, data.implicitReceiver).filter { it.requiresExtensionReceiver }.forEach {
+                        result.add(
+                            context.createCandidate(
+                                    it, ExplicitReceiverKind.NO_EXPLICIT_RECEIVER, extensionReceiver = data.implicitReceiver))
+
+                        if (data.implicitReceiver is CoroutineReceiverValue) {
+                            val newDescriptor = it.descriptor.createCoroutineSuspensionFunctionView() ?: return@forEach
+                            result.add(
+                                context.createCandidate(
+                                        it.copy(newDescriptor = newDescriptor),
+                                        ExplicitReceiverKind.NO_EXPLICIT_RECEIVER, extensionReceiver = data.implicitReceiver))
+                        }
                     }
+
+                    result
                 }
                 else -> emptyList()
             }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/tower/TowerUtils.kt

@@ -35,4 +35,7 @@ internal class CandidateWithBoundDispatchReceiverImpl<out D : CallableDescriptor
         override val dispatchReceiver: ReceiverValue?,
         override val descriptor: D,
         override val diagnostics: List<ResolutionDiagnostic>
-) : CandidateWithBoundDispatchReceiver<D>
+) : CandidateWithBoundDispatchReceiver<D> {
+    override fun copy(newDescriptor: @UnsafeVariance D) =
+            CandidateWithBoundDispatchReceiverImpl(dispatchReceiver, newDescriptor, diagnostics)
+}

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

@@ -0,0 +1,36 @@
+@AllowSuspendExtensions
+class Controller {
+    suspend fun String.suspendHere(x: Continuation<String>) {
+        x.resume(this)
+    }
+
+    inline suspend fun String.inlineSuspendHere(x: Continuation<String>) {
+        suspendHere(x)
+    }
+}
+
+suspend fun Controller.suspendExtension(v: String, x: Continuation<String>) {
+    v.suspendHere(x)
+}
+
+inline suspend fun Controller.inlineSuspendExtension(v: String, x: Continuation<String>) {
+    v.inlineSuspendHere(x)
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>) {
+    c(Controller()).resume(Unit)
+}
+
+fun box(): String {
+    var result = ""
+
+    builder {
+        if ("56".suspendHere() != "56") throw java.lang.RuntimeException("fail 1")
+        if ("28".inlineSuspendHere() != "28") throw java.lang.RuntimeException("fail 2")
+
+        if (suspendExtension("123") != "123") throw java.lang.RuntimeException("fail 3")
+        result = inlineSuspendExtension("OK")
+    }
+
+    return result
+}

compiler/testData/diagnostics/tests/coroutines/suspendExtensionFunctions.kt

@@ -0,0 +1,60 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// !CHECK_TYPE
+@AllowSuspendExtensions
+class Controller
+
+suspend fun Controller.noParams(c: Continuation<Unit>) {
+
+}
+suspend fun Controller.yieldString(value: String, c: Continuation<Unit>) {
+}
+
+suspend fun <V> Controller.await(f: () -> V, machine: Continuation<V>) {
+}
+
+suspend fun <V> Controller.await(f: Int, machine: Continuation<V>) {
+}
+
+suspend fun Controller.severalParams(x: String, y: Int, machine: Continuation<Double>) {
+}
+
+// These two must be prohibited because String and Any are not properly annotated
+<!INAPPLICABLE_MODIFIER!>suspend<!> fun String.wrongReceiver(y: Int, machine: Continuation<Double>) {
+}
+
+<!INAPPLICABLE_MODIFIER!>suspend<!> fun Any.anyReceiver(y: Int, machine: Continuation<Double>) {
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>) {}
+
+fun test() {
+    builder {
+        noParams()
+        yieldString("abc") checkType { _<Unit>() }
+        yieldString(<!CONSTANT_EXPECTED_TYPE_MISMATCH!>1<!>) checkType { _<Unit>() }
+
+        await<String> { "123" } checkType { _<String>() }
+
+        // Inference from lambda return type
+        await { 123 } checkType { _<Int>() }
+
+        // Inference from expected type
+        checkSubtype<String>(await(567))
+
+        await<Double>(123) checkType { _<Double>() }
+
+        severalParams("", 89) checkType { _<Double>() }
+        // TODO: prohibit such calls
+        severalParams("", 89, <!CONSTANT_EXPECTED_TYPE_MISMATCH!>6.9<!>) checkType { _<Unit>() }
+        severalParams("", 89, this <!CAST_NEVER_SUCCEEDS!>as<!> Continuation<Double>) checkType { _<Unit>() }
+
+        <!UNRESOLVED_REFERENCE_WRONG_RECEIVER!>wrongReceiver<!>(1)
+
+        with("") {
+            wrongReceiver(2<!NO_VALUE_FOR_PARAMETER!>)<!>
+        }
+
+        // Though such calls are allowed declarations with not-annotated receiver should be prohibited
+        anyReceiver(3)
+    }
+}

compiler/testData/diagnostics/tests/coroutines/suspendExtensionFunctions.txt

@@ -0,0 +1,18 @@
+package
+
+public fun builder(/*0*/ coroutine c: Controller.() -> kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Unit
+public fun test(): kotlin.Unit
+public suspend fun kotlin.Any.anyReceiver(/*0*/ y: kotlin.Int, /*1*/ machine: kotlin.coroutines.Continuation<kotlin.Double>): kotlin.Unit
+public suspend fun </*0*/ V> Controller.await(/*0*/ f: () -> V, /*1*/ machine: kotlin.coroutines.Continuation<V>): kotlin.Unit
+public suspend fun </*0*/ V> Controller.await(/*0*/ f: kotlin.Int, /*1*/ machine: kotlin.coroutines.Continuation<V>): kotlin.Unit
+public suspend fun Controller.noParams(/*0*/ c: kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Unit
+public suspend fun Controller.severalParams(/*0*/ x: kotlin.String, /*1*/ y: kotlin.Int, /*2*/ machine: kotlin.coroutines.Continuation<kotlin.Double>): kotlin.Unit
+public suspend fun kotlin.String.wrongReceiver(/*0*/ y: kotlin.Int, /*1*/ machine: kotlin.coroutines.Continuation<kotlin.Double>): kotlin.Unit
+public suspend fun Controller.yieldString(/*0*/ value: kotlin.String, /*1*/ c: kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Unit
+
+@kotlin.coroutines.AllowSuspendExtensions() public final class Controller {
+    public constructor Controller()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -3996,6 +3996,12 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("suspendExtensionFunctions.kt")
+            public void testSuspendExtensionFunctions() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/coroutines/suspendExtensionFunctions.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("suspendFunctions.kt")
             public void testSuspendFunctions() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/coroutines/suspendFunctions.kt");

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

@@ -4327,6 +4327,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("suspendExtension.kt")
+        public void testSuspendExtension() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendExtension.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("suspendFromInlineLambda.kt")
         public void testSuspendFromInlineLambda() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/suspendFromInlineLambda.kt");