From d89913dc8d754fbb7bbbf558e5ae860b2e1b518c Mon Sep 17 00:00:00 2001
From: Simon Ogorodnik <Simon.Ogorodnik@jetbrains.com>
Date: Thu, 18 Apr 2019 18:15:17 +0300
Subject: [PATCH] FIR: use type of last expression as type of branch, not type
 of block

---
 .../fir/resolve/transformers/FirBodyResolveTransformer.kt  | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirBodyResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirBodyResolveTransformer.kt
index c823e6581f1..1b9debea6b1 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirBodyResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirBodyResolveTransformer.kt
@@ -529,7 +529,12 @@ open class FirBodyResolveTransformer(val session: FirSession, val implicitTypeOn
         whenExpression.transformChildren(this, data)
         if (whenExpression.resultType !is FirResolvedTypeRef) {
             val type = commonSuperType(whenExpression.branches.mapNotNull {
-                it.result.resultType
+                val expression = it.result.statements.lastOrNull() as? FirExpression
+                if (expression != null) {
+                    (expression.resultType as? FirResolvedTypeRef) ?: FirErrorTypeRefImpl(session, null, "No type for when branch result")
+                } else {
+                    FirImplicitUnitTypeRef(session, null)
+                }
             })
             if (type != null) whenExpression.resultType = type
         }