K2: Emulate K1 behavior in case of when-elvis combination

^KT-55692 Related

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

@@ -14066,6 +14066,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/inference/dontThrowEmptyIntersectionException.kt");
             }
 
+            @Test
+            @TestMetadata("elvisInsideWhen.kt")
+            public void testElvisInsideWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/elvisInsideWhen.kt");
+            }
+
             @Test
             @TestMetadata("equalitySubstitutionInsideNonInvariantType.kt")
             public void testEqualitySubstitutionInsideNonInvariantType() throws Exception {

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

@@ -14072,6 +14072,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/inference/dontThrowEmptyIntersectionException.kt");
             }
 
+            @Test
+            @TestMetadata("elvisInsideWhen.kt")
+            public void testElvisInsideWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/elvisInsideWhen.kt");
+            }
+
             @Test
             @TestMetadata("equalitySubstitutionInsideNonInvariantType.kt")
             public void testEqualitySubstitutionInsideNonInvariantType() throws Exception {

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

@@ -14066,6 +14066,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/inference/dontThrowEmptyIntersectionException.kt");
             }
 
+            @Test
+            @TestMetadata("elvisInsideWhen.kt")
+            public void testElvisInsideWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/elvisInsideWhen.kt");
+            }
+
             @Test
             @TestMetadata("equalitySubstitutionInsideNonInvariantType.kt")
             public void testEqualitySubstitutionInsideNonInvariantType() throws Exception {

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

@@ -243,6 +243,13 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirAbstractBodyRes
         if (elvisExpression.calleeReference is FirResolvedNamedReference) return elvisExpression
         elvisExpression.transformAnnotations(transformer, data)
 
+        // Do not use expect type when it's came from if/when (when it doesn't require completion)
+        // It returns us to K1 behavior in the case of when-elvis combination (see testData/diagnostics/tests/inference/elvisInsideWhen.kt)
+        // But this is mostly a hack that I hope might be lifted once KT-55692 is considered
+        // NB: Currently situation `it is ResolutionMode.WithExpectedType && !it.forceFullCompletion` might only happen in case of when branches
+        @Suppress("NAME_SHADOWING")
+        val data = data.takeUnless { it is ResolutionMode.WithExpectedType && !it.forceFullCompletion } ?: ResolutionMode.ContextDependent
+
         val expectedType = data.expectedType?.coneTypeSafe<ConeKotlinType>()
         val mayBeCoercionToUnitApplied = (data as? ResolutionMode.WithExpectedType)?.mayBeCoercionToUnitApplied == true
 

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

@@ -0,0 +1,19 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+
+fun <X> select(vararg x: X): X = x[0]
+fun <E> myOut(): Out<E> = TODO()
+
+interface Out<out T>
+
+fun foo(x: Int, o: Out<String>, oNullable: Out<String>?) {
+    val y: Out<*> = when {
+        x > 0 -> when {
+            x == 1 -> o
+            // Once elvis is being analyzed with expected type Out<*>, because of the @Exact annotation it sticks to that type
+            // While here, it's better to use the Out<String> type from the main branch
+            else -> oNullable ?: myOut()
+        }
+        else -> myOut<Any?>()
+    }
+}

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

@@ -14072,6 +14072,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/inference/dontThrowEmptyIntersectionException.kt");
             }
 
+            @Test
+            @TestMetadata("elvisInsideWhen.kt")
+            public void testElvisInsideWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/inference/elvisInsideWhen.kt");
+            }
+
             @Test
             @TestMetadata("equalitySubstitutionInsideNonInvariantType.kt")
             public void testEqualitySubstitutionInsideNonInvariantType() throws Exception {