[FE] Allow lambda argument cast to contextual functional type

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/ExternalComponents.kt

@@ -76,6 +76,7 @@ interface KotlinResolutionCallbacks {
         lambdaArgument: LambdaKotlinCallArgument,
         isSuspend: Boolean,
         receiverType: UnwrappedType?,
+        contextReceiversTypes: List<UnwrappedType>,
         parameters: List<UnwrappedType>,
         expectedReturnType: UnwrappedType?, // null means, that return type is not proper i.e. it depends on some type variables
         annotations: Annotations,

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponeArgumentsChecks.kt

@@ -86,7 +86,7 @@ private fun preprocessLambdaArgument(
 
     if (expectedType != null) {
         val lambdaType = createFunctionType(
-            csBuilder.builtIns, Annotations.EMPTY, resolvedArgument.receiver, emptyList(), // TODO: Context receivers?
+            csBuilder.builtIns, Annotations.EMPTY, resolvedArgument.receiver, resolvedArgument.contextReceivers,
             resolvedArgument.parameters, null, resolvedArgument.returnType, resolvedArgument.isSuspend
         )
         csBuilder.addSubtypeConstraint(lambdaType, expectedType, ArgumentConstraintPositionImpl(argument))
@@ -114,6 +114,14 @@ private fun extraLambdaInfo(
         argumentAsFunctionExpression?.returnType ?: expectedType?.arguments?.singleOrNull()?.type?.unwrap()?.takeIf { isFunctionSupertype }
         ?: typeVariable.defaultType
 
+    val contextReceiversTypes = argumentAsFunctionExpression?.contextReceiversTypes?.mapIndexed { index, contextReceiverType ->
+        if (contextReceiverType != null) {
+            contextReceiverType
+        } else {
+            diagnosticsHolder.addDiagnostic(NotEnoughInformationForLambdaParameter(argument, index))
+            ErrorUtils.createErrorType("<Unknown lambda context receiver type>")
+        }
+    } ?: emptyList()
     val parameters = argument.parametersTypes?.mapIndexed { index, parameterType ->
         if (parameterType != null) {
             parameterType
@@ -130,6 +138,7 @@ private fun extraLambdaInfo(
         argument,
         isSuspend,
         receiverType,
+        contextReceiversTypes,
         parameters,
         returnType,
         typeVariable.takeIf { newTypeVariableUsed },
@@ -146,6 +155,7 @@ private fun extractLambdaInfoFromFunctionalType(
     val parametersTypes = argument.parametersTypes
     val expectedParameters = expectedType.getValueParameterTypesFromFunctionType()
     val expectedReceiver = expectedType.getReceiverTypeFromFunctionType()?.unwrap()
+    val expectedContextReceivers = expectedType.getContextReceiverTypesFromFunctionType().map { it.unwrap() }.toTypedArray()
     val argumentAsFunctionExpression = argument.safeAs<FunctionExpression>()
 
     val receiverFromExpected = argumentAsFunctionExpression?.receiverType == null && expectedReceiver != null
@@ -186,6 +196,7 @@ private fun extractLambdaInfoFromFunctionalType(
                 type.orExpected(index)
             } ?: expectedParameters.map { it.type.unwrap() }) to (if (receiverFromExpected) expectedReceiver else null)
     }
+    val contextReceivers = (argumentAsFunctionExpression?.contextReceiversTypes ?: expectedContextReceivers).filterNotNull()
 
     val returnType = argumentAsFunctionExpression?.returnType ?: expectedType.getReturnTypeFromFunctionType().unwrap()
 
@@ -193,6 +204,7 @@ private fun extractLambdaInfoFromFunctionalType(
         argument,
         expectedType.isSuspendFunctionType,
         receiver,
+        contextReceivers,
         parameters,
         returnType,
         typeVariableForLambdaReturnType = returnTypeVariable,

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponedArgumentsAnalyzer.kt

@@ -5,10 +5,7 @@
 
 package org.jetbrains.kotlin.resolve.calls.components
 
-import org.jetbrains.kotlin.builtins.StandardNames
-import org.jetbrains.kotlin.builtins.getReceiverTypeFromFunctionType
-import org.jetbrains.kotlin.builtins.getValueParameterTypesFromFunctionType
-import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
+import org.jetbrains.kotlin.builtins.*
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
@@ -94,10 +91,14 @@ class PostponedArgumentsAnalyzer(
 
         val expectedParameters = lambda.expectedType.valueParameters()
         val expectedReceiver = lambda.expectedType.receiver()
+        val expectedContextReceivers = lambda.expectedType.contextReceivers()
 
         val receiver = lambda.receiver?.let {
             expectedOrActualType(expectedReceiver ?: expectedParameters?.getOrNull(0), lambda.receiver)
         }
+        val contextReceivers = lambda.contextReceivers.mapIndexedNotNull { i, contextReceiver ->
+            expectedOrActualType(expectedContextReceivers?.getOrNull(i), contextReceiver)
+        }
 
         val expectedParametersToMatchAgainst = when {
             receiver == null && expectedReceiver != null && expectedParameters != null -> listOf(expectedReceiver) + expectedParameters
@@ -131,6 +132,7 @@ class PostponedArgumentsAnalyzer(
             lambda.atom,
             lambda.isSuspend,
             receiver,
+            contextReceivers,
             parameters,
             expectedTypeForReturnArguments,
             convertedAnnotations ?: Annotations.EMPTY,
@@ -220,6 +222,10 @@ class PostponedArgumentsAnalyzer(
         return forFunctionalType { getReceiverTypeFromFunctionType()?.unwrap() }
     }
 
+    private fun UnwrappedType?.contextReceivers(): List<UnwrappedType>? {
+        return forFunctionalType { getContextReceiverTypesFromFunctionType().map { it.unwrap() } }
+    }
+
     private fun UnwrappedType?.valueParameters(): List<UnwrappedType>? {
         return forFunctionalType { getValueParameterTypesFromFunctionType().map { it.type.unwrap() } }
     }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/KotlinCallArguments.kt

@@ -73,6 +73,8 @@ interface FunctionExpression : LambdaKotlinCallArgument {
     // null means that there function can not have receiver
     val receiverType: UnwrappedType?
 
+    val contextReceiversTypes: Array<UnwrappedType?>
+
     // null means that return type is not declared, for fun(){ ... } returnType == Unit
     val returnType: UnwrappedType?
 }