[NI] Fix inference of lambda with receiver when wrapped...

in when/if or another lambda.
#KT-37419 fixed

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirOldFrontendDiagnosticsTestGenerated.java

@@ -11478,6 +11478,11 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirOldFronte
                 runTest("compiler/testData/diagnostics/tests/inference/regressions/kt37043_2.kt");
             }
 
+            @TestMetadata("kt37419.kt")
+            public void testKt37419() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/regressions/kt37419.kt");
+            }
+
             @TestMetadata("kt37650.kt")
             public void testKt37650() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/inference/regressions/kt37650.kt");

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/components/KotlinConstraintSystemCompleter.kt

@@ -263,9 +263,12 @@ class KotlinConstraintSystemCompleter(
             c.notFixedTypeVariables.getValue(variable),
             TypeVariableDirectionCalculator.ResolveDirection.TO_SUPERTYPE
         ) as KotlinType
-        val isExtensionWithoutParameters =
-            functionalType.isExtensionFunctionType && functionalType.arguments.size == 2 && parameterTypes?.isEmpty() == true
-        if (parameterTypes?.all { type -> type != null } == true && !isExtensionWithoutParameters) return this
+
+        val isExtensionFunction = functionalType.isExtensionFunctionType
+        val isExtensionFunctionWithReceiverAsDeclaredParameter =
+            isExtensionFunction && functionalType.arguments.size - 1 == parameterTypes?.count { it != null }
+        if (parameterTypes?.all { it != null } == true && (!isExtensionFunction || isExtensionFunctionWithReceiverAsDeclaredParameter)) return this
+
         if (!functionalType.isSuitable()) return this
         val returnVariable = typeVariableCreator()
         csBuilder.registerVariable(returnVariable)

compiler/testData/diagnostics/tests/inference/regressions/kt37419.fir.kt

@@ -0,0 +1,45 @@
+// !LANGUAGE: +NewInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE
+
+interface Receiver
+interface Parameter
+typealias LambdaWithReceiver = Receiver.(Parameter) -> Unit
+
+fun Receiver.method(param: Parameter): LambdaWithReceiver = TODO()
+
+enum class E { VALUE }
+
+fun <K> id(x: K): K = x
+
+class SomeClass {
+    val e = E.VALUE
+
+    val withoutType: LambdaWithReceiver
+        get() = when (e) {
+            E.VALUE -> { param ->
+                <!INAPPLICABLE_CANDIDATE!>method<!>(param)
+            }
+        }
+
+    val withExplicitType: LambdaWithReceiver
+        get() = when (e) {
+            E.VALUE -> { param: Parameter ->
+                <!INAPPLICABLE_CANDIDATE!>method<!>(param)
+            }
+        }
+}
+
+class OtherClass {
+    val ok: LambdaWithReceiver
+        get() = { param: Parameter ->
+            method(param)
+        }
+}
+
+val e2 = E.VALUE
+val staticWithExplicitType: LambdaWithReceiver
+    get() = when (e2) {
+        E.VALUE -> { param: Parameter ->
+            <!UNRESOLVED_REFERENCE!>method<!>(param)
+        }
+    }

compiler/testData/diagnostics/tests/inference/regressions/kt37419.kt

@@ -0,0 +1,45 @@
+// !LANGUAGE: +NewInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE
+
+interface Receiver
+interface Parameter
+typealias LambdaWithReceiver = Receiver.(Parameter) -> Unit
+
+fun Receiver.method(param: Parameter): LambdaWithReceiver = TODO()
+
+enum class E { VALUE }
+
+fun <K> id(x: K): K = x
+
+class SomeClass {
+    val e = E.VALUE
+
+    val withoutType: LambdaWithReceiver
+        get() = when (e) {
+            E.VALUE -> { param ->
+                method(param)
+            }
+        }
+
+    val withExplicitType: LambdaWithReceiver
+        get() = when (e) {
+            E.VALUE -> { param: Parameter ->
+                method(param)
+            }
+        }
+}
+
+class OtherClass {
+    val ok: LambdaWithReceiver
+        get() = { param: Parameter ->
+            method(param)
+        }
+}
+
+val e2 = E.VALUE
+val staticWithExplicitType: LambdaWithReceiver
+    get() = when (e2) {
+        E.VALUE -> { param: Parameter ->
+            method(param)
+        }
+    }

compiler/testData/diagnostics/tests/inference/regressions/kt37419.txt

@@ -0,0 +1,56 @@
+package
+
+public val e2: E
+public val staticWithExplicitType: LambdaWithReceiver /* = Receiver.(Parameter) -> kotlin.Unit */
+public fun </*0*/ K> id(/*0*/ x: K): K
+public fun Receiver.method(/*0*/ param: Parameter): LambdaWithReceiver /* = Receiver.(Parameter) -> kotlin.Unit */
+
+public final enum class E : kotlin.Enum<E> {
+    enum entry VALUE
+
+    private constructor E()
+    public final override /*1*/ /*fake_override*/ val name: kotlin.String
+    public final override /*1*/ /*fake_override*/ val ordinal: kotlin.Int
+    protected final override /*1*/ /*fake_override*/ fun clone(): kotlin.Any
+    public final override /*1*/ /*fake_override*/ fun compareTo(/*0*/ other: E): kotlin.Int
+    public final override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    protected/*protected and package*/ final override /*1*/ /*fake_override*/ /*isHiddenForResolutionEverywhereBesideSupercalls*/ fun finalize(): kotlin.Unit
+    public final override /*1*/ /*fake_override*/ /*isHiddenForResolutionEverywhereBesideSupercalls*/ fun getDeclaringClass(): java.lang.Class<E!>!
+    public final override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+
+    // Static members
+    public final /*synthesized*/ fun valueOf(/*0*/ value: kotlin.String): E
+    public final /*synthesized*/ fun values(): kotlin.Array<E>
+}
+
+public final class OtherClass {
+    public constructor OtherClass()
+    public final val ok: LambdaWithReceiver /* = Receiver.(Parameter) -> kotlin.Unit */
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public interface Parameter {
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public interface Receiver {
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class SomeClass {
+    public constructor SomeClass()
+    public final val e: E
+    public final val withExplicitType: LambdaWithReceiver /* = Receiver.(Parameter) -> kotlin.Unit */
+    public final val withoutType: LambdaWithReceiver /* = Receiver.(Parameter) -> kotlin.Unit */
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+public typealias LambdaWithReceiver = Receiver.(Parameter) -> kotlin.Unit

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

@@ -11,6 +11,8 @@ fun test() {
     val f02: Sample.() -> Unit = id<Sample.() -> Unit> { s: Sample -> }
 }
 
+enum class E { VALUE }
+
 typealias E0 = Int.() -> Int
 class W1(val f: E0) {
     // overload ambiguity is not supported yet - see commented examples with "overload" keyword below
@@ -68,15 +70,18 @@ fun test2() { // to extension lambda 1
     val w23 = W2 { s -> s.length } // oi+ ni+
     val i23: E1 = id { s -> s.<!UNRESOLVED_REFERENCE!>length<!> } // oi+ ni+
     val w24 = W2 { s: String -> this + s.length } // oi+ ni+
-//    val i24: E1 = id { s: String -> this + s.length } //oi- ni-
+    val i24: E1 = id { s: String -> this + s.length } //oi- ni+
     val w25 = W2 { s: String -> s.length } // oi+ ni+
-//    val i25: E1 = id { s: String -> s.length } // oi- ni-
-
-    // yet unsupported cases with ambiguity for the lambda conversion (commented constructors in wrappers above)
-//    val w26 = W2 { i, s -> i + s.length } // overload oi- ni-
-//    val i26: E1 = id { i, s -> i + s.length } // overload oi- ni-
-//    val w27 = W2 { i, s: String -> i + s.length } // overload oi- ni-
-//    val i27: E1 = id { i, s: String -> i + s.length } // overload oi- ni-
+    val i25: E1 = id { s: String -> s.length } // oi- ni+
+    val w26 = W2(id { s: String -> this + s.length }) // oi- ni+
+    val w26a = W2(id { s -> this + s.<!UNRESOLVED_REFERENCE!>length<!> }) // oi+ ni+
+    val i26: E1 = id { s: String -> this + s.length } // oi- ni+
+    val i26a: E1 = id { s -> this + s.<!UNRESOLVED_REFERENCE!>length<!> } // oi+ ni+
+    val e = E.VALUE
+    val w27 = <!INAPPLICABLE_CANDIDATE!>W2<!>(when (e) { E.VALUE ->  { s: String -> this + s.length } }) // oi- ni+
+    val w27a = <!INAPPLICABLE_CANDIDATE!>W2<!>(when (e) { E.VALUE ->  { s -> this + s.<!UNRESOLVED_REFERENCE!>length<!> } }) // oi+ ni+
+    val i27: E1 = when (e) { E.VALUE ->  { s: String -> this + s.length } } // oi+ ni+
+    val i27a: E1 = when (e) { E.VALUE ->  { s -> this + s.<!UNRESOLVED_REFERENCE!>length<!> } } // oi+ ni+
 
     val w28 = <!INAPPLICABLE_CANDIDATE!>W2<!> { i: Int, s -> i <!AMBIGUITY!>+<!> s.<!UNRESOLVED_REFERENCE!>length<!> } // oi- ni-
     val i28: E1 = id { i: Int, s -> i + s.length } // oi- ni-
@@ -84,8 +89,14 @@ fun test2() { // to extension lambda 1
     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 o2a = W2 { i: Int -> i } // overload oi- ni+
-//    val o2b = W2 { i -> i } // overload oi- ni-
+//    val w2a = W2 { i, s -> i + s.length } // overload oi- ni-
+//    val i2a: E1 = id { i, s -> i + s.length } // overload oi- ni-
+//    val w2b = W2 { i, s: String -> i + s.length } // overload oi- ni-
+//    val i2b: E1 = id { i, s: String -> i + s.length } // overload oi- ni-
+
+    // yet unsupported cases with ambiguity for the lambda conversion (commented constructors in wrappers above)
+//    val o2c = W2 { i: Int -> i } // overload oi- ni+
+//    val o2d = W2 { i -> i } // overload oi- ni-
 }
 
 fun test3() { // to non-extension lambda 1

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

@@ -11,6 +11,8 @@ fun test() {
     val f02: Sample.() -> Unit = id<Sample.() -> Unit> { s: Sample -> }
 }
 
+enum class E { VALUE }
+
 typealias E0 = Int.() -> Int
 class W1(val f: E0) {
     // overload ambiguity is not supported yet - see commented examples with "overload" keyword below
@@ -68,15 +70,18 @@ fun test2() { // to extension lambda 1
     val w23 = W2 { s -> s.length } // oi+ ni+
     val i23: E1 = id { s -> s.length } // oi+ ni+
     val w24 = W2 { s: String -> this + s.length } // oi+ ni+
-//    val i24: E1 = id { s: String -> this + s.length } //oi- ni-
+    val i24: E1 = id { s: String -> this + s.length } //oi- ni+
     val w25 = W2 { s: String -> s.length } // oi+ ni+
-//    val i25: E1 = id { s: String -> s.length } // oi- ni-
-
-    // yet unsupported cases with ambiguity for the lambda conversion (commented constructors in wrappers above)
-//    val w26 = W2 { i, s -> i + s.length } // overload oi- ni-
-//    val i26: E1 = id { i, s -> i + s.length } // overload oi- ni-
-//    val w27 = W2 { i, s: String -> i + s.length } // overload oi- ni-
-//    val i27: E1 = id { i, s: String -> i + s.length } // overload oi- ni-
+    val i25: E1 = id { s: String -> s.length } // oi- ni+
+    val w26 = W2(id { s: String -> this + s.length }) // oi- ni+
+    val w26a = W2(id { s -> this + s.length }) // oi+ ni+
+    val i26: E1 = id { s: String -> this + s.length } // oi- ni+
+    val i26a: E1 = id { s -> this + s.length } // oi+ ni+
+    val e = E.VALUE
+    val w27 = W2(when (e) { E.VALUE ->  { s: String -> this + s.length } }) // oi- ni+
+    val w27a = W2(when (e) { E.VALUE ->  { s -> this + s.length } }) // oi+ ni+
+    val i27: E1 = when (e) { E.VALUE ->  { s: String -> this + s.length } } // oi+ ni+
+    val i27a: E1 = when (e) { E.VALUE ->  { s -> this + s.length } } // oi+ ni+
 
     val w28 = W2 <!TYPE_MISMATCH!>{ <!EXPECTED_PARAMETERS_NUMBER_MISMATCH!><!EXPECTED_PARAMETER_TYPE_MISMATCH, EXPECTED_PARAMETER_TYPE_MISMATCH!>i: Int<!>, <!CANNOT_INFER_PARAMETER_TYPE!>s<!><!> -> i + <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE!>s<!>.<!DEBUG_INFO_MISSING_UNRESOLVED!>length<!> }<!> // oi- ni-
     val i28: E1 = <!TYPE_MISMATCH!>id <!TYPE_MISMATCH, TYPE_MISMATCH!>{ <!EXPECTED_PARAMETERS_NUMBER_MISMATCH!><!EXPECTED_PARAMETER_TYPE_MISMATCH, EXPECTED_PARAMETER_TYPE_MISMATCH!>i: Int<!>, <!CANNOT_INFER_PARAMETER_TYPE!>s<!><!> -> i + <!DEBUG_INFO_ELEMENT_WITH_ERROR_TYPE!>s<!>.<!DEBUG_INFO_MISSING_UNRESOLVED!>length<!> }<!><!> // oi- ni-
@@ -84,8 +89,14 @@ fun test2() { // to extension lambda 1
     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 o2a = W2 { i: Int -> i } // overload oi- ni+
-//    val o2b = W2 { i -> i } // overload oi- ni-
+//    val w2a = W2 { i, s -> i + s.length } // overload oi- ni-
+//    val i2a: E1 = id { i, s -> i + s.length } // overload oi- ni-
+//    val w2b = W2 { i, s: String -> i + s.length } // overload oi- ni-
+//    val i2b: E1 = id { i, s: String -> i + s.length } // overload oi- ni-
+
+    // yet unsupported cases with ambiguity for the lambda conversion (commented constructors in wrappers above)
+//    val o2c = W2 { i: Int -> i } // overload oi- ni+
+//    val o2d = W2 { i -> i } // overload oi- ni-
 }
 
 fun test3() { // to non-extension lambda 1

compiler/testData/diagnostics/tests/regressions/kt30245.txt

@@ -22,6 +22,25 @@ public final class B : A {
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
 }
 
+public final enum class E : kotlin.Enum<E> {
+    enum entry VALUE
+
+    private constructor E()
+    public final override /*1*/ /*fake_override*/ val name: kotlin.String
+    public final override /*1*/ /*fake_override*/ val ordinal: kotlin.Int
+    protected final override /*1*/ /*fake_override*/ fun clone(): kotlin.Any
+    public final override /*1*/ /*fake_override*/ fun compareTo(/*0*/ other: E): kotlin.Int
+    public final override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    protected/*protected and package*/ final override /*1*/ /*fake_override*/ /*isHiddenForResolutionEverywhereBesideSupercalls*/ fun finalize(): kotlin.Unit
+    public final override /*1*/ /*fake_override*/ /*isHiddenForResolutionEverywhereBesideSupercalls*/ fun getDeclaringClass(): java.lang.Class<E!>!
+    public final override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+
+    // Static members
+    public final /*synthesized*/ fun valueOf(/*0*/ value: kotlin.String): E
+    public final /*synthesized*/ fun values(): kotlin.Array<E>
+}
+
 public final class Sample {
     public constructor Sample()
     public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -11485,6 +11485,11 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTestWithFirVali
                     runTest("compiler/testData/diagnostics/tests/inference/regressions/kt37043_2.kt");
                 }
 
+                @TestMetadata("kt37419.kt")
+                public void testKt37419() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/regressions/kt37419.kt");
+                }
+
                 @TestMetadata("kt37650.kt")
                 public void testKt37650() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/regressions/kt37650.kt");

compiler/tests/org/jetbrains/kotlin/checkers/javac/DiagnosticsUsingJavacTestGenerated.java

@@ -11480,6 +11480,11 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                     runTest("compiler/testData/diagnostics/tests/inference/regressions/kt37043_2.kt");
                 }
 
+                @TestMetadata("kt37419.kt")
+                public void testKt37419() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/regressions/kt37419.kt");
+                }
+
                 @TestMetadata("kt37650.kt")
                 public void testKt37650() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/inference/regressions/kt37650.kt");