From 79fa2b6db0659aa7c100a6c9bf4cd7cea9f2afdd Mon Sep 17 00:00:00 2001
From: Alexander Udalov <alexander.udalov@jetbrains.com>
Date: Wed, 28 Jun 2023 15:12:10 +0200
Subject: [PATCH] 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).
---
 .../jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt | 15 ++++++---------
 .../ir/irText/expressions/ifElseIf.fir.ir.txt     |  3 ++-
 .../ir/irText/expressions/ifElseIf.fir.kt.txt     |  3 ++-
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt
index 8e1954f1af5..15ab6a243d8 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt
+++ b/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(
-                    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
-                )
+                generateWhen(startOffset, endOffset, origin, subjectVariable, irBranches, whenExpressionType.toIrType())
             }
         }.also {
             whenExpression.accept(implicitCastInserter, it)
diff --git a/compiler/testData/ir/irText/expressions/ifElseIf.fir.ir.txt b/compiler/testData/ir/irText/expressions/ifElseIf.fir.ir.txt
index 8f6ca19d70e..62e53d91b88 100644
--- a/compiler/testData/ir/irText/expressions/ifElseIf.fir.ir.txt
+++ b/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
diff --git a/compiler/testData/ir/irText/expressions/ifElseIf.fir.kt.txt b/compiler/testData/ir/irText/expressions/ifElseIf.fir.kt.txt
index 4298ff4dd63..3f9258d223a 100644
--- a/compiler/testData/ir/irText/expressions/ifElseIf.fir.kt.txt
+++ b/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 */
 }
+