Replace common extension receiver of coroutine-lambda with magic one

Currently it only contains additional synthetic declarations for suspend-functions

E.g. `fun <V> await(f: CompletableFuture<V>): V` for `fun <V> await(f: CompletableFuture<V>, machine: Continuation<V>): Unit`

compiler/frontend/src/org/jetbrains/kotlin/coroutines/coroutineUtil.kt

@@ -16,7 +16,9 @@
 
 package org.jetbrains.kotlin.coroutines
 
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
@@ -28,3 +30,10 @@ fun SimpleFunctionDescriptor.getExpectedTypeForCoroutineControllerHandleResult()
 
     return valueParameters.getOrNull(0)?.type
 }
+
+val CallableDescriptor.controllerTypeIfCoroutine: KotlinType?
+    get() {
+        if (this !is AnonymousFunctionDescriptor || !this.isCoroutine) return null
+
+        return this.extensionReceiverParameter?.returnType
+    }

compiler/frontend/src/org/jetbrains/kotlin/descriptors/impl/AnonymousFunctionDescriptor.java

@@ -23,12 +23,20 @@ import org.jetbrains.kotlin.descriptors.annotations.Annotations;
 import org.jetbrains.kotlin.name.Name;
 
 public class AnonymousFunctionDescriptor extends SimpleFunctionDescriptorImpl {
+    private final boolean isCoroutine;
+
     public AnonymousFunctionDescriptor(
             @NotNull DeclarationDescriptor containingDeclaration,
             @NotNull Annotations annotations,
             @NotNull Kind kind,
-            @NotNull SourceElement source
+            @NotNull SourceElement source,
+            boolean isCoroutine
     ) {
         super(containingDeclaration, null, annotations, Name.special("<anonymous>"), kind, source);
+        this.isCoroutine = isCoroutine;
+    }
+
+    public boolean isCoroutine() {
+        return isCoroutine;
     }
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/FunctionDescriptorUtil.java

@@ -20,7 +20,11 @@ import kotlin.Unit;
 import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.descriptors.*;
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor;
+import org.jetbrains.kotlin.descriptors.impl.ReceiverParameterDescriptorImpl;
+import org.jetbrains.kotlin.resolve.coroutine.CoroutineReceiverValue;
 import org.jetbrains.kotlin.resolve.scopes.*;
+import org.jetbrains.kotlin.resolve.scopes.receivers.ExtensionReceiver;
 import org.jetbrains.kotlin.types.*;
 import org.jetbrains.kotlin.types.typeUtil.TypeUtilsKt;
 
@@ -66,6 +70,17 @@ public class FunctionDescriptorUtil {
             @NotNull LocalRedeclarationChecker redeclarationChecker
     ) {
         ReceiverParameterDescriptor receiver = descriptor.getExtensionReceiverParameter();
+
+        if (descriptor instanceof AnonymousFunctionDescriptor
+                && (((AnonymousFunctionDescriptor) descriptor).isCoroutine())
+                && receiver != null && receiver.getValue() instanceof ExtensionReceiver) {
+            receiver =
+                    new ReceiverParameterDescriptorImpl(
+                            descriptor,
+                            new CoroutineReceiverValue(
+                                    ((ExtensionReceiver) receiver.getValue()).getDeclarationDescriptor(), receiver.getValue().getType()));
+        }
+
         return new LexicalScopeImpl(outerScope, descriptor, true, receiver, LexicalScopeKind.FUNCTION_INNER_SCOPE, redeclarationChecker,
                                     new Function1<LexicalScopeImpl.InitializeHandler, Unit>() {
                                         @Override

compiler/frontend/src/org/jetbrains/kotlin/resolve/callableReferences/CallableReferencesResolutionUtils.kt

@@ -153,7 +153,7 @@ private fun bindFunctionReference(expression: KtCallableReferenceExpression, typ
             context.scope.ownerDescriptor,
             Annotations.EMPTY,
             CallableMemberDescriptor.Kind.DECLARATION,
-            expression.toSourceElement()
+            expression.toSourceElement(), /* isCoroutine = */ false
     )
 
     functionDescriptor.initialize(

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CallResolverUtil.kt

@@ -26,10 +26,14 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.incremental.KotlinLookupLocation
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.CallTransformer
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.calls.callUtil.getValueArgumentForExpression
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.constraintPosition.ConstraintPositionKind.EXPECTED_TYPE_POSITION
 import org.jetbrains.kotlin.resolve.calls.inference.getNestedTypeVariables
+import org.jetbrains.kotlin.resolve.calls.model.ArgumentMatch
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.calls.tasks.ResolutionCandidate
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
@@ -229,3 +233,15 @@ fun createResolutionCandidatesForConstructors(
 
     return constructors.map { ResolutionCandidate.create(call, it, dispatchReceiver, receiverKind, knownSubstitutor) }
 }
+
+fun KtLambdaExpression.getCorrespondingParameterForFunctionArgument(
+        bindingContext: BindingContext
+): ValueParameterDescriptor? {
+    val resolvedCall = KtPsiUtil.getParentCallIfPresent(this)?.getResolvedCall(bindingContext) ?: return null
+    val valueArgument =
+            resolvedCall.call.getValueArgumentForExpression(this)
+            ?: return null
+    val mapping = resolvedCall.getArgumentMapping(valueArgument) as? ArgumentMatch ?: return null
+
+    return mapping.valueParameter
+}

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/FunctionsTypingVisitor.kt

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.psi.psiUtil.checkReservedPrefixWord
 import org.jetbrains.kotlin.psi.psiUtil.getAnnotationEntries
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.BindingContext.EXPECTED_RETURN_TYPE
+import org.jetbrains.kotlin.resolve.calls.callResolverUtil.getCorrespondingParameterForFunctionArgument
 import org.jetbrains.kotlin.resolve.lazy.ForceResolveUtil
 import org.jetbrains.kotlin.resolve.scopes.LexicalWritableScope
 import org.jetbrains.kotlin.resolve.source.toSourceElement
@@ -170,7 +171,8 @@ internal class FunctionsTypingVisitor(facade: ExpressionTypingInternals) : Expre
         val functionDescriptor = AnonymousFunctionDescriptor(
             context.scope.ownerDescriptor,
             components.annotationResolver.resolveAnnotationsWithArguments(context.scope, expression.getAnnotationEntries(), context.trace),
-            CallableMemberDescriptor.Kind.DECLARATION, functionLiteral.toSourceElement()
+            CallableMemberDescriptor.Kind.DECLARATION, functionLiteral.toSourceElement(),
+            expression.getCorrespondingParameterForFunctionArgument(context.trace.bindingContext)?.isCoroutine ?: false
         )
         components.functionDescriptorResolver.
                 initializeFunctionDescriptorAndExplicitReturnType(context.scope.ownerDescriptor, context.scope, functionLiteral,

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

@@ -22,6 +22,8 @@ import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.smartcasts.getReceiverValueWithSmartCast
 import org.jetbrains.kotlin.resolve.calls.util.FakeCallableDescriptorForObject
+import org.jetbrains.kotlin.resolve.coroutine.CoroutineReceiverValue
+import org.jetbrains.kotlin.resolve.coroutine.createCoroutineSuspensionFunctionView
 import org.jetbrains.kotlin.resolve.descriptorUtil.HIDES_MEMBERS_NAME_LIST
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasClassValueDescriptor
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasHidesMembersAnnotation
@@ -116,6 +118,16 @@ internal class ReceiverScopeTowerLevel(
             }
         }
 
+        if (dispatchReceiver is CoroutineReceiverValue) {
+            result.addAll(result.mapNotNull {
+                val realDescriptor = it.descriptor
+                val suspensionFunctionView =
+                        (realDescriptor as? SimpleFunctionDescriptor)?.createCoroutineSuspensionFunctionView() ?: return@mapNotNull null
+                @Suppress("UNCHECKED_CAST")
+                createCandidateDescriptor(suspensionFunctionView as D, dispatchReceiver)
+            })
+        }
+
         return result
     }
 

compiler/resolution/src/org/jetbrains/kotlin/resolve/coroutine/coroutineUtil.kt

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.resolve.coroutine
+
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.resolve.scopes.receivers.ExtensionReceiver
+import org.jetbrains.kotlin.types.KotlinType
+
+// Returns suspension function as it's visible within coroutines:
+// E.g. `fun <V> await(f: CompletableFuture<V>): V` instead of `fun <V> await(f: CompletableFuture<V>, machine: Continuation<V>): Unit`
+fun SimpleFunctionDescriptor.createCoroutineSuspensionFunctionView(): SimpleFunctionDescriptor? {
+    if (!isSuspend) return null
+    val returnType = valueParameters.lastOrNull()?.returnType?.arguments?.getOrNull(0)?.type ?: return null
+
+    val newOriginal =
+            if (original !== this)
+                original.createCoroutineSuspensionFunctionView()
+            else null
+
+    return newCopyBuilder().apply {
+        setReturnType(returnType)
+        setOriginal(newOriginal)
+        setValueParameters(valueParameters.subList(0, valueParameters.size - 1))
+    }.build()!!
+}
+
+class CoroutineReceiverValue(callableDescriptor: CallableDescriptor, receiverType: KotlinType) : ExtensionReceiver(callableDescriptor, receiverType)

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

@@ -0,0 +1,43 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER
+// !CHECK_TYPE
+class Controller {
+    suspend fun noParams(c: Continuation<Unit>) {
+
+    }
+    suspend fun yieldString(value: String, c: Continuation<Unit>) {
+    }
+
+    suspend fun <V> await(f: () -> V, machine: Continuation<V>) {
+    }
+
+    suspend fun <V> await(f: Int, machine: Continuation<V>) {
+    }
+
+    suspend fun severalParams(x: String, 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>() }
+    }
+}

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

@@ -0,0 +1,16 @@
+package
+
+public fun builder(/*0*/ coroutine c: Controller.() -> kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Unit
+public fun test(): kotlin.Unit
+
+public final class Controller {
+    public constructor Controller()
+    public final suspend fun </*0*/ V> await(/*0*/ f: () -> V, /*1*/ machine: kotlin.coroutines.Continuation<V>): kotlin.Unit
+    public final suspend fun </*0*/ V> await(/*0*/ f: kotlin.Int, /*1*/ machine: kotlin.coroutines.Continuation<V>): kotlin.Unit
+    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 final suspend fun noParams(/*0*/ c: kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Unit
+    public final suspend fun severalParams(/*0*/ x: kotlin.String, /*1*/ y: kotlin.Int, /*2*/ machine: kotlin.coroutines.Continuation<kotlin.Double>): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+    public final suspend fun yieldString(/*0*/ value: kotlin.String, /*1*/ c: kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Unit
+}

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

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

core/descriptors/src/org/jetbrains/kotlin/descriptors/FunctionDescriptor.java

@@ -120,7 +120,7 @@ public interface FunctionDescriptor extends CallableMemberDescriptor {
         CopyBuilder<D> setDispatchReceiverParameter(@Nullable ReceiverParameterDescriptor dispatchReceiverParameter);
 
         @NotNull
-        CopyBuilder<D> setOriginal(@NotNull FunctionDescriptor original);
+        CopyBuilder<D> setOriginal(@Nullable FunctionDescriptor original);
 
         @NotNull
         CopyBuilder<D> setSignatureChange();

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

@@ -445,7 +445,7 @@ public abstract class FunctionDescriptorImpl extends DeclarationDescriptorNonRoo
 
         @Override
         @NotNull
-        public CopyConfiguration setOriginal(@NotNull FunctionDescriptor original) {
+        public CopyConfiguration setOriginal(@Nullable FunctionDescriptor original) {
             this.original = original;
             return this;
         }

core/descriptors/src/org/jetbrains/kotlin/types/error/ErrorSimpleFunctionDescriptorImpl.java

@@ -137,7 +137,7 @@ public class ErrorSimpleFunctionDescriptorImpl extends SimpleFunctionDescriptorI
 
             @NotNull
             @Override
-            public CopyBuilder<SimpleFunctionDescriptor> setOriginal(@NotNull FunctionDescriptor original) {
+            public CopyBuilder<SimpleFunctionDescriptor> setOriginal(@Nullable FunctionDescriptor original) {
                 return this;
             }