FE: Support assigning lambdas to receiver/context function types

It's been controversially working since 2.0, but was broken
once ResolveTopLevelLambdasAsSyntheticCallArgument became enabled
by default.

See KT-81115 and related issues for details.

^KT-81115 Fixed

Author: dzharkov

compiler/fir/analysis-tests/testData/resolve/inference/lambdasWithExplicitParametersViaReceiversFeatureEnabled.kt

@@ -3,13 +3,13 @@
 // ISSUE: KT-81115
 
 val l1: String.() -> String = fun(x: String) = ""
-val l2: String.() -> String = <!INITIALIZER_TYPE_MISMATCH!>{ x: String -> "" }<!>
+val l2: String.() -> String = { x: String -> "" }
 
 val l3: String.() -> String = fun(x) = x
-val l4: String.() -> String = <!INITIALIZER_TYPE_MISMATCH!>{ <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>x<!> -> x }<!>
+val l4: String.() -> String = { x -> x }
 
 val l5: String.(Int) -> String = fun(x: String, y) = ""
-val l6: String.(Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>{ x: String, <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>y<!> -> "" }<!>
+val l6: String.(Int) -> String = { x: String, y -> "" }
 
 val w1: String.(Int, Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>fun(x: String, y) = ""<!>
 val w2: String.(Int, Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>{ x: String, y -> "" }<!>

compiler/fir/analysis-tests/testData/resolve/inference/lambdasWithExplicitParametersWithContexts.kt

@@ -3,10 +3,10 @@
 // ISSUE: KT-81115
 
 val l1: context(String) (Int) -> String = fun(x: String, y) = ""
-val l2: context(String) (Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>{ x: String, <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>y<!> -> "" }<!>
+val l2: context(String) (Int) -> String = { x: String, y -> "" }
 
 val l3: context(String) Double.(Int) -> String = fun(x: String, y, z: Int) = ""
-val l4: context(String) Double.(Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>{ x: String, <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>y<!>, z: Int -> "" }<!>
+val l4: context(String) Double.(Int) -> String = { x: String, y, z: Int -> "" }
 
 val w1: context(String) Double.(Int, Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>fun(x: String, y, z: Int) = ""<!>
 val w2: context(String) Double.(Int, Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>{ x: String, y, z: Int -> "" }<!>

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSyntheticCallGenerator.kt

@@ -255,10 +255,13 @@ class FirSyntheticCallGenerator(
     ): FirExpression {
         val argumentList = buildUnaryArgumentList(anonymousFunctionExpression)
 
+        val expectedType =
+            expectedTypeData?.expectedType?.adaptExpectedTypeForLambdaIfNeeded(anonymousFunctionExpression.anonymousFunction)
+
         val reference = generateCalleeReferenceToFunctionWithExpectedTypeForArgument(
             anonymousFunctionExpression,
             argumentList,
-            expectedTypeData?.expectedType,
+            expectedType,
             context,
         )
 
@@ -298,6 +301,26 @@ class FirSyntheticCallGenerator(
         return resolvedArgument
     }
 
+    private fun ConeKotlinType.adaptExpectedTypeForLambdaIfNeeded(lambda: FirAnonymousFunction): ConeKotlinType {
+        if (!lambda.hasExplicitParameterList) return this
+        if (!isSomeFunctionType(session)) return this
+
+        val isThereReceiver = receiverType(session) != null
+        if (!isThereReceiver && contextParameterNumberForFunctionType == 0) return this
+
+        val classLikeType = unwrapLowerBound() as? ConeClassLikeType ?: return this
+
+        return when {
+            lambda.valueParameters.size == classLikeType.valueParameterTypesIncludingReceiver(session).size ->
+                classLikeType.withAttributes(
+                    classLikeType.attributes
+                        .remove(CompilerConeAttributes.ExtensionFunctionType::class)
+                        .remove(CompilerConeAttributes.ContextFunctionTypeParams::class)
+                )
+            else -> this
+        }
+    }
+
     /**
      * After resolution of a top-level lambda via synthetic call, we have some diagnostic, which in most of the cases says
      * that the type of the lambda can't be passed to the given expected type.

compiler/testData/diagnostics/tests/functionLiterals/kt56138.fir.kt

@@ -6,9 +6,9 @@ fun takeLambda1(f: String.() -> String) {}
 fun takeLambda2(f: String.(String) -> String) {}
 
 fun test_1() {
-    val x1: String.(String) -> String = <!INITIALIZER_TYPE_MISMATCH!>{ str, <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>str2<!> -> "this" }<!>
+    val x1: String.(String) -> String = { str, str2 -> "this" }
     val x2: String.() -> String = { <!UNRESOLVED_REFERENCE!>it<!> }
-    val x3: String.() -> String = <!INITIALIZER_TYPE_MISMATCH!>{ <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>x<!> -> "x" }<!>
+    val x3: String.() -> String = { x -> "x" }
 }
 
 fun test_2() {
@@ -19,9 +19,9 @@ fun test_2() {
 
 fun test_3(b: Boolean) {
     val x1: String.(String) -> String = if (b) {
-        { str, <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>str2<!> -> "this" }
+        { str, str2 -> "this" }
     } else {
-        { str, <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>str2<!> -> "this" }
+        { str, str2 -> "this" }
     }
 
     val x2: String.() -> String = if (b) {
@@ -31,9 +31,9 @@ fun test_3(b: Boolean) {
     }
 
     val x3: String.() -> String = if (b) {
-        <!ARGUMENT_TYPE_MISMATCH!>{ <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>x<!> -> "x" }<!>
+        { x -> "x" }
     } else {
-        <!ARGUMENT_TYPE_MISMATCH!>{ <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>x<!> -> "x" }<!>
+        { x -> "x" }
     }
 }
 

compiler/testData/diagnostics/tests/inference/extensionLambdasAndArrow.fir.kt

@@ -1,4 +1,4 @@
-// RUN_PIPELINE_TILL: FRONTEND
+// RUN_PIPELINE_TILL: BACKEND
 // DIAGNOSTICS: -UNUSED_VARIABLE -UNUSED_ANONYMOUS_PARAMETER
 
 fun <T> select(vararg x: T) = x[0]
@@ -9,9 +9,9 @@ fun main() {
     val x3: () -> String = if (true) {{ -> "this" }} else {{ -> "this" }}
     val x4: String.() -> String = if (true) {{ str: String -> "this" }} else {{ str: String -> "this" }}
     val x41: String.(String) -> String = if (true) {{ str: String, str2: String -> "this" }} else {{ str: String, str2: String -> "this" }}
-    val x42: String.(String) -> String = if (true) {{ str, <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>str2<!> -> "this" }} else {{ str, <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>str2<!> -> "this" }}
-    val x5: String.() -> String = if (true) {<!ARGUMENT_TYPE_MISMATCH!>{ <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>str<!> -> "this" }<!>} else {<!ARGUMENT_TYPE_MISMATCH!>{ <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>str<!> -> "this" }<!>}
-    val x6: String.() -> String = <!INITIALIZER_TYPE_MISMATCH!>if (true) {<!ARGUMENT_TYPE_MISMATCH!>{ <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>str<!> -> "this" }<!>} else {{ "this" }}<!>
+    val x42: String.(String) -> String = if (true) {{ str, str2 -> "this" }} else {{ str, str2 -> "this" }}
+    val x5: String.() -> String = if (true) {{ str -> "this" }} else {{ str -> "this" }}
+    val x6: String.() -> String = if (true) {{ str -> "this" }} else {{ "this" }}
     val x7: String.() -> String = select({ -> this }, { -> this })
     val x8: String.() -> String = select({ this }, { this })
 }

compiler/testData/diagnostics/tests/inference/extensionLambdasAndArrow.kt

@@ -1,4 +1,4 @@
-// RUN_PIPELINE_TILL: FRONTEND
+// RUN_PIPELINE_TILL: BACKEND
 // DIAGNOSTICS: -UNUSED_VARIABLE -UNUSED_ANONYMOUS_PARAMETER
 
 fun <T> select(vararg x: T) = x[0]

compiler/testData/diagnostics/tests/inference/lambdasWithExplicitParametersViaReceivers.fir.kt

@@ -2,13 +2,13 @@
 // ISSUE: KT-81115
 
 val l1: String.() -> String = fun(x: String) = ""
-val l2: String.() -> String = <!INITIALIZER_TYPE_MISMATCH!>{ x: String -> "" }<!>
+val l2: String.() -> String = { x: String -> "" }
 
 val l3: String.() -> String = fun(x) = x
-val l4: String.() -> String = <!INITIALIZER_TYPE_MISMATCH!>{ <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>x<!> -> x }<!>
+val l4: String.() -> String = { x -> x }
 
 val l5: String.(Int) -> String = fun(x: String, y) = ""
-val l6: String.(Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>{ x: String, <!CANNOT_INFER_VALUE_PARAMETER_TYPE!>y<!> -> "" }<!>
+val l6: String.(Int) -> String = { x: String, y -> "" }
 
 val w1: String.(Int, Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>fun(x: String, y) = ""<!>
 val w2: String.(Int, Int) -> String = <!INITIALIZER_TYPE_MISMATCH!>{ x: String, y -> "" }<!>