K2: Apply constraints from forks at UNTIL_FIRST_LAMBDA completion phase

Otherwise, OVERLOAD_RESOLUTION_AMBIGUITY is reported for the calls to
the functions annotated as @OverloadResolutionByLambdaReturnType

^KT-43296 In Progress

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -13946,6 +13946,18 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/inference/inferInFunctionLiteralsWithReturn.kt");
             }
 
+            @Test
+            @TestMetadata("inferenceForkRegression.kt")
+            public void testInferenceForkRegression() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/inferenceForkRegression.kt");
+            }
+
+            @Test
+            @TestMetadata("inferenceForkRegressionSimple.kt")
+            public void testInferenceForkRegressionSimple() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/inferenceForkRegressionSimple.kt");
+            }
+
             @Test
             @TestMetadata("intersectionTypeMultipleBoundsAsReceiver.kt")
             public void testIntersectionTypeMultipleBoundsAsReceiver() throws Exception {
@@ -15972,6 +15984,22 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 }
             }
 
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/inference/forks")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Forks {
+                @Test
+                public void testAllFilesPresentInForks() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/inference/forks"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                }
+
+                @Test
+                @TestMetadata("overloadResolutionByLambdaReturnTypeAndExpectedType.kt")
+                public void testOverloadResolutionByLambdaReturnTypeAndExpectedType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/forks/overloadResolutionByLambdaReturnTypeAndExpectedType.kt");
+                }
+            }
+
             @Nested
             @TestMetadata("compiler/testData/diagnostics/tests/inference/nestedCalls")
             @TestDataPath("$PROJECT_ROOT")

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

@@ -13946,6 +13946,18 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/inference/inferInFunctionLiteralsWithReturn.kt");
             }
 
+            @Test
+            @TestMetadata("inferenceForkRegression.kt")
+            public void testInferenceForkRegression() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/inferenceForkRegression.kt");
+            }
+
+            @Test
+            @TestMetadata("inferenceForkRegressionSimple.kt")
+            public void testInferenceForkRegressionSimple() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/inferenceForkRegressionSimple.kt");
+            }
+
             @Test
             @TestMetadata("intersectionTypeMultipleBoundsAsReceiver.kt")
             public void testIntersectionTypeMultipleBoundsAsReceiver() throws Exception {
@@ -15972,6 +15984,22 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 }
             }
 
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/inference/forks")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Forks {
+                @Test
+                public void testAllFilesPresentInForks() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/inference/forks"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                }
+
+                @Test
+                @TestMetadata("overloadResolutionByLambdaReturnTypeAndExpectedType.kt")
+                public void testOverloadResolutionByLambdaReturnTypeAndExpectedType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/forks/overloadResolutionByLambdaReturnTypeAndExpectedType.kt");
+                }
+            }
+
             @Nested
             @TestMetadata("compiler/testData/diagnostics/tests/inference/nestedCalls")
             @TestDataPath("$PROJECT_ROOT")

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java

@@ -13946,6 +13946,18 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/inference/inferInFunctionLiteralsWithReturn.kt");
             }
 
+            @Test
+            @TestMetadata("inferenceForkRegression.kt")
+            public void testInferenceForkRegression() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/inferenceForkRegression.kt");
+            }
+
+            @Test
+            @TestMetadata("inferenceForkRegressionSimple.kt")
+            public void testInferenceForkRegressionSimple() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/inferenceForkRegressionSimple.kt");
+            }
+
             @Test
             @TestMetadata("intersectionTypeMultipleBoundsAsReceiver.kt")
             public void testIntersectionTypeMultipleBoundsAsReceiver() throws Exception {
@@ -15972,6 +15984,22 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 }
             }
 
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/inference/forks")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Forks {
+                @Test
+                public void testAllFilesPresentInForks() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/inference/forks"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                }
+
+                @Test
+                @TestMetadata("overloadResolutionByLambdaReturnTypeAndExpectedType.kt")
+                public void testOverloadResolutionByLambdaReturnTypeAndExpectedType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/forks/overloadResolutionByLambdaReturnTypeAndExpectedType.kt");
+                }
+            }
+
             @Nested
             @TestMetadata("compiler/testData/diagnostics/tests/inference/nestedCalls")
             @TestDataPath("$PROJECT_ROOT")

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

@@ -62,7 +62,7 @@ class ConstraintSystemCompleter(components: BodyResolveComponents, private val c
     ) {
         val topLevelTypeVariables = topLevelType.extractTypeVariables()
 
-        if (completionMode == ConstraintSystemCompletionMode.FULL) {
+        if (completionMode == ConstraintSystemCompletionMode.FULL || completionMode == ConstraintSystemCompletionMode.UNTIL_FIRST_LAMBDA) {
             resolveForkPointsConstraints()
         }
 

compiler/testData/diagnostics/tests/inference/forks/overloadResolutionByLambdaReturnTypeAndExpectedType.fir.kt

@@ -0,0 +1,39 @@
+// SKIP_TXT
+// WITH_STDLIB
+import kotlin.experimental.ExperimentalTypeInference
+
+interface MyList<E>
+interface MySequence<E>
+
+fun <E> myListOf(e: E): MyList<E> = TODO()
+
+interface C : MyList<String>
+
+fun <E> foo(m: MyList<E>, c: C) {
+    if (c === m) {
+        val x1: MyList<String> = m.noOverloadResolutionByLambdaReturnType { x ->
+            myListOf(x)
+        } // ok in K1 and K2
+
+        val x2: MyList<E> = m.noOverloadResolutionByLambdaReturnType { x ->
+            myListOf(x)
+        } // ok in K2, error in k1
+
+        val y1: MyList<String> = m.limitedFlatMap { x ->
+            myListOf(x)
+        } // ok in K1 and K2
+
+        val y2: MyList<E> = <!INITIALIZER_TYPE_MISMATCH, TYPE_MISMATCH!>m.limitedFlatMap { x ->
+            myListOf(x)
+        }<!> // error in K1 and K2
+    }
+}
+
+fun <T> MyList<T>.noOverloadResolutionByLambdaReturnType(producer: (T) -> MyList<T>): MyList<T> = TODO()
+
+fun <T> MyList<T>.limitedFlatMap(producer: (T) -> MyList<T>): MyList<T> = TODO()
+
+@OptIn(ExperimentalTypeInference::class)
+@OverloadResolutionByLambdaReturnType
+@kotlin.jvm.JvmName("limitedFlatMapSeq")
+fun <T> MyList<T>.limitedFlatMap(producer: (T) -> MySequence<T>): MyList<T> = TODO()

compiler/testData/diagnostics/tests/inference/forks/overloadResolutionByLambdaReturnTypeAndExpectedType.kt

@@ -0,0 +1,39 @@
+// SKIP_TXT
+// WITH_STDLIB
+import kotlin.experimental.ExperimentalTypeInference
+
+interface MyList<E>
+interface MySequence<E>
+
+fun <E> myListOf(e: E): MyList<E> = TODO()
+
+interface C : MyList<String>
+
+fun <E> foo(m: MyList<E>, c: C) {
+    if (c === m) {
+        val x1: MyList<String> = <!DEBUG_INFO_SMARTCAST!>m<!>.noOverloadResolutionByLambdaReturnType { x ->
+            myListOf(x)
+        } // ok in K1 and K2
+
+        val x2: MyList<E> = <!TYPE_MISMATCH!><!DEBUG_INFO_SMARTCAST!>m<!>.<!TYPE_MISMATCH, TYPE_MISMATCH!>noOverloadResolutionByLambdaReturnType { x ->
+            myListOf(x)
+        }<!><!> // ok in K2, error in k1
+
+        val y1: MyList<String> = <!DEBUG_INFO_SMARTCAST!>m<!>.limitedFlatMap { x ->
+            myListOf(x)
+        } // ok in K1 and K2
+
+        val y2: MyList<E> = <!TYPE_MISMATCH!><!DEBUG_INFO_SMARTCAST!>m<!>.<!TYPE_MISMATCH!>limitedFlatMap { x ->
+            myListOf(x)
+        }<!><!> // error in K1 and K2
+    }
+}
+
+fun <T> MyList<T>.noOverloadResolutionByLambdaReturnType(producer: (T) -> MyList<T>): MyList<T> = TODO()
+
+fun <T> MyList<T>.limitedFlatMap(producer: (T) -> MyList<T>): MyList<T> = TODO()
+
+@OptIn(ExperimentalTypeInference::class)
+@OverloadResolutionByLambdaReturnType
+@kotlin.jvm.JvmName("limitedFlatMapSeq")
+fun <T> MyList<T>.limitedFlatMap(producer: (T) -> MySequence<T>): MyList<T> = TODO()

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

@@ -0,0 +1,26 @@
+// SKIP_TXT
+// FIR_IDENTICAL
+// FULL_JDK
+// WITH_STDLIB
+
+import java.util.concurrent.ConcurrentHashMap
+
+interface ModificationData
+
+fun main() {
+    val updatedFiles = ConcurrentHashMap<String, ModificationData>()
+
+    // HashMap effectively have two Map supertypes: one like Map<K, V> and one Map<K!, V!> (through AbstractMap)
+    // But here, since `updatedFiles` has only a supertype Map<String!, ModificationData!>
+    // `updatedFilesSnapshot` has two supertypes relevant to Map: Map<String!, ModificationData!>  and one Map<String!, ModificationData!>
+    // Actually, we might collapse them, because they're equivalent, but for now we don't
+    val updatedFilesSnapshot = HashMap(updatedFiles)
+
+    // Thus, it leads to constraint system fork during `flatMap` inference
+    // And during overload resolution between two `flatMap` versions with @OverloadResolutionByLambdaReturnType we have to run lambda analysis
+    // So, we have to apply forks to the system on this completion phase too, but still not for the PARTIAL completion mode
+    // Otherwise, OVERLOAD_RESOLUTION_AMBIGUITY happens
+    updatedFilesSnapshot.flatMap { _ ->
+        listOf("")
+    }
+}

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

@@ -0,0 +1,21 @@
+// SKIP_TXT
+// FULL_JDK
+// WITH_STDLIB
+
+interface C : MutableMap<String, Int>
+
+fun <K, V> foo(m: MutableMap<K, V>, c: C) {
+    if (c === m) {
+        // `m` has a type C & MutableMap<K, V>
+        // So it has two instances of Map supertypes: Map<String, Int> and Map<K, V>
+        //
+        // Thus, it leads to constraint system fork during `flatMap` inference
+        // And during overload resolution between two `flatMap` versions with @OverloadResolutionByLambdaReturnType we have to run lambda analysis
+        // So, we have to apply forks to the system on this completion phase too, so we would have enough information for the input types of lambda
+        // Otherwise, OVERLOAD_RESOLUTION_AMBIGUITY happens
+        // But we don't do it for PARTIAL completion mode still
+        c.flatMap { _ ->
+            listOf("")
+        }
+    }
+}

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

@@ -0,0 +1,21 @@
+// SKIP_TXT
+// FULL_JDK
+// WITH_STDLIB
+
+interface C : MutableMap<String, Int>
+
+fun <K, V> foo(m: MutableMap<K, V>, c: C) {
+    if (c === m) {
+        // `m` has a type C & MutableMap<K, V>
+        // So it has two instances of Map supertypes: Map<String, Int> and Map<K, V>
+        //
+        // Thus, it leads to constraint system fork during `flatMap` inference
+        // And during overload resolution between two `flatMap` versions with @OverloadResolutionByLambdaReturnType we have to run lambda analysis
+        // So, we have to apply forks to the system on this completion phase too, so we would have enough information for the input types of lambda
+        // Otherwise, OVERLOAD_RESOLUTION_AMBIGUITY happens
+        // But we don't do it for PARTIAL completion mode still
+        <!DEBUG_INFO_SMARTCAST!>c<!>.flatMap { _ ->
+            listOf("")
+        }
+    }
+}

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

@@ -13952,6 +13952,18 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/inference/inferInFunctionLiteralsWithReturn.kt");
             }
 
+            @Test
+            @TestMetadata("inferenceForkRegression.kt")
+            public void testInferenceForkRegression() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/inferenceForkRegression.kt");
+            }
+
+            @Test
+            @TestMetadata("inferenceForkRegressionSimple.kt")
+            public void testInferenceForkRegressionSimple() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/inferenceForkRegressionSimple.kt");
+            }
+
             @Test
             @TestMetadata("intersectionTypeMultipleBoundsAsReceiver.kt")
             public void testIntersectionTypeMultipleBoundsAsReceiver() throws Exception {
@@ -15978,6 +15990,22 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 }
             }
 
+            @Nested
+            @TestMetadata("compiler/testData/diagnostics/tests/inference/forks")
+            @TestDataPath("$PROJECT_ROOT")
+            public class Forks {
+                @Test
+                public void testAllFilesPresentInForks() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/inference/forks"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
+                }
+
+                @Test
+                @TestMetadata("overloadResolutionByLambdaReturnTypeAndExpectedType.kt")
+                public void testOverloadResolutionByLambdaReturnTypeAndExpectedType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/forks/overloadResolutionByLambdaReturnTypeAndExpectedType.kt");
+                }
+            }
+
             @Nested
             @TestMetadata("compiler/testData/diagnostics/tests/inference/nestedCalls")
             @TestDataPath("$PROJECT_ROOT")