K2: Do not force completion for 'when' branches with expected type

Otherwise, it leads to branches inference run fully independent,
while there are cases when it's necessary to flow type information from
one of the branch to another (see the new test).

NB. In K1, it worked differently: if branches were inferred altogether
only for Any/Any? expect types (otherwise they're analyzed independently)
See foo2/foo4 in the test.

To avoid breaking change we need to support foo1/foo3, but we're trying
not to have some special rule for Any, so we've got a new resolution mode
that provides expect type, but doesn't require full completion.

^KT-45989 Fixed
^KT-56563 Fixed
^KT-54709 Related

For change in specialCallWithMaterializeAndExpectedType.kt
At first, see at KT-36776

Long time ago, it's been decided that if/when resolution
should look similar to similar "select()" calls,
but it's a breaking change (see KT-36776), and we were ready for that back then.

But then, there were too many broken cases found, thus we reverted it at
100a6f70ca

But probably, it would be better to try to infer `String?`
instead of `Nothing?` (see next commits)

Note that change in specialCallWithMaterializeAndExpectedType.kt
will be addressed in later commits, too

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

@@ -14168,6 +14168,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/inference/hasErrorInConstrainingTypes.kt");
             }
 
+            @Test
+            @TestMetadata("ifWithDependentBranches.kt")
+            public void testIfWithDependentBranches() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/ifWithDependentBranches.kt");
+            }
+
             @Test
             @TestMetadata("illegalUnderscoredTypeArgument.kt")
             public void testIllegalUnderscoredTypeArgument() throws Exception {

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

@@ -14174,6 +14174,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/inference/hasErrorInConstrainingTypes.kt");
             }
 
+            @Test
+            @TestMetadata("ifWithDependentBranches.kt")
+            public void testIfWithDependentBranches() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/ifWithDependentBranches.kt");
+            }
+
             @Test
             @TestMetadata("illegalUnderscoredTypeArgument.kt")
             public void testIllegalUnderscoredTypeArgument() throws Exception {

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

@@ -14168,6 +14168,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/inference/hasErrorInConstrainingTypes.kt");
             }
 
+            @Test
+            @TestMetadata("ifWithDependentBranches.kt")
+            public void testIfWithDependentBranches() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/ifWithDependentBranches.kt");
+            }
+
             @Test
             @TestMetadata("illegalUnderscoredTypeArgument.kt")
             public void testIllegalUnderscoredTypeArgument() throws Exception {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolutionMode.kt

@@ -41,7 +41,18 @@ sealed class ResolutionMode(val forceFullCompletion: Boolean) {
         // In these examples we should try using the property type information while resolving the initializer,
         // but it's ok if it's not applicable
         val shouldBeStrictlyEnforced: Boolean = true,
-    ) : ResolutionMode(forceFullCompletion = true) {
+        // Currently the only case for expected type when we don't force completion are when's branches
+        forceFullCompletion: Boolean = true,
+    ) : ResolutionMode(forceFullCompletion) {
+
+        fun copy(
+            mayBeCoercionToUnitApplied: Boolean = this.mayBeCoercionToUnitApplied,
+            forceFullCompletion: Boolean = this.forceFullCompletion
+        ): WithExpectedType = WithExpectedType(
+            expectedTypeRef, mayBeCoercionToUnitApplied, expectedTypeMismatchIsReportedInChecker, fromCast, shouldBeStrictlyEnforced,
+            forceFullCompletion
+        )
+
         override fun toString(): String {
             return "WithExpectedType: ${expectedTypeRef.prettyString()}, " +
                     "mayBeCoercionToUnitApplied=${mayBeCoercionToUnitApplied}, " +

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

@@ -74,12 +74,15 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirAbstractBodyRes
                         // when with one branch cannot be completed if it's not already complete in the first place
                     }
                     else -> {
+                        val resolutionModeForBranches =
+                            (data as? ResolutionMode.WithExpectedType)
+                                // Currently we don't use information from cast, but probably we could have
+                                ?.takeUnless { it.fromCast }
+                                ?.copy(forceFullCompletion = false)
+                                ?: ResolutionMode.ContextDependent
                         whenExpression = whenExpression.transformBranches(
                             transformer,
-                            data.takeIf {
-                                val expectedType = it.expectedType
-                                expectedType != null && expectedType !is FirImplicitTypeRef
-                            } ?: ResolutionMode.ContextDependent,
+                            resolutionModeForBranches,
                         )
 
                         whenExpression = syntheticCallGenerator.generateCalleeForWhenExpression(whenExpression, resolutionContext) ?: run {

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

@@ -518,7 +518,7 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
             val value =
                 if (index == numberOfStatements - 1)
                     if (data is ResolutionMode.WithExpectedType)
-                        ResolutionMode.WithExpectedType(data.expectedTypeRef, mayBeCoercionToUnitApplied = true)
+                        data.copy(mayBeCoercionToUnitApplied = true)
                     else
                         data
                 else

compiler/testData/diagnostics/tests/controlStructures/ifElseIntersection.fir.kt

@@ -26,7 +26,7 @@ fun <T> bind2(r: Option<T>): Option<T> {
 fun <T, R> bind3(r: Option<T>): Option<T> {
     return <!RETURN_TYPE_MISMATCH!>if (r is Some) {
         // Diagnoses an error correctly
-        <!TYPE_MISMATCH!>if (true) <!TYPE_MISMATCH, TYPE_MISMATCH, TYPE_MISMATCH!>None<R>()<!> else r<!>
+        if (true) None<R>() else r
     }
     else r<!>
 }

compiler/testData/diagnostics/tests/generics/kt53656.fir.kt

@@ -5,7 +5,7 @@ interface Entity<T>
 abstract class SecuredEntity<out E>(val entity: E) where E : Entity<Int>, E : SecurityCodeAware<*,*>
 interface SecurityCodeAware<out E, R : SecuredEntity<E>> where E : Entity<Int>, E : SecurityCodeAware<E, R>
 fun <E, R : SecuredEntity<E>> SecurityCodeAware<E, R>.secured() : R where E : Entity<Int>, E : SecurityCodeAware<E, R> = <!RETURN_TYPE_MISMATCH!>when(this) {
-    is Order -> <!TYPE_MISMATCH!>SecuredOrder(this)<!>
+    is Order -> SecuredOrder(this)
     else -> null!!
 }<!>
 class Order : Entity<Int>

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

@@ -0,0 +1,62 @@
+// SKIP_TXT
+
+interface Additional
+interface A<T> : Additional
+
+fun <U> aOf(): A<U> = TODO()
+
+interface B<E>
+
+fun <F> B<F>.convert(): A<F> = TODO()
+
+fun foo1(x: B<String>): Any {
+    return if (x.hashCode() == 0) aOf() else x.convert()
+}
+
+fun foo2(x: B<String>): Additional {
+    return if (x.hashCode() == 0) aOf() else x.convert()
+}
+
+fun foo3(x: B<String>): Any {
+    return when {
+        x.hashCode() == 0 -> aOf()
+        else -> x.convert()
+    }
+}
+
+fun foo4(x: B<String>): Additional {
+    return when {
+        x.hashCode() == 0 -> aOf()
+        else -> x.convert()
+    }
+}
+
+fun foo5(x: B<String>): Any {
+    return if (x.hashCode() == 0) {
+        aOf()
+    } else {
+        x.convert()
+    }
+}
+
+fun foo6(x: B<String>): Additional {
+    return if (x.hashCode() == 0) {
+        aOf()
+    } else {
+        x.convert()
+    }
+}
+
+fun foo7(x: B<String>): Any {
+    return when {
+        x.hashCode() == 0 -> { aOf() }
+        else -> { x.convert() }
+    }
+}
+
+fun foo8(x: B<String>): Additional {
+    return when {
+        x.hashCode() == 0 -> { aOf() }
+        else -> { x.convert() }
+    }
+}

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

@@ -0,0 +1,62 @@
+// SKIP_TXT
+
+interface Additional
+interface A<T> : Additional
+
+fun <U> aOf(): A<U> = TODO()
+
+interface B<E>
+
+fun <F> B<F>.convert(): A<F> = TODO()
+
+fun foo1(x: B<String>): Any {
+    return if (x.hashCode() == 0) aOf() else x.convert()
+}
+
+fun foo2(x: B<String>): Additional {
+    return if (x.hashCode() == 0) <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>aOf<!>() else x.convert()
+}
+
+fun foo3(x: B<String>): Any {
+    return when {
+        x.hashCode() == 0 -> aOf()
+        else -> x.convert()
+    }
+}
+
+fun foo4(x: B<String>): Additional {
+    return when {
+        x.hashCode() == 0 -> <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>aOf<!>()
+        else -> x.convert()
+    }
+}
+
+fun foo5(x: B<String>): Any {
+    return if (x.hashCode() == 0) {
+        aOf()
+    } else {
+        x.convert()
+    }
+}
+
+fun foo6(x: B<String>): Additional {
+    return if (x.hashCode() == 0) {
+        <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>aOf<!>()
+    } else {
+        x.convert()
+    }
+}
+
+fun foo7(x: B<String>): Any {
+    return when {
+        x.hashCode() == 0 -> { aOf() }
+        else -> { x.convert() }
+    }
+}
+
+fun foo8(x: B<String>): Additional {
+    return when {
+        x.hashCode() == 0 -> { <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>aOf<!>() }
+        else -> { x.convert() }
+    }
+}

compiler/testData/diagnostics/tests/inference/nothingType/specialCallWithMaterializeAndExpectedType.fir.kt

@@ -0,0 +1,7 @@
+// !DIAGNOSTICS: -UNUSED_VARIABLE
+
+fun foo() {
+    val s: String? = if (true) <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Nothing?")!>materialize()<!> else null
+}
+
+fun <K> materialize(): K = TODO()

compiler/testData/diagnostics/tests/inference/nothingType/specialCallWithMaterializeAndExpectedType.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // !DIAGNOSTICS: -UNUSED_VARIABLE
 
 fun foo() {

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

@@ -14174,6 +14174,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/inference/hasErrorInConstrainingTypes.kt");
             }
 
+            @Test
+            @TestMetadata("ifWithDependentBranches.kt")
+            public void testIfWithDependentBranches() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/ifWithDependentBranches.kt");
+            }
+
             @Test
             @TestMetadata("illegalUnderscoredTypeArgument.kt")
             public void testIllegalUnderscoredTypeArgument() throws Exception {