Replace expected type for coroutine-lambda

Use 'handleResult' method of controller to determine what lambda should return

If original declaration was
fun builder(coroutine c: Controller.() -> Continuation<Unit>) = 1
and there is 'fun handleResult(x: X, c: Continuation<Nothing>)',
then expected type for lambda is 'Controller.() -> X'

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

@@ -0,0 +1,41 @@
+/*
+ * 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.coroutines
+
+import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
+import org.jetbrains.kotlin.types.KotlinType
+
+@JvmField
+val CONTINUATION_INTERFACE_FQ_NAME = FqName("kotlin.coroutines.Continuation")
+
+@JvmField
+val HANDLE_RESULT_NAME = Name.identifier("handleResult")
+
+/**
+ * @returns type of first value parameter if function is 'operator handleResult' in coroutine controller
+ */
+fun SimpleFunctionDescriptor.getExpectedTypeForCoroutineControllerHandleResult(): KotlinType? {
+    if (name != HANDLE_RESULT_NAME) return null
+    if (valueParameters.size != 2) return null
+
+    if (valueParameters[1].type.constructor.declarationDescriptor?.fqNameUnsafe != CONTINUATION_INTERFACE_FQ_NAME.toUnsafe()) return null
+
+    return valueParameters[0].type
+}

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

@@ -19,7 +19,11 @@ package org.jetbrains.kotlin.resolve.calls.callResolverUtil
 import com.google.common.collect.Lists
 import com.intellij.util.containers.ContainerUtil
 import org.jetbrains.kotlin.builtins.ReflectionTypes
+import org.jetbrains.kotlin.builtins.getReceiverTypeFromFunctionType
+import org.jetbrains.kotlin.builtins.isExtensionFunctionType
+import org.jetbrains.kotlin.coroutines.getExpectedTypeForCoroutineControllerHandleResult
 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.calls.CallTransformer
@@ -28,6 +32,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.constraintPosition.Constrain
 import org.jetbrains.kotlin.resolve.calls.inference.getNestedTypeVariables
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.calls.tasks.ResolutionCandidate
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
@@ -36,6 +41,7 @@ import org.jetbrains.kotlin.resolve.validation.InfixValidator
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.TypeUtils.DONT_CARE
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
+import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
 import org.jetbrains.kotlin.types.typeUtil.contains
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
@@ -165,6 +171,28 @@ fun getEffectiveExpectedType(parameterDescriptor: ValueParameterDescriptor, argu
         return varargElementType
     }
 
+    if (parameterDescriptor.isCoroutine &&
+            argument.getArgumentExpression() is KtLambdaExpression &&
+            parameterDescriptor.type.isExtensionFunctionType
+    ) {
+        val receiverType = getReceiverTypeFromFunctionType(parameterDescriptor.type)!!
+
+        val newExpectedLambdaReturnType =
+                receiverType.memberScope
+                        .getContributedFunctions(HANDLE_RESULT_NAME, KotlinLookupLocation(argument.asElement())).mapNotNull {
+                                it.getExpectedTypeForCoroutineControllerHandleResult()
+                        }.singleOrNull()
+                // If no handleResult function found, then expected return type for lambda is Unit
+                ?: parameterDescriptor.module.builtIns.unitType
+
+        // replace return type for lambda with the one we got from single 'handleResult'
+        val newFunctionTypeArguments = parameterDescriptor.type.arguments.toMutableList()
+        newFunctionTypeArguments[newFunctionTypeArguments.lastIndex] = newExpectedLambdaReturnType.asTypeProjection()
+
+        return parameterDescriptor.type.replace(
+                newArguments = newFunctionTypeArguments)
+    }
+
     return parameterDescriptor.type
 }
 
@@ -198,4 +226,4 @@ fun createResolutionCandidatesForConstructors(
     }
 
     return constructors.map { ResolutionCandidate.create(call, it, dispatchReceiver, receiverKind, knownSubstitutor) }
-}
+}

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

@@ -0,0 +1,64 @@
+// !CHECK_TYPE
+// !DIAGNOSTICS: -UNUSED_PARAMETER -ASSIGNED_BUT_NEVER_ACCESSED_VARIABLE -UNUSED_VALUE -UNUSED_VARIABLE
+class IntController {
+    fun handleResult(x: Int, c: Continuation<Nothing>) { }
+}
+
+class GenericController<T> {
+    fun handleResult(x: T, c: Continuation<Nothing>) { }
+}
+
+class UnitController {
+    fun handleResult(x: Unit, c: Continuation<Nothing>) { }
+}
+
+class EmptyController
+
+fun builder(coroutine c: IntController.() -> Continuation<Unit>) = 1
+fun <T> genericBuilder(coroutine c: GenericController<T>.() -> Continuation<Unit>): T = null!!
+fun unitBuilder(coroutine c: UnitController.() -> Continuation<Unit>) = 1
+fun emptyBuilder(coroutine c: EmptyController.() -> Continuation<Unit>) = 1
+
+fun <T> manyArgumentsBuilder(
+        coroutine c1: UnitController.() -> Continuation<Unit>,
+        coroutine c2: GenericController<T>.() -> Continuation<Unit>,
+        coroutine c3: IntController.() -> Continuation<Int>
+):T = null!!
+
+fun severalParamsInLambda(coroutine c: UnitController.(String, Int) -> Continuation<Unit>) {}
+
+fun foo() {
+    builder({ 1 })
+    builder { 1 }
+
+    val x = { 1 }
+    builder(<!TYPE_MISMATCH!>x<!>)
+    builder(<!UNCHECKED_CAST!>{1} as (IntController.() -> Continuation<Unit>)<!>)
+
+    var i: Int = 1
+    i = genericBuilder({ 1 })
+    i = genericBuilder { 1 }
+    <!TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>genericBuilder<!> { 1 }
+    genericBuilder<Int> { 1 }
+    genericBuilder<Int> { <!TYPE_MISMATCH!>""<!> }
+
+    val y = { 1 }
+    <!TYPE_INFERENCE_PARAMETER_CONSTRAINT_ERROR!>genericBuilder<!>(<!TYPE_MISMATCH!>y<!>)
+
+    unitBuilder {}
+    unitBuilder { <!UNUSED_EXPRESSION!>1<!> }
+    unitBuilder({})
+    unitBuilder({ <!UNUSED_EXPRESSION!>1<!> })
+
+    <!TYPE_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>manyArgumentsBuilder<!>({}, { "" }) { 1 }
+
+    val s: String = manyArgumentsBuilder({}, { "" }) { 1 }
+
+    manyArgumentsBuilder<String>({}, { "" }, { 1 })
+    manyArgumentsBuilder<String>({}, { <!CONSTANT_EXPECTED_TYPE_MISMATCH!>1<!> }, { 2 })
+
+    severalParamsInLambda { x, y ->
+        x checkType { _<String>() }
+        y checkType { _<Int>() }
+    }
+}

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

@@ -0,0 +1,40 @@
+package
+
+public fun builder(/*0*/ coroutine c: IntController.() -> kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Int
+public fun emptyBuilder(/*0*/ coroutine c: EmptyController.() -> kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Int
+public fun foo(): kotlin.Unit
+public fun </*0*/ T> genericBuilder(/*0*/ coroutine c: GenericController<T>.() -> kotlin.coroutines.Continuation<kotlin.Unit>): T
+public fun </*0*/ T> manyArgumentsBuilder(/*0*/ coroutine c1: UnitController.() -> kotlin.coroutines.Continuation<kotlin.Unit>, /*1*/ coroutine c2: GenericController<T>.() -> kotlin.coroutines.Continuation<kotlin.Unit>, /*2*/ coroutine c3: IntController.() -> kotlin.coroutines.Continuation<kotlin.Int>): T
+public fun severalParamsInLambda(/*0*/ coroutine c: UnitController.(kotlin.String, kotlin.Int) -> kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Unit
+public fun unitBuilder(/*0*/ coroutine c: UnitController.() -> kotlin.coroutines.Continuation<kotlin.Unit>): kotlin.Int
+
+public final class EmptyController {
+    public constructor EmptyController()
+    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
+}
+
+public final class GenericController</*0*/ T> {
+    public constructor GenericController</*0*/ T>()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public final fun handleResult(/*0*/ x: T, /*1*/ c: kotlin.coroutines.Continuation<kotlin.Nothing>): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class IntController {
+    public constructor IntController()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public final fun handleResult(/*0*/ x: kotlin.Int, /*1*/ c: kotlin.coroutines.Continuation<kotlin.Nothing>): kotlin.Unit
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class UnitController {
+    public constructor UnitController()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public final fun handleResult(/*0*/ x: kotlin.Unit, /*1*/ c: kotlin.coroutines.Continuation<kotlin.Nothing>): kotlin.Unit
+    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

@@ -3865,6 +3865,21 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
             }
         }
 
+        @TestMetadata("compiler/testData/diagnostics/tests/coroutines")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class Coroutines extends AbstractDiagnosticsTest {
+            public void testAllFilesPresentInCoroutines() throws Exception {
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/diagnostics/tests/coroutines"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+
+            @TestMetadata("lambdaExpectedType.kt")
+            public void testLambdaExpectedType() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/coroutines/lambdaExpectedType.kt");
+                doTest(fileName);
+            }
+        }
+
         @TestMetadata("compiler/testData/diagnostics/tests/cyclicHierarchy")
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)