[FE] Allow lambda argument cast to contextual functional type

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt

@@ -106,6 +106,7 @@ class KotlinResolutionCallbacksImpl(
         lambdaArgument: LambdaKotlinCallArgument,
         isSuspend: Boolean,
         receiverType: UnwrappedType?,
+        contextReceiversTypes: List<UnwrappedType>,
         parameters: List<UnwrappedType>,
         expectedReturnType: UnwrappedType?,
         annotations: Annotations,
@@ -170,10 +171,12 @@ class KotlinResolutionCallbacksImpl(
         val refinedReceiverType = receiverType?.let {
             @OptIn(TypeRefinement::class) callComponents.kotlinTypeChecker.kotlinTypeRefiner.refineType(it)
         }
+        val refinedContextReceiverTypes = contextReceiversTypes.map {
+            @OptIn(TypeRefinement::class) callComponents.kotlinTypeChecker.kotlinTypeRefiner.refineType(it)
+        }
 
-        // TODO: Context receivers?
         val expectedType = createFunctionType(
-            builtIns, annotations, refinedReceiverType, emptyList(), parameters, null,
+            builtIns, annotations, refinedReceiverType, refinedContextReceiverTypes, parameters, null,
             lambdaInfo.expectedType, isSuspend
         )
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/NewCallArguments.kt

@@ -125,6 +125,7 @@ class FunctionExpressionImpl(
     val containingBlockForFunction: KtExpression,
     override val ktFunction: KtNamedFunction,
     override val receiverType: UnwrappedType?,
+    override val contextReceiversTypes: Array<UnwrappedType?>,
     override val parametersTypes: Array<UnwrappedType?>,
     override val returnType: UnwrappedType?
 ) : FunctionExpression, PSIFunctionKotlinCallArgument(outerCallContext, valueArgument, dataFlowInfoBeforeThisArgument, argumentName) {
@@ -250,13 +251,14 @@ fun processFunctionalExpression(
             // if function is a not anonymous function, resolve it as simple expression
             if (!postponedExpression.isFunctionalExpression()) return null
             val receiverType = resolveType(outerCallContext, postponedExpression.receiverTypeReference, typeResolver)
+            val contextReceiversTypes = resolveContextReceiversTypes(outerCallContext, postponedExpression, typeResolver)
             val parametersTypes = resolveParametersTypes(outerCallContext, postponedExpression, typeResolver) ?: emptyArray()
             val returnType = resolveType(outerCallContext, postponedExpression.typeReference, typeResolver)
                 ?: if (postponedExpression.hasBlockBody()) builtIns.unitType else null
 
             FunctionExpressionImpl(
                 outerCallContext, valueArgument, startDataFlowInfo, argumentName,
-                argumentExpression, postponedExpression, receiverType, parametersTypes, returnType
+                argumentExpression, postponedExpression, receiverType, contextReceiversTypes, parametersTypes, returnType
             )
         }
 
@@ -286,6 +288,18 @@ private fun resolveParametersTypes(
     }
 }
 
+private fun resolveContextReceiversTypes(
+    context: BasicCallResolutionContext,
+    ktFunction: KtFunction,
+    typeResolver: TypeResolver
+): Array<UnwrappedType?> {
+    val contextReceivers = ktFunction.contextReceivers
+
+    return Array(contextReceivers.size) {
+        contextReceivers[it]?.typeReference()?.let { typeRef -> resolveType(context, typeRef, typeResolver) }
+    }
+}
+
 @JvmName("resolveTypeWithGivenTypeReference")
 internal fun resolveType(
     context: BasicCallResolutionContext,