psi2ir: keep the type of when when possible

Ideally, the type of `IrWhen` should be provided by type inference for
a consistent behavior. `USED_AS_EXPRESSION` from CFG isn't always
consistent with type inference, unfortunately.

The behavior is now aligned with `if`. The type of `when` is kept when
it *can* be an expression, instead of whether it is used or not.

compiler/fir/fir2ir/tests/org/jetbrains/kotlin/fir/Fir2IrTextTestGenerated.java

@@ -1267,6 +1267,11 @@ public class Fir2IrTextTestGenerated extends AbstractFir2IrTextTest {
             runTest("compiler/testData/ir/irText/expressions/whenReturn.kt");
         }
 
+        @TestMetadata("whenUnusedExpression.kt")
+        public void testWhenUnusedExpression() throws Exception {
+            runTest("compiler/testData/ir/irText/expressions/whenUnusedExpression.kt");
+        }
+
         @TestMetadata("whenWithSubjectVariable.kt")
         public void testWhenWithSubjectVariable() throws Exception {
             runTest("compiler/testData/ir/irText/expressions/whenWithSubjectVariable.kt");

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/BranchingExpressionGenerator.kt

@@ -104,11 +104,9 @@ class BranchingExpressionGenerator(statementGenerator: StatementGenerator) : Sta
 
         val inferredType = getInferredTypeWithImplicitCastsOrFail(expression)
 
-        // TODO relies on ControlFlowInformationProvider, get rid of it
-        val isUsedAsExpression = get(BindingContext.USED_AS_EXPRESSION, expression) ?: false
-
         val resultType = when {
-            isUsedAsExpression -> inferredType.toIrType()
+            // Non-exhaustive when can only be used as statement.
+            expression.isExhaustiveWhen() -> inferredType.toIrType()
             KotlinBuiltIns.isNothing(inferredType) -> inferredType.toIrType()
             else -> context.irBuiltIns.unitType
         }

compiler/testData/codegen/box/constants/constantsInWhen.kt

@@ -1,7 +1,3 @@
-// Even before any IR lowerings, the type of `when` is determined to be
-// Unit even though the outer `if` still returns `Int?`. This results
-// in a ClassCastException when that Unit is converted into a Number.
-// IGNORE_BACKEND: JVM_IR
 fun test(
         b: Boolean,
         i: Int

compiler/testData/ir/irText/expressions/whenUnusedExpression.fir.txt

@@ -0,0 +1,22 @@
+FILE fqName:<root> fileName:/whenUnusedExpression.kt
+  FUN name:test visibility:public modality:FINAL <> (b:kotlin.Boolean, i:kotlin.Int) returnType:kotlin.Unit 
+    VALUE_PARAMETER name:b index:0 type:kotlin.Boolean 
+    VALUE_PARAMETER name:i index:1 type:kotlin.Int 
+    BLOCK_BODY
+      WHEN type=kotlin.Unit origin=IF
+        BRANCH
+          if: GET_VAR 'b: kotlin.Boolean declared in <root>.test' type=kotlin.Boolean origin=null
+          then: BLOCK type=kotlin.Unit origin=WHEN
+            VAR IR_TEMPORARY_VARIABLE name:tmp0_subject type:kotlin.Int [val] 
+            WHEN type=kotlin.Unit origin=WHEN
+              BRANCH
+                if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
+                  arg0: GET_VAR 'val tmp0_subject: kotlin.Int [val] declared in <root>.test' type=kotlin.Int origin=null
+                  arg1: CONST Int type=kotlin.Unit value=0
+                then: CONST Int type=kotlin.Unit value=1
+              BRANCH
+                if: CONST Boolean type=kotlin.Boolean value=true
+                then: CONST Null type=kotlin.Nothing? value=null
+        BRANCH
+          if: CONST Boolean type=kotlin.Boolean value=true
+          then: CONST Null type=kotlin.Nothing? value=null

compiler/testData/ir/irText/expressions/whenUnusedExpression.kt

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+
+fun test(b: Boolean, i: Int) {
+    if (b) {
+        when (i) {
+            0 -> 1
+            else -> null
+        }
+    } else null
+}

compiler/testData/ir/irText/expressions/whenUnusedExpression.txt

@@ -0,0 +1,25 @@
+FILE fqName:<root> fileName:/whenUnusedExpression.kt
+  FUN name:test visibility:public modality:FINAL <> (b:kotlin.Boolean, i:kotlin.Int) returnType:kotlin.Unit 
+    VALUE_PARAMETER name:b index:0 type:kotlin.Boolean 
+    VALUE_PARAMETER name:i index:1 type:kotlin.Int 
+    BLOCK_BODY
+      TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+        WHEN type=kotlin.Int? origin=IF
+          BRANCH
+            if: GET_VAR 'b: kotlin.Boolean declared in <root>.test' type=kotlin.Boolean origin=null
+            then: BLOCK type=kotlin.Int? origin=null
+              BLOCK type=kotlin.Int? origin=WHEN
+                VAR IR_TEMPORARY_VARIABLE name:tmp0_subject type:kotlin.Int [val] 
+                  GET_VAR 'i: kotlin.Int declared in <root>.test' type=kotlin.Int origin=null
+                WHEN type=kotlin.Int? origin=WHEN
+                  BRANCH
+                    if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ
+                      arg0: GET_VAR 'val tmp0_subject: kotlin.Int [val] declared in <root>.test' type=kotlin.Int origin=null
+                      arg1: CONST Int type=kotlin.Int value=0
+                    then: CONST Int type=kotlin.Int value=1
+                  BRANCH
+                    if: CONST Boolean type=kotlin.Boolean value=true
+                    then: CONST Null type=kotlin.Nothing? value=null
+          BRANCH
+            if: CONST Boolean type=kotlin.Boolean value=true
+            then: CONST Null type=kotlin.Nothing? value=null

compiler/tests/org/jetbrains/kotlin/ir/IrTextTestCaseGenerated.java

@@ -1267,6 +1267,11 @@ public class IrTextTestCaseGenerated extends AbstractIrTextTestCase {
             runTest("compiler/testData/ir/irText/expressions/whenReturn.kt");
         }
 
+        @TestMetadata("whenUnusedExpression.kt")
+        public void testWhenUnusedExpression() throws Exception {
+            runTest("compiler/testData/ir/irText/expressions/whenUnusedExpression.kt");
+        }
+
         @TestMetadata("whenWithSubjectVariable.kt")
         public void testWhenWithSubjectVariable() throws Exception {
             runTest("compiler/testData/ir/irText/expressions/whenWithSubjectVariable.kt");