Get rid of excessive replaces

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiFactory.java

@@ -254,17 +254,9 @@ public class JetPsiFactory {
         return (JetBinaryExpression) createExpression(project, lhs + " " + op + " " + rhs);
     }
 
-    @SuppressWarnings("ConstantConditions")
     @NotNull
     public static JetBinaryExpression createBinaryExpression(Project project, @NotNull JetExpression lhs, @NotNull String op, @NotNull JetExpression rhs) {
-        JetBinaryExpression assignment = createBinaryExpression(project, "_", op, "_");
+        return createBinaryExpression(project, lhs.getText(), op, rhs.getText());
-
-        assert assignment.getRight() != null;
-
-        assignment = (JetBinaryExpression)assignment.getLeft().replace(lhs).getParent();
-        assignment = (JetBinaryExpression)assignment.getRight().replace(rhs).getParent();
-
-        return assignment;
     }
 
     public static JetTypeCodeFragment createTypeCodeFragment(Project project, String text, PsiElement context) {
@@ -280,15 +272,9 @@ public class JetPsiFactory {
         return (JetIsExpression) createExpression(project, lhs + " " + (negated ? "!is" : "is") + " " + rhs);
     }
 
-    @SuppressWarnings("ConstantConditions")
     @NotNull
     public static JetIsExpression createIsExpression(Project project, @NotNull JetExpression lhs, @NotNull JetTypeReference rhs, boolean negated) {
-        JetIsExpression isExpression = createIsExpression(project, "_", "_", negated);
+        return createIsExpression(project, lhs.getText(), rhs.getText(), negated);
-
-        isExpression = (JetIsExpression)isExpression.getLeftHandSide().replace(lhs).getParent();
-        isExpression = (JetIsExpression)isExpression.getTypeRef().replace(rhs).getParent();
-
-        return isExpression;
     }
 
     @NotNull
@@ -296,65 +282,84 @@ public class JetPsiFactory {
         return (JetReturnExpression) createExpression(project, "return " + text);
     }
 
-    @SuppressWarnings("ConstantConditions")
     @NotNull
     public static JetReturnExpression createReturn(Project project, @NotNull JetExpression expression) {
-        JetReturnExpression returnExpr = createReturn(project, "_");
+        return createReturn(project, expression.getText());
-
-        assert returnExpr.getReturnedExpression() != null;
-
-        return (JetReturnExpression)returnExpr.getReturnedExpression().replace(expression).getParent();
     }
 
     @NotNull
-    public static JetIfExpression createIf(
+    public static JetIfExpression createIf(Project project,
-            Project project,
+            @NotNull JetExpression condition, @NotNull JetExpression thenExpr, @Nullable JetExpression elseExpr) {
-            @NotNull String condText, @NotNull String thenText, @Nullable String elseText,
+        return (JetIfExpression) createExpression(project, JetPsiUnparsingUtils.toIf(condition, thenExpr, elseExpr));
-            boolean thenNewLine, boolean elseNewLine) {
-        String thenSpace = thenNewLine ? "\n" : " ";
-        String elseSpace = elseNewLine ? "\n" : " ";
-        return (JetIfExpression) createExpression(project, "if (" + condText + ")" + thenSpace + thenText + (elseText != null ? elseSpace + "else " + elseText : ""));
     }
 
-    @SuppressWarnings("ConstantConditions")
+    public static class IfChainBuilder {
-    @NotNull
+        private final StringBuilder sb = new StringBuilder();
-    public static JetIfExpression createIf(
+        private boolean first = true;
-            Project project,
+        private boolean frozen = false;
-            @NotNull JetExpression condition, @NotNull JetExpression thenExpr, @Nullable JetExpression elseExpr,
-            boolean thenNewLine, boolean elseNewLine) {
-        JetIfExpression ifExpr = createIf(project, "_", "_", elseExpr != null ? "_" : null, thenNewLine, elseNewLine);
 
-        assert ifExpr.getCondition() != null;
+        public IfChainBuilder() {
-        assert ifExpr.getThen() != null;
-        assert elseExpr == null || ifExpr.getElse() != null;
-
-        ifExpr = (JetIfExpression)ifExpr.getCondition().replace(condition).getParent().getParent();
-        ifExpr = (JetIfExpression)ifExpr.getThen().replace(thenExpr).getParent().getParent();
-        if (elseExpr != null) {
-            ifExpr = (JetIfExpression)ifExpr.getElse().replace(elseExpr).getParent().getParent();
         }
 
-        return ifExpr;
+        @NotNull
+        public IfChainBuilder ifBranch(@NotNull String conditionText, @NotNull String expressionText) {
+            if (first) {
+                first = false;
+            } else {
+                sb.append("else ");
+            }
+
+            sb.append("if (").append(conditionText).append(") ").append(expressionText).append("\n");
+            return this;
+        }
+
+        @NotNull
+        public IfChainBuilder ifBranch(@NotNull JetExpression condition, @NotNull JetExpression expression) {
+            return ifBranch(condition.getText(), expression.getText());
+        }
+
+        @NotNull
+        public IfChainBuilder elseBranch(@NotNull String expressionText) {
+            sb.append("else ").append(expressionText);
+            return this;
+        }
+
+        @NotNull
+        public IfChainBuilder elseBranch(@NotNull JetExpression expression) {
+            return elseBranch(expression.getText());
+        }
+
+        @NotNull
+        public JetIfExpression toExpression(Project project) {
+            if (!frozen) {
+                frozen = true;
+            }
+            return (JetIfExpression) createExpression(project, sb.toString());
+        }
     }
 
-    public static class WhenTemplateBuilder {
+    public static class WhenBuilder {
         private final StringBuilder sb = new StringBuilder("when ");
         private boolean frozen = false;
         private boolean inCondition = false;
 
-        public WhenTemplateBuilder(boolean addSubject) {
+        public WhenBuilder() {
-            if (addSubject) {
+            this((String)null);
-                sb.append("(_) ");
+        }
+
+        public WhenBuilder(@Nullable String subjectText) {
+            if (subjectText != null) {
+                sb.append("(").append(subjectText).append(") ");
             }
             sb.append("{\n");
         }
 
-        private void finishAndFreeze() {
+        public WhenBuilder(@Nullable JetExpression subject) {
-            sb.append("else -> _\n}");
+            this(subject != null ? subject.getText() : null);
-            frozen = true;
         }
 
-        private WhenTemplateBuilder addCondition(String text) {
+        @NotNull
+        public WhenBuilder condition(@NotNull String text) {
             assert !frozen;
 
             if (!inCondition) {
@@ -367,84 +372,95 @@ public class JetPsiFactory {
             return this;
         }
 
-        public WhenTemplateBuilder addBranchWithSingleCondition() {
+        @NotNull
-            assert !frozen;
+        public WhenBuilder condition(@NotNull JetExpression expression) {
-
+            return condition(expression.getText());
-            sb.append("_ -> _\n");
-
-            return this;
         }
 
-        public WhenTemplateBuilder addBranchesWithSingleCondition(int count) {
+        @NotNull
-            for (int i = 0; i < count; i++) {
+        public WhenBuilder pattern(@NotNull String typeReferenceText, boolean negated) {
-                addBranchWithSingleCondition();
+            return condition((negated ? "!is" : "is") + " " + typeReferenceText);
-            }
-
-            return this;
         }
 
-        public WhenTemplateBuilder addBranchWithMultiCondition(int conditionCount) {
+        @NotNull
-            assert !frozen;
+        public WhenBuilder pattern(@NotNull JetTypeReference typeReference, boolean negated) {
-            assert conditionCount > 0;
+            return pattern(typeReference.getText(), negated);
-
-            sb.append("_");
-            for (int i = 0; i < conditionCount - 1; i++) {
-                sb.append(", _");
-            }
-            sb.append(" -> _\n");
-
-            return this;
         }
 
-        public WhenTemplateBuilder addExpressionCondition() {
+        @NotNull
-            return addCondition("_");
+        public WhenBuilder range(@NotNull String argumentText, boolean negated) {
+            return condition((negated ? "!in" : "in") + " " + argumentText);
         }
 
-        public WhenTemplateBuilder addIsCondition(boolean negated) {
+        @NotNull
-            return addCondition((negated ? "!is" : "is") + " _");
+        public WhenBuilder range(@NotNull JetExpression argument, boolean negated) {
+            return range(argument.getText(), negated);
         }
 
-        public WhenTemplateBuilder addInCondition(boolean negated) {
+        @NotNull
-            return addCondition((negated ? "!in" : "in") + " _");
+        public WhenBuilder branchExpression(@NotNull String expressionText) {
-        }
-
-        public WhenTemplateBuilder finishBranch() {
             assert !frozen;
             assert inCondition;
 
             inCondition = false;
-            sb.append(" -> _\n");
+            sb.append(" -> ").append(expressionText).append("\n");
 
             return this;
         }
 
+        @NotNull
+        public WhenBuilder branchExpression(@NotNull JetExpression expression) {
+            return branchExpression(expression.getText());
+        }
+
+        @NotNull
+        public WhenBuilder entry(@NotNull String entryText) {
+            assert !frozen;
+            assert !inCondition;
+
+            sb.append(entryText).append("\n");
+
+            return this;
+        }
+
+        @NotNull
+        public WhenBuilder entry(@NotNull JetWhenEntry whenEntry) {
+            return entry(whenEntry.getText());
+        }
+
+        @NotNull
+        public WhenBuilder elseEntry(@NotNull String text) {
+            return entry("else -> " + text);
+        }
+
+        @NotNull
+        public WhenBuilder elseEntry(@NotNull JetExpression expression) {
+            return elseEntry(expression.getText());
+        }
+
+        @NotNull
         public JetWhenExpression toExpression(Project project) {
             if (!frozen) {
-                finishAndFreeze();
+                sb.append("}");
+                frozen = true;
             }
             return (JetWhenExpression) createExpression(project, sb.toString());
         }
     }
 
-    public static JetWhenExpression createWhenWithoutSubject(Project project, int positiveBranchCount) {
+    @NotNull
-        return new WhenTemplateBuilder(false).addBranchesWithSingleCondition(positiveBranchCount).toExpression(project);
-    }
-
     public static JetParenthesizedExpression createParenthesizedExpression(Project project, @NotNull String text) {
-        return (JetParenthesizedExpression)createExpression(project, "(" + text + ")");
+        return (JetParenthesizedExpression) createExpression(project, "(" + text + ")");
     }
 
-    @SuppressWarnings("ConstantConditions")
+    @NotNull
     public static JetParenthesizedExpression createParenthesizedExpression(Project project, @NotNull JetExpression expression) {
-        JetParenthesizedExpression parenthesizedExpression = createParenthesizedExpression(project, "_");
+        return createParenthesizedExpression(project, expression.getText());
-        parenthesizedExpression.getExpression().replace(expression);
-
-        return parenthesizedExpression;
     }
 
+    @NotNull
     public static JetParenthesizedExpression createParenthesizedExpressionIfNeeded(Project project, @NotNull JetExpression expression) {
         return (expression instanceof JetParenthesizedExpression) ?
-               (JetParenthesizedExpression) expression :
+               (JetParenthesizedExpression) expression : createParenthesizedExpression(project, expression);
-               createParenthesizedExpression(project, expression);
     }
 }

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiUnparsingUtils.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.lang.psi;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public class JetPsiUnparsingUtils {
+    private JetPsiUnparsingUtils() {
+    }
+
+    @NotNull
+    public static String toIf(@NotNull JetExpression condition, @NotNull JetExpression thenExpression, @Nullable JetExpression elseExpression) {
+        return toIf(condition.getText(), thenExpression.getText(), elseExpression != null ? elseExpression.getText() : null);
+    }
+
+    @NotNull
+    public static String toIf(@NotNull String condition, @NotNull String thenExpression, @Nullable String elseExpression) {
+        return "if " + parenthesizeTextIfNeeded(condition) + " " + thenExpression + (elseExpression != null ? " else " + elseExpression : "");
+    }
+
+    @NotNull
+    public static String toBinaryExpression(@NotNull JetExpression left, @NotNull String op, @NotNull JetElement right) {
+        return toBinaryExpression(left.getText(), op, right.getText());
+    }
+
+    @NotNull
+    public static String toBinaryExpression(@NotNull String left, @NotNull String op, @NotNull String right) {
+        return left + " " + op + " " + right;
+    }
+
+    @NotNull
+    public static String parenthesizeIfNeeded(@NotNull JetExpression expression) {
+        String text = expression.getText();
+        return (expression instanceof JetParenthesizedExpression ||
+                expression instanceof JetConstantExpression ||
+                expression instanceof JetSimpleNameExpression)
+               ? text : "(" + text + ")";
+    }
+
+    @NotNull
+    public static String parenthesizeTextIfNeeded(@NotNull String expressionText) {
+        return (expressionText.startsWith("(") && expressionText.endsWith(")")) ? expressionText : "(" + expressionText + ")";
+    }
+}

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetPsiUtil.java

@@ -683,6 +683,16 @@ public class JetPsiUtil {
         return false;
     }
 
+    public static <C extends Collection<JetVariableDeclaration>> C getBlockVariableDeclarations(@NotNull JetBlockExpression block, @NotNull C collection) {
+        for (JetElement element : block.getStatements()) {
+            if (element instanceof JetVariableDeclaration) {
+                collection.add((JetVariableDeclaration) element);
+            }
+        }
+
+        return collection;
+    }
+
     public static boolean checkWhenExpressionHasSingleElse(JetWhenExpression whenExpression) {
         int elseCount = 0;
         for (JetWhenEntry entry : whenExpression.getEntries()) {

idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/BranchedFoldingUtils.java

@@ -47,6 +47,7 @@ public class BranchedFoldingUtils {
             }
 
             if (assignment.getParent() instanceof JetBlockExpression) {
+                //noinspection ConstantConditions
                 return !JetPsiUtil.checkVariableDeclarationInBlock((JetBlockExpression) assignment.getParent(), assignment.getLeft().getText());
             }
 
@@ -239,7 +240,7 @@ public class BranchedFoldingUtils {
         assertNotNull(elseRoot);
 
         //noinspection ConstantConditions
-        JetIfExpression newIfExpression = JetPsiFactory.createIf(project, condition, thenRoot, elseRoot, false, false);
+        JetIfExpression newIfExpression = JetPsiFactory.createIf(project, condition, thenRoot, elseRoot);
         JetReturnExpression newReturnExpression = JetPsiFactory.createReturn(project, newIfExpression);
 
         newIfExpression = (JetIfExpression)newReturnExpression.getReturnedExpression();

idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/BranchedUnfoldingUtils.java

@@ -35,11 +35,10 @@ public class BranchedUnfoldingUtils {
 
         if (JetPsiUtil.isAssignment(root)) {
             JetBinaryExpression assignment = (JetBinaryExpression)root;
-            JetExpression lhs = assignment.getLeft();
+
+            assertNotNull(assignment.getLeft());
+
             JetExpression rhs = assignment.getRight();
-
-            if (!(lhs instanceof JetSimpleNameExpression)) return null;
-
             if (rhs instanceof JetIfExpression) return UnfoldableKind.ASSIGNMENT_TO_IF;
             if (rhs instanceof JetWhenExpression) return UnfoldableKind.ASSIGNMENT_TO_WHEN;
         } else if (root instanceof JetReturnExpression) {
@@ -75,6 +74,7 @@ public class BranchedUnfoldingUtils {
         assertNotNull(thenExpr);
         assertNotNull(elseExpr);
 
+        //noinspection ConstantConditions
         thenExpr.replace(JetPsiFactory.createBinaryExpression(project, lhs, op, thenExpr));
         elseExpr.replace(JetPsiFactory.createBinaryExpression(project, lhs, op, elseExpr));
 
@@ -97,6 +97,7 @@ public class BranchedUnfoldingUtils {
 
             assertNotNull(currExpr);
 
+            //noinspection ConstantConditions
             currExpr.replace(JetPsiFactory.createBinaryExpression(project, lhs, op, currExpr));
         }
 

idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/GuardedExpression.java

@@ -1,27 +0,0 @@
-package org.jetbrains.jet.plugin.codeInsight.codeTransformations.branchedTransformations;
-
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.jet.lang.psi.JetExpression;
-
-public class GuardedExpression {
-    @NotNull
-    private final JetExpression condition;
-
-    @NotNull
-    private final JetExpression baseExpression;
-
-    public GuardedExpression(@NotNull JetExpression condition, @NotNull JetExpression baseExpression) {
-        this.condition = condition;
-        this.baseExpression = baseExpression;
-    }
-
-    @NotNull
-    public JetExpression getCondition() {
-        return condition;
-    }
-
-    @NotNull
-    public JetExpression getBaseExpression() {
-        return baseExpression;
-    }
-}

idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/IfWhenUtils.java

@@ -1,6 +1,5 @@
 package org.jetbrains.jet.plugin.codeInsight.codeTransformations.branchedTransformations;
 
-import com.intellij.openapi.project.Project;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lexer.JetTokens;
@@ -8,6 +7,8 @@ import org.jetbrains.jet.lexer.JetTokens;
 import java.util.ArrayList;
 import java.util.List;
 
+import static org.jetbrains.jet.lang.psi.JetPsiUnparsingUtils.*;
+
 public class IfWhenUtils {
 
     public static final String TRANSFORM_WITHOUT_CHECK =
@@ -71,8 +72,7 @@ public class IfWhenUtils {
     }
 
     public static void transformIfToWhen(@NotNull JetIfExpression ifExpression) {
-        List<MultiGuardedExpression> positiveBranches = new ArrayList<MultiGuardedExpression>();
+        JetPsiFactory.WhenBuilder builder = new JetPsiFactory.WhenBuilder();
-        JetExpression elseExpression = null;
 
         JetIfExpression currIfExpression = ifExpression;
         do {
@@ -84,123 +84,73 @@ public class IfWhenUtils {
             assertNotNull(thenBranch);
             assertNotNull(elseBranch);
 
+            List<JetExpression> orBranches = splitExpressionToOrBranches(condition);
+            for (JetExpression orBranch : orBranches) {
+                builder.condition(orBranch);
+            }
             //noinspection ConstantConditions
-            positiveBranches.add(new MultiGuardedExpression(splitExpressionToOrBranches(condition), thenBranch));
+            builder.branchExpression(thenBranch);
 
             if (elseBranch instanceof JetIfExpression) {
                 currIfExpression = (JetIfExpression) elseBranch;
             }
             else {
                 currIfExpression = null;
-                elseExpression = elseBranch;
+                //noinspection ConstantConditions
+                builder.elseEntry(elseBranch);
             }
         } while (currIfExpression != null);
 
-        JetPsiFactory.WhenTemplateBuilder builder = new JetPsiFactory.WhenTemplateBuilder(false);
+        ifExpression.replace(builder.toExpression(ifExpression.getProject()));
-        for (MultiGuardedExpression positiveBranch : positiveBranches) {
-            builder.addBranchWithMultiCondition(positiveBranch.getConditions().size());
-        }
-
-        JetWhenExpression whenExpression = builder.toExpression(ifExpression.getProject());
-
-        int i = 0;
-        List<JetWhenEntry> entries = whenExpression.getEntries();
-        for (JetWhenEntry entry : entries) {
-            if (entry.isElse()) {
-                //noinspection ConstantConditions
-                entry.getExpression().replace(elseExpression);
-                break;
-            }
-
-            MultiGuardedExpression branch = positiveBranches.get(i++);
-
-            //noinspection ConstantConditions
-            entry.getExpression().replace(branch.getBaseExpression());
-
-            int j = 0;
-            JetWhenCondition[] conditions = entry.getConditions();
-            for (JetWhenCondition condition : conditions) {
-                assert condition instanceof JetWhenConditionWithExpression : TRANSFORM_WITHOUT_CHECK;
-
-                JetExpression conditionExpression = ((JetWhenConditionWithExpression) condition).getExpression();
-                assertNotNull(conditionExpression);
-
-                //noinspection ConstantConditions
-                conditionExpression.replace(branch.getConditions().get(j++));
-            }
-        }
-
-        ifExpression.replace(whenExpression);
     }
 
-    @SuppressWarnings("ConstantConditions")
+    private static String combineWhenConditions(JetWhenCondition[] conditions, JetExpression subject) {
-    private static JetExpression combineWhenConditions(Project project, JetWhenCondition[] conditions, JetExpression subject) {
         int n = conditions.length;
-        assert n > 0 : TRANSFORM_WITHOUT_CHECK;
+        if (n == 0) return "";
 
-        JetWhenCondition condition = conditions[n - 1];
+        JetWhenCondition condition = conditions[0];
         assert condition != null : TRANSFORM_WITHOUT_CHECK;
 
-        JetExpression resultExpr = WhenUtils.whenConditionToExpression(condition, subject);
+        StringBuilder sb = new StringBuilder();
+
+        String text = WhenUtils.whenConditionToExpressionText(condition, subject);
         if (n > 1) {
-            resultExpr = JetPsiFactory.createParenthesizedExpressionIfNeeded(project, resultExpr);
+            text = parenthesizeTextIfNeeded(text);
         }
+        sb.append(text);
 
-        for (int i = n - 2; i >= 0; i--) {
+        for (int i = 1; i < n; i++) {
             JetWhenCondition currCondition = conditions[i];
-
             assert currCondition != null : TRANSFORM_WITHOUT_CHECK;
 
-            resultExpr = JetPsiFactory.createBinaryExpression(
+            sb.append(" || ").append(parenthesizeTextIfNeeded(WhenUtils.whenConditionToExpressionText(currCondition, subject)));
-                    project,
-                    JetPsiFactory.createParenthesizedExpressionIfNeeded(project, WhenUtils.whenConditionToExpression(currCondition, subject)),
-                    "||",
-                    resultExpr);
         }
 
-        return resultExpr;
+        return sb.toString();
     }
 
     public static void transformWhenToIf(@NotNull JetWhenExpression whenExpression) {
-        Project project = whenExpression.getProject();
+        JetPsiFactory.IfChainBuilder builder = new JetPsiFactory.IfChainBuilder();
-
-        JetExpression elseExpression = null;
-        List<GuardedExpression> positiveBranches = new ArrayList<GuardedExpression>();
 
         List<JetWhenEntry> entries = whenExpression.getEntries();
         for (JetWhenEntry entry : entries) {
             JetExpression branch = entry.getExpression();
-
             assertNotNull(branch);
 
             if (entry.isElse()) {
-                elseExpression = branch;
+                //noinspection ConstantConditions
+                builder.elseBranch(branch);
             } else {
-                JetExpression branchCondition = combineWhenConditions(project, entry.getConditions(), whenExpression.getSubjectExpression());
+                String branchConditionText = combineWhenConditions(entry.getConditions(), whenExpression.getSubjectExpression());
-                JetExpression branchExpression = entry.getExpression();
 
+                JetExpression branchExpression = entry.getExpression();
                 assertNotNull(branchExpression);
 
                 //noinspection ConstantConditions
-                positiveBranches.add(new GuardedExpression(branchCondition, branchExpression));
+                builder.ifBranch(branchConditionText, branchExpression.getText());
             }
         }
 
-        assertNotNull(elseExpression);
+        whenExpression.replace(builder.toExpression(whenExpression.getProject()));
-        assert !positiveBranches.isEmpty() : TRANSFORM_WITHOUT_CHECK;
-
-        JetExpression outerExpression = elseExpression;
-
-        for (int i = positiveBranches.size() - 1; i >= 0; i--) {
-            GuardedExpression branch = positiveBranches.get(i);
-
-            outerExpression = JetPsiFactory.createIf(
-                    project,
-                    branch.getCondition(), branch.getBaseExpression(), outerExpression,
-                    !(branch.getBaseExpression() instanceof JetBlockExpression), !(outerExpression instanceof JetBlockExpression));
-        }
-
-        //noinspection ConstantConditions
-        whenExpression.replace(outerExpression);
     }
 }

idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/MultiGuardedExpression.java

@@ -1,29 +0,0 @@
-package org.jetbrains.jet.plugin.codeInsight.codeTransformations.branchedTransformations;
-
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.jet.lang.psi.JetExpression;
-
-import java.util.List;
-
-public class MultiGuardedExpression {
-    @NotNull
-    private final List<JetExpression> conditions;
-
-    @NotNull
-    private final JetExpression baseExpression;
-
-    public MultiGuardedExpression(@NotNull List<JetExpression> conditions, @NotNull JetExpression baseExpression) {
-        this.conditions = conditions;
-        this.baseExpression = baseExpression;
-    }
-
-    @NotNull
-    public List<JetExpression> getConditions() {
-        return conditions;
-    }
-
-    @NotNull
-    public JetExpression getBaseExpression() {
-        return baseExpression;
-    }
-}

idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/WhenUtils.java

@@ -1,11 +1,12 @@
 package org.jetbrains.jet.plugin.codeInsight.codeTransformations.branchedTransformations;
 
-import com.intellij.openapi.project.Project;
 import com.intellij.psi.tree.IElementType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lexer.JetTokens;
 
+import static org.jetbrains.jet.lang.psi.JetPsiUnparsingUtils.*;
+
 import java.util.List;
 
 public class WhenUtils {
@@ -87,7 +88,6 @@ public class WhenUtils {
 
     public static void flattenWhen(@NotNull JetWhenExpression whenExpression) {
         JetExpression subjectExpression = whenExpression.getSubjectExpression();
-        boolean hasSubject = subjectExpression != null;
 
         JetExpression elseBranch = JetPsiUtil.getWhenElseBranch(whenExpression);
         assert elseBranch instanceof JetWhenExpression : TRANSFORM_WITHOUT_CHECK;
@@ -97,38 +97,36 @@ public class WhenUtils {
         List<JetWhenEntry> outerEntries = whenExpression.getEntries();
         List<JetWhenEntry> innerEntries = nestedWhenExpression.getEntries();
 
-        JetWhenExpression newWhenExpression = new JetPsiFactory.WhenTemplateBuilder(hasSubject)
+        JetPsiFactory.WhenBuilder builder = new JetPsiFactory.WhenBuilder(subjectExpression);
-                .addBranchesWithSingleCondition(outerEntries.size() + innerEntries.size() - 2)
-                .toExpression(whenExpression.getProject());
 
-        if (hasSubject) {
-            JetExpression dummySubjectExpression = newWhenExpression.getSubjectExpression();
-            assertNotNull(dummySubjectExpression);
-
-            //noinspection ConstantConditions
-            dummySubjectExpression.replace(subjectExpression);
-        }
-
-        List<JetWhenEntry> newEntries = newWhenExpression.getEntries();
-
-        int i = 0;
         for (JetWhenEntry entry : outerEntries) {
-            if (!entry.isElse()) {
+            if (entry.isElse()) continue;
-                newEntries.get(i++).replace(entry);
+
-            }
+            builder.entry(entry);
-        }
-        for (JetWhenEntry entry : innerEntries) {
-            newEntries.get(i++).replace(entry);
         }
 
-        whenExpression.replace(newWhenExpression);
+        for (JetWhenEntry entry : innerEntries) {
+            builder.entry(entry);
+        }
+
+        whenExpression.replace(builder.toExpression(whenExpression.getProject()));
     }
 
-    private static JetWhenExpression createWhenTemplateWithSubject(@NotNull JetWhenExpression whenExpression) {
+    public static void introduceWhenSubject(@NotNull JetWhenExpression whenExpression) {
-        JetPsiFactory.WhenTemplateBuilder builder = new JetPsiFactory.WhenTemplateBuilder(true);
+        JetExpression subject = getWhenSubjectCandidate(whenExpression);
+        assertNotNull(subject);
+
+        JetPsiFactory.WhenBuilder builder = new JetPsiFactory.WhenBuilder(subject);
 
         for (JetWhenEntry entry : whenExpression.getEntries()) {
-            if (entry.isElse()) continue;
+            JetExpression branchExpression = entry.getExpression();
+            assertNotNull(branchExpression);
+
+            if (entry.isElse()) {
+                //noinspection ConstantConditions
+                builder.elseEntry(branchExpression);
+                continue;
+            }
 
             for (JetWhenCondition condition : entry.getConditions()) {
                 assert condition instanceof JetWhenConditionWithExpression : TRANSFORM_WITHOUT_CHECK;
@@ -136,81 +134,13 @@ public class WhenUtils {
                 JetExpression conditionExpression = ((JetWhenConditionWithExpression) condition).getExpression();
 
                 if (conditionExpression instanceof JetIsExpression) {
-                    builder.addIsCondition(((JetIsExpression) conditionExpression).isNegated());
+                    JetIsExpression isExpression = (JetIsExpression) conditionExpression;
-                }
-                else if (conditionExpression instanceof JetBinaryExpression) {
-                    JetBinaryExpression binaryExpression = (JetBinaryExpression) conditionExpression;
 
-                    IElementType op = binaryExpression.getOperationToken();
+                    JetTypeReference typeReference = isExpression.getTypeRef();
-                    if (op == JetTokens.IN_KEYWORD) {
-                        builder.addInCondition(false);
-                    }
-                    else if (op == JetTokens.NOT_IN) {
-                        builder.addInCondition(true);
-                    }
-                    else if (op == JetTokens.EQEQ) {
-                        builder.addExpressionCondition();
-                    }
-                    else assert false : TRANSFORM_WITHOUT_CHECK;
-                }
-                else assert false : TRANSFORM_WITHOUT_CHECK;
-            }
-
-            builder.finishBranch();
-        }
-
-        return builder.toExpression(whenExpression.getProject());
-    }
-
-    public static void introduceWhenSubject(@NotNull JetWhenExpression whenExpression) {
-        JetExpression subject = getWhenSubjectCandidate(whenExpression);
-        assertNotNull(subject);
-
-        JetWhenExpression newWhenExpression = createWhenTemplateWithSubject(whenExpression);
-
-        JetExpression newSubject = newWhenExpression.getSubjectExpression();
-        assertNotNull(newSubject);
-
-        //noinspection ConstantConditions
-        newSubject.replace(subject);
-
-        int i = 0;
-        List<JetWhenEntry> entries = whenExpression.getEntries();
-        List<JetWhenEntry> newEntries = newWhenExpression.getEntries();
-        for (JetWhenEntry newEntry : newEntries) {
-            JetWhenEntry entry = entries.get(i++);
-
-            JetExpression branchExpression = entry.getExpression();
-            assertNotNull(branchExpression);
-
-            JetExpression newBranchExpression = newEntry.getExpression();
-            assertNotNull(newBranchExpression);
-
-            //noinspection ConstantConditions
-            newBranchExpression.replace(branchExpression);
-
-            int j = 0;
-            JetWhenCondition[] conditions = entry.getConditions();
-            JetWhenCondition[] newConditions = newEntry.getConditions();
-
-            for (JetWhenCondition newCondition : newConditions) {
-                JetWhenCondition condition = conditions[j++];
-
-                assert condition instanceof JetWhenConditionWithExpression : TRANSFORM_WITHOUT_CHECK;
-
-                JetExpression conditionExpression = ((JetWhenConditionWithExpression) condition).getExpression();
-
-                if (conditionExpression instanceof JetIsExpression) {
-                    assert newCondition instanceof JetWhenConditionIsPattern : TRANSFORM_WITHOUT_CHECK;
-
-                    JetTypeReference typeReference = ((JetIsExpression) conditionExpression).getTypeRef();
                     assertNotNull(typeReference);
 
-                    JetTypeReference newTypeReference = ((JetWhenConditionIsPattern) newCondition).getTypeRef();
-                    assertNotNull(newTypeReference);
-
                     //noinspection ConstantConditions
-                    newTypeReference.replace(typeReference);
+                    builder.pattern(typeReference, isExpression.isNegated());
                 }
                 else if (conditionExpression instanceof JetBinaryExpression) {
                     JetBinaryExpression binaryExpression = (JetBinaryExpression) conditionExpression;
@@ -219,48 +149,31 @@ public class WhenUtils {
                     assertNotNull(rhs);
 
                     IElementType op = binaryExpression.getOperationToken();
-                    if (op == JetTokens.IN_KEYWORD || op == JetTokens.NOT_IN) {
+                    if (op == JetTokens.IN_KEYWORD) {
-                        assert newCondition instanceof JetWhenConditionInRange : TRANSFORM_WITHOUT_CHECK;
-
-                        JetExpression newRangeExpression = ((JetWhenConditionInRange) newCondition).getRangeExpression();
-                        assertNotNull(newRangeExpression);
-
                         //noinspection ConstantConditions
-                        newRangeExpression.replace(rhs);
+                        builder.range(rhs, false);
+                    }
+                    else if (op == JetTokens.NOT_IN) {
+                        //noinspection ConstantConditions
+                        builder.range(rhs, true);
                     }
                     else if (op == JetTokens.EQEQ) {
-                        assert newCondition instanceof JetWhenConditionWithExpression : TRANSFORM_WITHOUT_CHECK;
-
-                        JetExpression newConditionExpression = ((JetWhenConditionWithExpression) newCondition).getExpression();
-                        assertNotNull(newConditionExpression);
-
                         //noinspection ConstantConditions
-                        newConditionExpression.replace(rhs);
+                        builder.condition(rhs);
                     }
                     else assert false : TRANSFORM_WITHOUT_CHECK;
                 }
                 else assert false : TRANSFORM_WITHOUT_CHECK;
             }
+
+            //noinspection ConstantConditions
+            builder.branchExpression(branchExpression);
         }
 
-        whenExpression.replace(newWhenExpression);
+        whenExpression.replace(builder.toExpression(whenExpression.getProject()));
     }
 
-    private static JetWhenExpression createWhenTemplateWithoutSubject(@NotNull JetWhenExpression whenExpression) {
+    static String whenConditionToExpressionText(@NotNull JetWhenCondition condition, JetExpression subject) {
-        JetPsiFactory.WhenTemplateBuilder builder = new JetPsiFactory.WhenTemplateBuilder(false);
-
-        for (JetWhenEntry entry : whenExpression.getEntries()) {
-            if (!entry.isElse()) {
-                builder.addBranchWithMultiCondition(entry.getConditions().length);
-            }
-        }
-
-        return builder.toExpression(whenExpression.getProject());
-    }
-
-    static JetExpression whenConditionToExpression(@NotNull JetWhenCondition condition, JetExpression subject) {
-        Project project = condition.getProject();
-
         if (condition instanceof JetWhenConditionIsPattern) {
             JetWhenConditionIsPattern patternCondition = (JetWhenConditionIsPattern) condition;
 
@@ -270,7 +183,7 @@ public class WhenUtils {
             assertNotNull(subject);
 
             //noinspection ConstantConditions
-            return JetPsiFactory.createIsExpression(project, subject, typeReference, patternCondition.isNegated());
+            return toBinaryExpression(subject, (patternCondition.isNegated() ? "!is" : "is"), typeReference);
         }
 
         if (condition instanceof JetWhenConditionInRange) {
@@ -282,7 +195,7 @@ public class WhenUtils {
             assertNotNull(subject);
 
             //noinspection ConstantConditions
-            return JetPsiFactory.createBinaryExpression(project, subject, rangeCondition.getOperationReference().getText(), rangeExpression);
+            return toBinaryExpression(subject, rangeCondition.getOperationReference().getText(), rangeExpression);
         }
 
         assert condition instanceof JetWhenConditionWithExpression : TRANSFORM_WITHOUT_CHECK;
@@ -291,45 +204,36 @@ public class WhenUtils {
         assertNotNull(conditionExpression);
 
         //noinspection ConstantConditions
-        return subject != null ? JetPsiFactory.createBinaryExpression(project, subject, "==", conditionExpression) : conditionExpression;
+        return subject != null ?
+               toBinaryExpression(parenthesizeIfNeeded(subject), "==", parenthesizeIfNeeded(conditionExpression))
+                               : conditionExpression.getText();
     }
 
     public static void eliminateWhenSubject(@NotNull JetWhenExpression whenExpression) {
         JetExpression subject = whenExpression.getSubjectExpression();
         assertNotNull(subject);
 
-        JetWhenExpression newWhenExpression = createWhenTemplateWithoutSubject(whenExpression);
+        JetPsiFactory.WhenBuilder builder = new JetPsiFactory.WhenBuilder();
-
-        int i = 0;
-        List<JetWhenEntry> entries = whenExpression.getEntries();
-        List<JetWhenEntry> newEntries = newWhenExpression.getEntries();
-        for (JetWhenEntry newEntry : newEntries) {
-            JetWhenEntry entry = entries.get(i++);
 
+        for (JetWhenEntry entry : whenExpression.getEntries()) {
             JetExpression branchExpression = entry.getExpression();
             assertNotNull(branchExpression);
 
-            JetExpression newBranchExpression = newEntry.getExpression();
+            if (entry.isElse()) {
-            assertNotNull(newBranchExpression);
+                //noinspection ConstantConditions
+                builder.elseEntry(branchExpression);
+
+                continue;
+            }
+
+            for (JetWhenCondition condition : entry.getConditions()) {
+                builder.condition(whenConditionToExpressionText(condition, subject));
+            }
 
             //noinspection ConstantConditions
-            newBranchExpression.replace(branchExpression);
+            builder.branchExpression(branchExpression);
-
-            int j = 0;
-            JetWhenCondition[] conditions = entry.getConditions();
-            JetWhenCondition[] newConditions = newEntry.getConditions();
-
-            for (JetWhenCondition newCondition : newConditions) {
-                JetWhenCondition condition = conditions[j++];
-
-                JetExpression newConditionExpression = ((JetWhenConditionWithExpression) newCondition).getExpression();
-                assertNotNull(newConditionExpression);
-
-                //noinspection ConstantConditions
-                newConditionExpression.replace(whenConditionToExpression(condition, subject));
-            }
         }
 
-        whenExpression.replace(newWhenExpression);
+        whenExpression.replace(builder.toExpression(whenExpression.getProject()));
     }
 }

idea/testData/codeInsight/codeTransformations/branched/folding/whenToReturn/simpleWhen.kt.after

@@ -1,6 +1,6 @@
 fun test(n: Int): String {
     <caret>return when (n) {
-      1 -> "one"
+        1 -> "one"
-      else -> "two"
+        else -> "two"
     }
 }

idea/testData/codeInsight/codeTransformations/branched/folding/whenToReturn/simpleWhenWithBlocks.kt

@@ -1,11 +1,12 @@
 fun test(n: Int): String {
     <caret>when(n) {
-    1 -> {
+        1 -> {
-        println("***")
+            println("***")
-        return "one"
+            return "one"
-    }
+        }
-    else -> {
+        else -> {
-        println("***")
+            println("***")
-        return "two"
+            return "two"
+        }
     }
 }

idea/testData/codeInsight/codeTransformations/branched/folding/whenToReturn/simpleWhenWithBlocks.kt.after

@@ -1,11 +1,12 @@
 fun test(n: Int): String {
     <caret>return when(n) {
-    1 -> {
+        1 -> {
-        println("***")
+            println("***")
-        "one"
+            "one"
-    }
+        }
-    else -> {
+        else -> {
-        println("***")
+            println("***")
-        "two"
+            "two"
+        }
     }
 }

idea/testData/codeInsight/codeTransformations/branched/ifWhen/whenToIf/whenWithEqualityTests.kt.after

@@ -1,9 +1,6 @@
 fun test(n: Int): String {
-    return <caret>if (n == 0)
+    return <caret>if (n == 0) "zero"
-        "zero"
+    else if (n == 1) "one"
-    else if (n == 1)
+    else if (n == 2) "two"
-        "one"
-    else if (n == 2)
-        "two"
     else "unknown"
 }

idea/testData/codeInsight/codeTransformations/branched/ifWhen/whenToIf/whenWithMultiConditions.kt.after

@@ -1,9 +1,6 @@
 fun test(n: Int): String {
-    return <caret>if ((n < 0) || (n > 1000))
+    return <caret>if ((n < 0) || (n > 1000)) "unknown"
-        "unknown"
+    else if (n <= 10) "small"
-    else if (n <= 10)
+    else if (n <= 100) "average"
-        "small"
-    else if (n <= 100)
-        "average"
     else "big"
 }

idea/testData/codeInsight/codeTransformations/branched/ifWhen/whenToIf/whenWithNegativePatterns.kt.after

@@ -1,9 +1,6 @@
 fun test(obj: Any): String {
-    return <caret>if (obj !is Iterable<*>)
+    return <caret>if (obj !is Iterable<*>) "not iterable"
-        "not iterable"
+    else if (obj !is Collection<*>) "not collection"
-    else if (obj !is Collection<*>)
+    else if (obj !is MutableCollection<*>) "not mutable collection"
-        "not collection"
-    else if (obj !is MutableCollection<*>)
-        "not mutable collection"
     else "unknown"
 }

idea/testData/codeInsight/codeTransformations/branched/ifWhen/whenToIf/whenWithNegativeRangeTests.kt.after

@@ -1,9 +1,6 @@
 fun test(n: Int): String {
-    return <caret>if (n !in 0..1000)
+    return <caret>if (n !in 0..1000) "unknown"
-        "unknown"
+    else if (n !in 0..100) "big"
-    else if (n !in 0..100)
+    else if (n !in 0..10) "average"
-        "big"
-    else if (n !in 0..10)
-        "average"
     else "small"
 }

idea/testData/codeInsight/codeTransformations/branched/ifWhen/whenToIf/whenWithPatterns.kt.after

@@ -1,9 +1,6 @@
 fun test(obj: Any): String {
-    return <caret>if (obj is String)
+    return <caret>if (obj is String) "string"
-        "string"
+    else if (obj is Int) "int"
-    else if (obj is Int)
+    else if (obj is Class<*>) "class"
-        "int"
-    else if (obj is Class<*>)
-        "class"
     else "unknown"
 }

idea/testData/codeInsight/codeTransformations/branched/ifWhen/whenToIf/whenWithRangeTests.kt.after

@@ -1,9 +1,6 @@
 fun test(n: Int): String {
-    return <caret>if (n in 0..10)
+    return <caret>if (n in 0..10) "small"
-        "small"
+    else if (n in 10..100) "average"
-    else if (n in 10..100)
+    else if (n in 100..1000) "big"
-        "average"
-    else if (n in 100..1000)
-        "big"
     else "unknown"
 }

idea/testData/codeInsight/codeTransformations/branched/ifWhen/whenToIf/whenWithRangeTestsAndMultiConditions.kt.after

@@ -1,9 +1,6 @@
 fun test(n: Int): String {
-    return <caret>if ((n in 0..5) || (n in 5..10))
+    return <caret>if ((n in 0..5) || (n in 5..10)) "small"
-        "small"
+    else if ((n in 10..50) || (n in 50..100)) "average"
-    else if ((n in 10..50) || (n in 50..100))
+    else if ((n in 100..500) || (n in 500..1000)) "big"
-        "average"
-    else if ((n in 100..500) || (n in 500..1000))
-        "big"
     else "unknown"
 }

idea/testData/codeInsight/codeTransformations/branched/ifWhen/whenToIf/whenWithoutSubject.kt.after

@@ -1,9 +1,6 @@
 fun test(n: Int): String {
-    return <caret>if (n in 0..10)
+    return <caret>if (n in 0..10) "small"
-        "small"
+    else if (n in 10..100) "average"
-    else if (n in 10..100)
+    else if (n in 100..1000) "big"
-        "average"
-    else if (n in 100..1000)
-        "big"
     else "unknown"
 }