From 71e76bc6906205a2e66f7a6b86934b03410cc55e Mon Sep 17 00:00:00 2001
From: Alexey Sedunov <alexey.sedunov@jetbrains.com>
Date: Wed, 15 May 2013 18:54:01 +0400
Subject: [PATCH] Add 'else' check for 'when' expressions

---
 .../branchedTransformations/BranchedUnfoldingUtils.java   | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/BranchedUnfoldingUtils.java b/idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/BranchedUnfoldingUtils.java
index e45924b8b12..755ce457657 100644
--- a/idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/BranchedUnfoldingUtils.java
+++ b/idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/BranchedUnfoldingUtils.java
@@ -42,12 +42,16 @@ public class BranchedUnfoldingUtils {
 
             JetExpression rhs = assignment.getRight();
             if (rhs instanceof JetIfExpression) return UnfoldableKind.ASSIGNMENT_TO_IF;
-            if (rhs instanceof JetWhenExpression) return UnfoldableKind.ASSIGNMENT_TO_WHEN;
+            if (rhs instanceof JetWhenExpression && JetPsiUtil.checkWhenExpressionHasSingleElse((JetWhenExpression) rhs)) {
+                return UnfoldableKind.ASSIGNMENT_TO_WHEN;
+            }
         } else if (root instanceof JetReturnExpression) {
             JetExpression resultExpr = ((JetReturnExpression)root).getReturnedExpression();
 
             if (resultExpr instanceof JetIfExpression) return UnfoldableKind.RETURN_TO_IF;
-            if (resultExpr instanceof JetWhenExpression) return UnfoldableKind.RETURN_TO_WHEN;
+            if (resultExpr instanceof JetWhenExpression && JetPsiUtil.checkWhenExpressionHasSingleElse((JetWhenExpression) resultExpr)) {
+                return UnfoldableKind.RETURN_TO_WHEN;
+            }
         }
 
         return null;