[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 {
+    val valueParametersTypesIncludingReceiver = expectedType.valueParameterTypesIncludingReceiver(session)
-        if (receiverType != null && expectedType.isExtensionFunctionType(session)) it.drop(1) else it
+    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
+            this,
-        ) ?: extraLambdaInfo(expectedType, anonymousFunction, csBuilder, context.session, 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

compiler/testData/diagnostics/tests/inference/completion/postponedArgumentsAnalysis/basic.fir.kt

@@ -149,14 +149,14 @@ fun main() {
     select(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function2<kotlin.String, kotlin.String, kotlin.Unit>")!>id(fun String.(x: String) {})<!>, <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function2<kotlin.String, kotlin.String, kotlin.Unit>")!>id(fun(x: String, y: String) {})<!>)
     select(id(fun String.(x: String) {}), id(fun(x: String, y: String) { }), { x: String -> <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>this<!> })
     select(id(fun String.(x: String) {}), id(fun(x: String, y: String) { }), { x -> <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>this<!> })
-    select(id(fun String.(x: String) {}), id(fun(x: String, y: String) { }), <!ARGUMENT_TYPE_MISMATCH!>{ x: String, y: String -> x }<!>)
+    select(id(fun String.(x: String) {}), id(fun(x: String, y: String) { }), { x: String, y: String -> x })
     // Convert to extension lambda is impossible because the lambda parameter types aren't specified explicitly
     select(id(fun String.(x: String) {}), id(fun(x: String, y: String) { }), <!ARGUMENT_TYPE_MISMATCH!>{ x, <!CANNOT_INFER_PARAMETER_TYPE!>y<!> -> x }<!>)
     select(id(id(fun(x: String, y: String) { }), <!TOO_MANY_ARGUMENTS!>fun String.(x: String) {}<!>), { x, y -> x })
     val x26: Int.(String) -> Int = fun (x: String) = 10 // it must be error, see KT-38439
     // Receiver must be specified in anonymous function declaration
     val x27: Int.(String) -> Int = id(<!ARGUMENT_TYPE_MISMATCH!>fun (x: String) = 10<!>)
-    select(id<Int.(String) -> Unit> {}, <!ARGUMENT_TYPE_MISMATCH!>{ x: Int, y: String -> x }<!>)
+    select(id<Int.(String) -> Unit> {}, { x: Int, y: String -> x })
 
     // Inferring lambda parameter types by partially specified parameter types of other lambdas
     select(id { x, y -> x.inv() + y.toByte() }, { x: Int, y -> y.toByte() }, { x, y: Number -> x.inv() })
@@ -218,7 +218,7 @@ fun main() {
     <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function<kotlin.Any>")!>select({ x: Int -> TODO() }, id { x: Int, y: Number -> Any() })<!>
     <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<*, kotlin.String?>")!>select(id { x: Int -> null }, id { x: String -> "" })<!>
     <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function<kotlin.Int>")!>select(id { x: Int -> 10 }, id { x: Int, y: Number -> TODO() })<!>
-    val x68: String.(String) -> String = select(id <!ARGUMENT_TYPE_MISMATCH!>{ x: String, y: String -> "10" }<!>, id <!ARGUMENT_TYPE_MISMATCH!>{ x: String, y: String -> "TODO()" }<!>)
+    val x68: String.(String) -> String = select(id { x: String, y: String -> "10" }, id { x: String, y: String -> "TODO()" })
 
     // Anonymous functions
     val x69: (C) -> Unit = selectB({ it }, { }, id(fun (x) { <!DEBUG_INFO_EXPRESSION_TYPE("A")!>x<!> }))
@@ -229,7 +229,7 @@ fun main() {
     val x71: String.() -> Unit = select(<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.Unit>")!>id(fun String.() { })<!>, <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.String, kotlin.Unit>")!>id(fun(x: String) {})<!>)
     val x72: String.() -> Unit = select(fun String.() { }, fun(x: String) {}) // must be error
     select(id(fun String.(x: String) {}), id(fun(x: String, y: String) { }), fun (x, y) { <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>x<!>;<!DEBUG_INFO_EXPRESSION_TYPE("kotlin.String")!>y<!> })
-    select(id<Int.(String) -> Unit>(fun (x, y) {}), <!ARGUMENT_TYPE_MISMATCH!>{ x: Int, y: String -> x }<!>) // receiver of anonymous function must be specified explicitly
+    select(id<Int.(String) -> Unit>(fun (x, y) {}), { x: Int, y: String -> x }) // receiver of anonymous function must be specified explicitly
-    select(id<Int.(String) -> Unit>(fun Int.(y) {}), <!ARGUMENT_TYPE_MISMATCH!>{ x: Int, y: String -> x }<!>)
+    select(id<Int.(String) -> Unit>(fun Int.(y) {}), { x: Int, y: String -> x })
     <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Function1<kotlin.Number, java.io.Serializable>")!>select(A3(), fun (x) = "", { a -> <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Number")!>a<!> })<!>
 }

compiler/testData/diagnostics/tests/regressions/kt30245.fir.kt

@@ -7,8 +7,8 @@ fun <K> id(x: K): K = x
 
 fun test() {
     val f00: Sample.() -> Unit = id { val a = 1 }
-    val f01: Sample.() -> Unit = id <!ARGUMENT_TYPE_MISMATCH!>{ s: Sample -> }<!>
+    val f01: Sample.() -> Unit = id { s: Sample -> }
-    val f02: Sample.() -> Unit = id<Sample.() -> Unit> <!ARGUMENT_TYPE_MISMATCH!>{ s: Sample -> }<!>
+    val f02: Sample.() -> Unit = id<Sample.() -> Unit> { s: Sample -> }
 }
 
 enum class E { VALUE }
@@ -45,7 +45,7 @@ fun test1() { // to extension lambda 0
     val g10: E0 = id(fun Int.(): Int = this) // oi+ ni+
 
     val w11 = W1 <!ARGUMENT_TYPE_MISMATCH!>{ i: Int -> i }<!> // oi- ni-
-    val i11: E0 = id <!ARGUMENT_TYPE_MISMATCH!>{ i: Int -> i }<!> // o1+ ni+
+    val i11: E0 = id { i: Int -> i } // o1+ ni+
     val w12 = W1 <!ARGUMENT_TYPE_MISMATCH!>{ <!CANNOT_INFER_PARAMETER_TYPE!>i<!> -> i }<!> // oi- ni-
     val i12: E0 = id <!ARGUMENT_TYPE_MISMATCH!>{ <!CANNOT_INFER_PARAMETER_TYPE!>i<!> -> i }<!> // oi- ni-
     val j12 = id<E0> <!ARGUMENT_TYPE_MISMATCH!>{ <!CANNOT_INFER_PARAMETER_TYPE!>i<!> -> i }<!> // oi- ni-
@@ -84,9 +84,9 @@ fun test2() { // to extension lambda 1
     val i27a: E1 = <!INITIALIZER_TYPE_MISMATCH!>when (e) { E.VALUE ->  { s -> <!NO_THIS!>this<!> + s.<!UNRESOLVED_REFERENCE!>length<!> } }<!> // oi+ ni+
 
     val w28 = W2 <!ARGUMENT_TYPE_MISMATCH!>{ i: Int, <!CANNOT_INFER_PARAMETER_TYPE!>s<!> -> i <!OVERLOAD_RESOLUTION_AMBIGUITY!>+<!> s.<!UNRESOLVED_REFERENCE!>length<!> }<!> // oi- ni-
-    val i28: E1 = id <!ARGUMENT_TYPE_MISMATCH!>{ i: Int, <!CANNOT_INFER_PARAMETER_TYPE!>s<!> -> i <!OVERLOAD_RESOLUTION_AMBIGUITY!>+<!> s.<!UNRESOLVED_REFERENCE!>length<!> }<!> // oi- ni-
+    val i28: E1 = id { i: Int, <!CANNOT_INFER_PARAMETER_TYPE!>s<!> -> <!ARGUMENT_TYPE_MISMATCH!>i <!OVERLOAD_RESOLUTION_AMBIGUITY!>+<!> s.<!UNRESOLVED_REFERENCE!>length<!><!> } // oi- ni-
     val w29 = W2 <!ARGUMENT_TYPE_MISMATCH!>{ i: Int, s: String -> i + s.length }<!> // oi- ni-
-    val i29: E1 = id <!ARGUMENT_TYPE_MISMATCH!>{ i: Int, s: String -> i + s.length }<!> // oi+ ni+
+    val i29: E1 = id { i: Int, s: String -> i + s.length } // oi+ ni+
 
     // yet unsupported cases with ambiguity for the lambda conversion (commented constructors in wrappers above)
 //    val w2a = W2 { i, s -> i + s.length } // overload oi- ni-