[NI] Fix detecting expected type variable for postponed atoms

#KT-36819 Fixed

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirOldFrontendDiagnosticsTestGenerated.java

@@ -10101,6 +10101,11 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirOldFronte
             runTest("compiler/testData/diagnostics/tests/inference/kt35702.kt");
         }
 
+        @TestMetadata("kt36819.kt")
+        public void testKt36819() throws Exception {
+            runTest("compiler/testData/diagnostics/tests/inference/kt36819.kt");
+        }
+
         @TestMetadata("kt6175.kt")
         public void testKt6175() throws Exception {
             runTest("compiler/testData/diagnostics/tests/inference/kt6175.kt");

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

@@ -167,9 +167,11 @@ class KotlinConstraintSystemCompleter(
 
         val builtIns = (variable as TypeVariableTypeConstructor).builtIns
         val csBuilder = (c as NewConstraintSystemImpl).getBuilder()
+
+        val expectedTypeVariable = postponedAtom.expectedType?.constructor?.takeIf { it in c.allTypeVariables } ?: variable
         val atomToAnalyze = when (postponedAtom) {
             is PostponedCallableReferenceAtom -> postponedAtom.preparePostponedAtomWithTypeVariableAsExpectedType(
-                c, csBuilder, variable,
+                c, csBuilder, expectedTypeVariable,
                 condition = { true },
                 isSuitable = KotlinType::isBuiltinFunctionalTypeOrSubtype,
                 typeVariableCreator = { TypeVariableForCallableReferenceReturnType(builtIns, "_Q") },
@@ -180,7 +182,7 @@ class KotlinConstraintSystemCompleter(
                 }
             )
             is LambdaWithTypeVariableAsExpectedTypeAtom -> postponedAtom.preparePostponedAtomWithTypeVariableAsExpectedType(
-                c, csBuilder, variable,
+                c, csBuilder, expectedTypeVariable,
                 condition = { it.atom.parametersTypes?.all { type -> type != null } != true },
                 isSuitable = KotlinType::isBuiltinFunctionalType,
                 typeVariableCreator = { TypeVariableForLambdaReturnType(postponedAtom.atom, builtIns, "_R") },

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

@@ -0,0 +1,18 @@
+// !LANGUAGE: +NewInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE
+// ISSUE: KT-36819
+
+fun <K> select(vararg x: K) = x[0]
+interface A
+class B: A
+class C: A
+fun <T> id1(x: T): T = x
+fun <R> id2(x: R): R = x
+
+class Out<out R>(x: R)
+
+fun main() {
+    val x1 = select(id1 { B() }, id2 { C() })
+    val x2 = select({ B() }, { C() }) // OK, CST = () -> A
+    val x3 = select(id1(Out(B())), id2(Out(C()))) // OK, CST = Out<A>
+}

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

@@ -0,0 +1,18 @@
+// !LANGUAGE: +NewInference
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE
+// ISSUE: KT-36819
+
+fun <K> select(vararg x: K) = x[0]
+interface A
+class B: A
+class C: A
+fun <T> id1(x: T): T = x
+fun <R> id2(x: R): R = x
+
+class Out<out R>(x: R)
+
+fun main() {
+    val x1 = select(id1 { B() }, id2 { <!TYPE_MISMATCH, TYPE_MISMATCH!>C()<!> })
+    val x2 = select({ B() }, { C() }) // OK, CST = () -> A
+    val x3 = select(id1(Out(B())), id2(Out(C()))) // OK, CST = Out<A>
+}

compiler/testData/diagnostics/tests/inference/kt36819.txt

@@ -0,0 +1,33 @@
+package
+
+public fun </*0*/ T> id1(/*0*/ x: T): T
+public fun </*0*/ R> id2(/*0*/ x: R): R
+public fun main(): kotlin.Unit
+public fun </*0*/ K> select(/*0*/ vararg x: K /*kotlin.Array<out K>*/): K
+
+public interface A {
+    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 B : A {
+    public constructor B()
+    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 C : A {
+    public constructor C()
+    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 Out</*0*/ out R> {
+    public constructor Out</*0*/ out R>(/*0*/ x: R)
+    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
+}

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

@@ -10108,6 +10108,11 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 runTest("compiler/testData/diagnostics/tests/inference/kt35702.kt");
             }
 
+            @TestMetadata("kt36819.kt")
+            public void testKt36819() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/kt36819.kt");
+            }
+
             @TestMetadata("kt6175.kt")
             public void testKt6175() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/inference/kt6175.kt");

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

@@ -10103,6 +10103,11 @@ public class DiagnosticsUsingJavacTestGenerated extends AbstractDiagnosticsUsing
                 runTest("compiler/testData/diagnostics/tests/inference/kt35702.kt");
             }
 
+            @TestMetadata("kt36819.kt")
+            public void testKt36819() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/kt36819.kt");
+            }
+
             @TestMetadata("kt6175.kt")
             public void testKt6175() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/inference/kt6175.kt");