diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt index 1c2a5ff94a1..248f13793ed 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt @@ -27,6 +27,7 @@ fun resolveArgumentExpression( csBuilder: ConstraintSystemBuilder, argument: FirExpression, expectedType: ConeKotlinType, + expectedTypeRef: FirTypeRef, sink: CheckerSink, isReceiver: Boolean, acceptLambdaAtoms: (PostponedResolvedAtomMarker) -> Unit, @@ -42,7 +43,7 @@ fun resolveArgumentExpression( typeProvider ) // TODO:! - is FirAnonymousFunction -> preprocessLambdaArgument(csBuilder, argument, expectedType, acceptLambdaAtoms) + is FirAnonymousFunction -> preprocessLambdaArgument(csBuilder, argument, expectedType, expectedTypeRef, acceptLambdaAtoms) // TODO:! is FirCallableReferenceAccess -> Unit // TODO:! @@ -51,6 +52,7 @@ fun resolveArgumentExpression( csBuilder, argument.expression, expectedType, + expectedTypeRef, sink, isReceiver, acceptLambdaAtoms, @@ -60,6 +62,7 @@ fun resolveArgumentExpression( csBuilder, argument.expression, expectedType, + expectedTypeRef, sink, isReceiver, acceptLambdaAtoms, @@ -120,6 +123,7 @@ internal fun Candidate.resolveArgument( this.system.getBuilder(), argument, expectedType, + parameter.returnTypeRef, sink, isReceiver, { this.postponedAtoms += it }, diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArguments.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArguments.kt index e9d06d4a296..6c78710e0e9 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArguments.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArguments.kt @@ -15,6 +15,7 @@ fun preprocessLambdaArgument( csBuilder: ConstraintSystemBuilder, argument: FirAnonymousFunction, expectedType: ConeKotlinType?, + expectedTypeRef: FirTypeRef, acceptLambdaAtoms: (PostponedResolvedAtomMarker) -> Unit, forceResolution: Boolean = false ) { @@ -22,7 +23,9 @@ fun preprocessLambdaArgument( //return LambdaWithTypeVariableAsExpectedTypeAtom(argument, expectedType) } - val resolvedArgument = extractLambdaInfoFromFunctionalType(expectedType, argument) ?: extraLambdaInfo(expectedType, argument, csBuilder) + val resolvedArgument = + extractLambdaInfoFromFunctionalType(expectedType, expectedTypeRef, argument) + ?: extraLambdaInfo(expectedType, expectedTypeRef, argument, csBuilder) if (expectedType != null) { // val lambdaType = createFunctionType( @@ -58,7 +61,19 @@ private val ConeKotlinType.isSuspendFunctionType: Boolean } } -private val ConeKotlinType.receiverType: ConeKotlinType? get() = null +private fun ConeKotlinType.receiverType(expectedTypeRef: FirTypeRef): ConeKotlinType? { + if (isFunctionalTypeWithReceiver(expectedTypeRef)) { + return (this.typeArguments.first() as ConeTypedProjection).type + } + return null +} + +private fun isFunctionalTypeWithReceiver(typeRef: FirTypeRef) = + typeRef.annotations.any { + val coneTypeSafe = it.annotationTypeRef.coneTypeSafe() ?: return@any false + coneTypeSafe.lookupTag.classId.asString() == "kotlin/ExtensionFunctionType" + } + private val ConeKotlinType.returnType: ConeKotlinType get() { require(this is ConeClassType) @@ -81,6 +96,7 @@ private val FirAnonymousFunction.receiverType get() = receiverTypeRef?.coneTypeS private fun extraLambdaInfo( expectedType: ConeKotlinType?, + expectedTypeRef: FirTypeRef, argument: FirAnonymousFunction, csBuilder: ConstraintSystemBuilder ): ResolvedLambdaAtom { @@ -93,7 +109,7 @@ private fun extraLambdaInfo( val typeVariable = TypeVariableForLambdaReturnType(argument, "_L") - val receiverType = null//argumentAsFunctionExpression?.receiverType + val receiverType = argumentAsFunctionExpression?.receiverType val returnType = argumentAsFunctionExpression?.returnType ?: expectedType?.typeArguments?.singleOrNull()?.safeAs()?.type?.takeIf { isFunctionSupertype } @@ -107,12 +123,16 @@ private fun extraLambdaInfo( return ResolvedLambdaAtom(argument, isSuspend, receiverType, parameters, returnType, typeVariable.takeIf { newTypeVariableUsed }) } -private fun extractLambdaInfoFromFunctionalType(expectedType: ConeKotlinType?, argument: FirAnonymousFunction): ResolvedLambdaAtom? { +private fun extractLambdaInfoFromFunctionalType( + expectedType: ConeKotlinType?, + expectedTypeRef: FirTypeRef, + argument: FirAnonymousFunction +): ResolvedLambdaAtom? { if (expectedType == null || !expectedType.isBuiltinFunctionalType) return null val parameters = extractLambdaParameters(expectedType, argument) val argumentAsFunctionExpression = argument//.safeAs() - val receiverType = argumentAsFunctionExpression?.receiverType ?: expectedType.receiverType + val receiverType = argumentAsFunctionExpression?.receiverType ?: expectedType.receiverType(expectedTypeRef) val returnType = argumentAsFunctionExpression?.returnType ?: expectedType.returnType return ResolvedLambdaAtom( diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArgumentsAnalyzer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArgumentsAnalyzer.kt index 7c24b3b389e..d9d02a04913 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArgumentsAnalyzer.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArgumentsAnalyzer.kt @@ -106,6 +106,7 @@ class PostponedArgumentsAnalyzer( c.getBuilder(), it, lambda.returnType.let(::substitute), + lambda.atom.returnTypeRef, // TODO: proper ref checkerSink, false, { atom = it }, diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolverParts.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolverParts.kt index ace40838a9a..9da7ccf4dc8 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolverParts.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolverParts.kt @@ -95,6 +95,7 @@ internal sealed class CheckReceivers : ResolutionStage() { resolveArgumentExpression( candidate.csBuilder, explicitReceiverExpression, candidate.substitutor.substituteOrSelf(receiverParameterValue.type), + explicitReceiverExpression.typeRef, sink, isReceiver = true, typeProvider = callInfo.typeProvider, acceptLambdaAtoms = { candidate.postponedAtoms += it } ) }