[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,

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.*
-import org.jetbrains.kotlin.builtins.getReceiverTypeFromFunctionType
-import org.jetbrains.kotlin.builtins.getValueParameterTypesFromFunctionType
-import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
 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?
 }

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/autoCloseScope.kt

@@ -17,7 +17,7 @@ fun File.open(): InputStream = TODO()
 fun withAutoClose(block: context(AutoCloseScope) () -> Unit) {
     val scope = AutoCloseScopeImpl() // Not shown here
     try {
-        with(scope) { block() }
+        with(scope) { block(<!NO_VALUE_FOR_PARAMETER!>)<!> }
     } finally {
         scope.close()
     }

compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.fir.kt

@@ -45,9 +45,26 @@ fun test() {
     with(C()) {
         with(R()) {
             f1(lf1)
+            f1 { _ ->
+                r
+                c
+            }
+
             f2(lf2)
+            f2 { _ ->
+                c
+            }
+
             f3(lf3)
+            f3 {
+                r
+                c
+            }
+
             f4(lf4)
+            f4 {
+                c
+            }
         }
     }
 }

compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.kt

@@ -45,9 +45,26 @@ fun test() {
     with(C()) {
         with(R()) {
             f1(lf1)
+            f1 { _ ->
+                r
+                c
+            }
+
             f2(lf2)
+            f2 { _ ->
+                c
+            }
+
             f3(lf3)
+            f3 {
+                r
+                c
+            }
+
             f4(lf4)
+            f4 {
+                c
+            }
         }
     }
 }