[FIR] Coerce first parameter to extension receiver ^KT-46371 Fixed

Fix typos

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/InferenceUtils.kt

@@ -241,7 +241,8 @@ fun extractLambdaInfoFromFunctionalType(
     argument: FirAnonymousFunction,
     returnTypeVariable: ConeTypeVariableForLambdaReturnType?,
     components: BodyResolveComponents,
-    candidate: Candidate?
+    candidate: Candidate?,
+    duringCompletion: Boolean,
 ): ResolvedLambdaAtom? {
     val session = components.session
     if (expectedType == null) return null
@@ -252,7 +253,8 @@ fun extractLambdaInfoFromFunctionalType(
             argument,
             returnTypeVariable,
             components,
-            candidate
+            candidate,
+            duringCompletion
         )
     }
     if (!expectedType.isBuiltinFunctionalType(session)) return null
@@ -270,13 +272,31 @@ fun extractLambdaInfoFromFunctionalType(
     // For lambdas, the existence of the receiver is always implied by the expected type, and a value parameter
     // can never fill its role.
     val receiverType = if (argument.isLambda) expectedType.receiverType(session) else argument.receiverType
-    val expectedParameters = expectedType.valueParameterTypesIncludingReceiver(session).let {
-        if (receiverType != null && expectedType.isExtensionFunctionType(session)) it.drop(1) else it
+    val valueParametersTypesIncludingReceiver = expectedType.valueParameterTypesIncludingReceiver(session)
+    val isExtensionFunctionType = expectedType.isExtensionFunctionType(session)
+    val expectedParameters = valueParametersTypesIncludingReceiver.let {
+        if (receiverType != null && isExtensionFunctionType) it.drop(1) else it
     }
-    val parameters = if (argument.isLambda && argument.valueParameters.isEmpty() && expectedParameters.size < 2) {
+
+    var coerceFirstParameterToExtensionReceiver = false
+    val argumentValueParameters = argument.valueParameters
+    val parameters = if (argument.isLambda && argumentValueParameters.isEmpty() && expectedParameters.size < 2) {
         expectedParameters // Infer existence of a parameter named `it` of an appropriate type.
     } else {
-        argument.valueParameters.mapIndexed { index, parameter ->
+        if (duringCompletion &&
+            argument.isLambda &&
+            isExtensionFunctionType &&
+            valueParametersTypesIncludingReceiver.size == argumentValueParameters.size
+        ) {
+            // (T, ...) -> V can be converter to T.(...) -> V
+            val firstValueParameter = argumentValueParameters.firstOrNull()
+            val extensionParameter = valueParametersTypesIncludingReceiver.firstOrNull()
+            if (firstValueParameter?.returnTypeRef?.coneTypeSafe<ConeKotlinType>() == extensionParameter?.type) {
+                coerceFirstParameterToExtensionReceiver = true
+            }
+        }
+
+        argumentValueParameters.mapIndexed { index, parameter ->
             parameter.returnTypeRef.coneTypeSafe()
                 ?: expectedParameters.getOrNull(index)
                 ?: ConeClassErrorType(
@@ -293,6 +313,7 @@ fun extractLambdaInfoFromFunctionalType(
         parameters,
         returnType,
         typeVariableForLambdaReturnType = returnTypeVariable,
-        candidate
+        candidate,
+        coerceFirstParameterToExtensionReceiver
     )
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt

@@ -45,13 +45,15 @@ fun Candidate.preprocessLambdaArgument(
             anonymousFunction,
             returnTypeVariable,
             context.bodyResolveComponents,
-            this
-        ) ?: extraLambdaInfo(expectedType, anonymousFunction, csBuilder, context.session, this)
+            this,
+            duringCompletion || sink == null
+        ) ?: extractLambdaInfo(expectedType, anonymousFunction, csBuilder, context.session, this)
 
     if (expectedType != null) {
         // TODO: add SAM conversion processing
+        val parameters = resolvedArgument.parameters
         val lambdaType = createFunctionalType(
-            resolvedArgument.parameters,
+            if (resolvedArgument.coerceFirstParameterToExtensionReceiver) parameters.drop(1) else parameters,
             resolvedArgument.receiver,
             resolvedArgument.returnType,
             isSuspend = resolvedArgument.isSuspend
@@ -86,7 +88,7 @@ fun Candidate.preprocessCallableReference(
     postponedAtoms += ResolvedCallableReferenceAtom(argument, expectedType, lhs, context.session)
 }
 
-private fun extraLambdaInfo(
+private fun extractLambdaInfo(
     expectedType: ConeKotlinType?,
     argument: FirAnonymousFunction,
     csBuilder: ConstraintSystemBuilder,
@@ -123,6 +125,7 @@ private fun extraLambdaInfo(
         parameters,
         returnType,
         typeVariable.takeIf { newTypeVariableUsed },
-        candidate
+        candidate,
+        coerceFirstParameterToExtensionReceiver = false
     )
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedAtoms.kt

@@ -49,7 +49,8 @@ class ResolvedLambdaAtom(
     val parameters: List<ConeKotlinType>,
     var returnType: ConeKotlinType,
     typeVariableForLambdaReturnType: ConeTypeVariableForLambdaReturnType?,
-    candidateOfOuterCall: Candidate?
+    candidateOfOuterCall: Candidate?,
+    val coerceFirstParameterToExtensionReceiver: Boolean
 ) : PostponedResolvedAtom() {
     init {
         candidateOfOuterCall?.let {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -731,7 +731,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
     ): FirAnonymousFunction {
         val resolvedLambdaAtom = (expectedTypeRef as? FirResolvedTypeRef)?.let {
             extractLambdaInfoFromFunctionalType(
-                it.type, it, anonymousFunction, returnTypeVariable = null, components, candidate = null
+                it.type, it, anonymousFunction, returnTypeVariable = null, components, candidate = null, duringCompletion = false
             )
         }
         var lambda = anonymousFunction