Fir2Ir: generate coercion to Unit for when branches if needed

Before this change, it could happen that `when` of type Unit has a
branch whose type is not Unit. This can lead to problems in IR
lowerings, for example PolymorphicSignatureLowering which is very
reliant on the correct types of expressions and placement of coercions
to Unit (KT-59218).

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -1092,8 +1092,12 @@ class Fir2IrVisitor(
                 if (whenExpression.branches.isEmpty()) {
                     return@convertWithOffsets IrBlockImpl(startOffset, endOffset, irBuiltIns.unitType, origin)
                 }
+                val whenExpressionType =
+                    if (whenExpression.isProperlyExhaustive && whenExpression.branches.none {
+                            it.condition is FirElseIfTrueCondition && it.result.statements.isEmpty()
+                        }) whenExpression.typeRef else session.builtinTypes.unitType
                 val irBranches = whenExpression.branches.mapTo(mutableListOf()) { branch ->
-                    branch.toIrWhenBranch(whenExpression.typeRef)
+                    branch.toIrWhenBranch(whenExpressionType)
                 }
                 if (whenExpression.isProperlyExhaustive && whenExpression.branches.none { it.condition is FirElseIfTrueCondition }) {
                     val irResult = IrCallImpl(
@@ -1106,14 +1110,7 @@ class Fir2IrVisitor(
                         IrConstImpl.boolean(startOffset, endOffset, irBuiltIns.booleanType, true), irResult
                     )
                 }
-                generateWhen(
+                generateWhen(startOffset, endOffset, origin, subjectVariable, irBranches, whenExpressionType.toIrType())
-                    startOffset, endOffset, origin,
-                    subjectVariable, irBranches,
-                    if (whenExpression.isProperlyExhaustive && whenExpression.branches.none {
-                            it.condition is FirElseIfTrueCondition && it.result.statements.isEmpty()
-                        }
-                    ) whenExpression.typeRef.toIrType() else irBuiltIns.unitType
-                )
             }
         }.also {
             whenExpression.accept(implicitCastInserter, it)

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

@@ -47,7 +47,8 @@ FILE fqName:<root> fileName:/ifElseIf.kt
       WHEN type=kotlin.Unit origin=IF
         BRANCH
           if: GET_VAR 'flag: kotlin.Boolean declared in <root>.testEmptyBranches2' type=kotlin.Boolean origin=null
-          then: CONST Boolean type=kotlin.Boolean value=true
+          then: TYPE_OP type=kotlin.Unit origin=IMPLICIT_COERCION_TO_UNIT typeOperand=kotlin.Unit
+            CONST Boolean type=kotlin.Boolean value=true
         BRANCH
           if: CONST Boolean type=kotlin.Boolean value=true
           then: BLOCK type=kotlin.Unit origin=null

compiler/testData/ir/irText/expressions/ifElseIf.fir.kt.txt

@@ -24,7 +24,7 @@ fun testEmptyBranches2(flag: Boolean) {
     else -> true
   } /*~> Unit */
   when {
-    flag -> true
+    flag -> true /*~> Unit */
     else -> { // BLOCK
     }
   }
@@ -37,3 +37,4 @@ fun testEmptyBranches3(flag: Boolean) {
     else -> true
   } /*~> Unit */
 }
+