[FIR] Fix mutant extension function types

Consider the following example from
`extensionLambdasAndArrow.kt`:

```
val x4: String.() -> String = if (true) {
    { str: String -> "this" }
} else {
    { str: String -> "this" }
}
```

Because of
`coerceFirstParameterToExtensionReceiver`
the given lambdas must be of the type
`String.() -> String`, but because of a bug
they are `String.(String) -> String`. At the
same time, during inference their expected
types are, indeed, calculated correctly as
`String.() -> String`.

^KT-59394 Declined
(no more compiler crashes, #potential-feature)

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirLightTreeBlackBoxCodegenTestGenerated.java

@@ -18799,6 +18799,12 @@ public class FirLightTreeBlackBoxCodegenTestGenerated extends AbstractFirLightTr
             runTest("compiler/testData/codegen/box/fir/SamWithReceiverMavenProjectImportHandler.kt");
         }
 
+        @Test
+        @TestMetadata("selectingLambdas.kt")
+        public void testSelectingLambdas() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/selectingLambdas.kt");
+        }
+
         @Test
         @TestMetadata("staticImportFromEnum.kt")
         public void testStaticImportFromEnum() throws Exception {

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirPsiBlackBoxCodegenTestGenerated.java

@@ -18799,6 +18799,12 @@ public class FirPsiBlackBoxCodegenTestGenerated extends AbstractFirPsiBlackBoxCo
             runTest("compiler/testData/codegen/box/fir/SamWithReceiverMavenProjectImportHandler.kt");
         }
 
+        @Test
+        @TestMetadata("selectingLambdas.kt")
+        public void testSelectingLambdas() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/selectingLambdas.kt");
+        }
+
         @Test
         @TestMetadata("staticImportFromEnum.kt")
         public void testStaticImportFromEnum() throws Exception {

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

@@ -277,17 +277,16 @@ class FirCallCompleter(
 
             val expectedReturnTypeRef = expectedReturnType?.let { lambdaArgument.returnTypeRef.resolvedTypeFromPrototype(it) }
 
-            if (receiverType == null) {
-                lambdaArgument.replaceReceiverParameter(null)
-            } else {
-                lambdaArgument.receiverParameter?.apply {
-                    replaceTypeRef(
-                        typeRef.resolvedTypeFromPrototype(
-                            receiverType.approximateLambdaInputType(),
-                            source?.fakeElement(KtFakeSourceElementKind.ImplicitTypeRef),
-                        )
-                    )
+            when {
+                receiverType == null -> lambdaArgument.replaceReceiverParameter(null)
+                !lambdaAtom.coerceFirstParameterToExtensionReceiver -> {
+                    lambdaArgument.receiverParameter?.apply {
+                        val type = receiverType.approximateLambdaInputType()
+                        val source = source?.fakeElement(KtFakeSourceElementKind.ImplicitTypeRef)
+                        replaceTypeRef(typeRef.resolvedTypeFromPrototype(type, source))
+                    }
                 }
+                else -> lambdaArgument.replaceReceiverParameter(null)
             }
 
             if (contextReceivers.isNotEmpty()) {
@@ -304,8 +303,15 @@ class FirCallCompleter(
 
             val lookupTracker = session.lookupTracker
             val fileSource = components.file.source
+            val theParameters = when {
+                lambdaAtom.coerceFirstParameterToExtensionReceiver -> when (receiverType) {
+                    null -> error("Coercion to extension receiver while no receiver present")
+                    else -> listOf(receiverType) + parameters
+                }
+                else -> parameters
+            }
             lambdaArgument.valueParameters.forEachIndexed { index, parameter ->
-                val newReturnType = parameters[index].approximateLambdaInputType()
+                val newReturnType = theParameters[index].approximateLambdaInputType()
                 val newReturnTypeRef = if (parameter.returnTypeRef is FirImplicitTypeRef) {
                     newReturnType.toFirResolvedTypeRef(parameter.source?.fakeElement(KtFakeSourceElementKind.ImplicitReturnTypeOfLambdaValueParameter))
                 } else parameter.returnTypeRef.resolvedTypeFromPrototype(newReturnType)

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

@@ -24,7 +24,7 @@ fun extractLambdaInfoFromFunctionType(
     returnTypeVariable: ConeTypeVariableForLambdaReturnType?,
     components: BodyResolveComponents,
     candidate: Candidate?,
-    duringCompletion: Boolean,
+    allowCoercionToExtensionReceiver: Boolean,
 ): ResolvedLambdaAtom? {
     val session = components.session
     if (expectedType == null) return null
@@ -35,7 +35,7 @@ fun extractLambdaInfoFromFunctionType(
             returnTypeVariable,
             components,
             candidate,
-            duringCompletion
+            allowCoercionToExtensionReceiver,
         )
     }
     val expectedFunctionKind = expectedType.functionTypeKind(session) ?: return null
@@ -77,7 +77,7 @@ fun extractLambdaInfoFromFunctionType(
     val parameters = if (argument.isLambda && !argument.hasExplicitParameterList && expectedParameters.size < 2) {
         expectedParameters // Infer existence of a parameter named `it` of an appropriate type.
     } else {
-        if (duringCompletion &&
+        if (allowCoercionToExtensionReceiver &&
             argument.isLambda &&
             isExtensionFunctionType &&
             valueParametersTypesIncludingReceiver.size == argumentValueParameters.size
@@ -90,7 +90,11 @@ fun extractLambdaInfoFromFunctionType(
             }
         }
 
-        argumentValueParameters.mapIndexed { index, parameter ->
+        if (coerceFirstParameterToExtensionReceiver) {
+            argumentValueParameters.drop(1)
+        } else {
+            argumentValueParameters
+        }.mapIndexed { index, parameter ->
             parameter.returnTypeRef.coneTypeSafe()
                 ?: expectedParameters.getOrNull(index)
                 ?: ConeErrorType(ConeCannotInferValueParameterType(parameter.symbol))

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

@@ -66,7 +66,7 @@ fun Candidate.preprocessLambdaArgument(
             ?: FunctionTypeKind.Function
         val lambdaType = createFunctionType(
             functionTypeKind,
-            if (resolvedArgument.coerceFirstParameterToExtensionReceiver) parameters.drop(1) else parameters,
+            parameters,
             resolvedArgument.receiver,
             resolvedArgument.returnType,
             contextReceivers = resolvedArgument.contextReceivers,

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

@@ -816,7 +816,8 @@ open class FirDeclarationsResolveTransformer(
     ): FirAnonymousFunction {
         val resolvedLambdaAtom = (expectedTypeRef as? FirResolvedTypeRef)?.let {
             extractLambdaInfoFromFunctionType(
-                it.type, anonymousFunction, returnTypeVariable = null, components, candidate = null, duringCompletion = false
+                it.type, anonymousFunction, returnTypeVariable = null, components, candidate = null,
+                allowCoercionToExtensionReceiver = true,
             )
         }
         var lambda = anonymousFunction
@@ -826,7 +827,9 @@ open class FirDeclarationsResolveTransformer(
         }
         lambda = buildAnonymousFunctionCopy(lambda) {
             receiverParameter = lambda.receiverParameter?.takeIf { it.typeRef !is FirImplicitTypeRef }
-                ?: resolvedLambdaAtom?.receiver?.let { coneKotlinType ->
+                ?: resolvedLambdaAtom?.receiver?.takeIf {
+                    !resolvedLambdaAtom.coerceFirstParameterToExtensionReceiver
+                }?.let { coneKotlinType ->
                     lambda.receiverParameter?.apply {
                         replaceTypeRef(typeRef.resolvedTypeFromPrototype(coneKotlinType))
                     }
@@ -903,7 +906,16 @@ open class FirDeclarationsResolveTransformer(
                 listOf(itParam)
             }
 
-            else -> obtainValueParametersFromExpectedParameterTypes(resolvedLambdaAtom.parameters, lambda)
+            else -> {
+                val parameters = if (resolvedLambdaAtom.coerceFirstParameterToExtensionReceiver) {
+                    val receiver = resolvedLambdaAtom.receiver ?: error("Coercion to an extension function type, but no receiver found")
+                    listOf(receiver) + resolvedLambdaAtom.parameters
+                } else {
+                    resolvedLambdaAtom.parameters
+                }
+
+                obtainValueParametersFromExpectedParameterTypes(parameters, lambda)
+            }
         }
     }
 

compiler/testData/codegen/box/fir/selectingLambdas.kt

@@ -0,0 +1,25 @@
+// TARGET_BACKEND: JVM_IR
+// IGNORE_BACKEND_K1: JVM_IR
+
+fun <K> id(x: K): K = x
+
+fun box(): String {
+    val x4: String.() -> String = if (true) {fun String.(): String { return "this" }} else {{ str: String -> "that" }}
+    val x51: String.() -> String = if (false) {fun String.(): String { return "this" }} else {{ "that" }}
+    val x52: String.() -> String = if (false) {fun String.(): String { return "this" }} else {{ str: String -> "that" }}
+    val x53: String.() -> String = if (false) {fun String.(): String { return "this" }} else {fun(str: String) = "that"}
+
+    val i28: Int.(String) -> Int = id { i: Int, s -> i + s.length }
+
+    val result = "test".x4() +
+            "::" + "test".x51() +
+            "::" + "test".x52() +
+            "::" + "test".x53() +
+            "::" + 10.i28("test")
+
+    return if (result == "this::that::that::that::14") {
+        "OK"
+    } else {
+        "Fail: $result"
+    }
+}

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

@@ -83,7 +83,7 @@ fun test2() { // to extension lambda 1
     val i27a: E1 = when (e) { E.VALUE ->  { s -> this + s.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 { i: Int, <!CANNOT_INFER_PARAMETER_TYPE!>s<!> -> i <!OVERLOAD_RESOLUTION_AMBIGUITY!>+<!> s.<!UNRESOLVED_REFERENCE!>length<!> } // oi- ni-
+    val i28: E1 = id { i: Int, s -> i + s.length } // oi- ni-
     val w29 = W2 <!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+
 

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -18799,6 +18799,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/fir/SamWithReceiverMavenProjectImportHandler.kt");
         }
 
+        @Test
+        @TestMetadata("selectingLambdas.kt")
+        public void testSelectingLambdas() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/selectingLambdas.kt");
+        }
+
         @Test
         @TestMetadata("staticImportFromEnum.kt")
         public void testStaticImportFromEnum() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenWithIrInlinerTestGenerated.java

@@ -18799,6 +18799,12 @@ public class IrBlackBoxCodegenWithIrInlinerTestGenerated extends AbstractIrBlack
             runTest("compiler/testData/codegen/box/fir/SamWithReceiverMavenProjectImportHandler.kt");
         }
 
+        @Test
+        @TestMetadata("selectingLambdas.kt")
+        public void testSelectingLambdas() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/selectingLambdas.kt");
+        }
+
         @Test
         @TestMetadata("staticImportFromEnum.kt")
         public void testStaticImportFromEnum() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -15377,6 +15377,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/fir/emptyIntersectionWarning.kt");
         }
 
+        @TestMetadata("selectingLambdas.kt")
+        public void ignoreSelectingLambdas() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/selectingLambdas.kt");
+        }
+
         @TestMetadata("typeParameterInClashingAccessor.kt")
         public void ignoreTypeParameterInClashingAccessor() throws Exception {
             runTest("compiler/testData/codegen/box/fir/typeParameterInClashingAccessor.kt");