FIR: detect extension function types by reference annotation

Later it should be replaced by cone type annotation analysis

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

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<ConeClassType>() ?: 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<ConeTypedProjection>()?.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<FunctionExpression>()
-    val receiverType = argumentAsFunctionExpression?.receiverType ?: expectedType.receiverType
+    val receiverType = argumentAsFunctionExpression?.receiverType ?: expectedType.receiverType(expectedTypeRef)
     val returnType = argumentAsFunctionExpression?.returnType ?: expectedType.returnType
 
     return ResolvedLambdaAtom(

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

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 }
                 )
             }