New formatter option: Align in columns 'case' branches

idea/src/org/jetbrains/jet/plugin/formatter/JetBlock.java

@@ -17,6 +17,7 @@
 package org.jetbrains.jet.plugin.formatter;
 
 import com.intellij.formatting.*;
+import com.intellij.formatting.alignment.AlignmentStrategy;
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.TokenType;
 import com.intellij.psi.codeStyle.CodeStyleSettings;
@@ -37,6 +38,7 @@ import static org.jetbrains.jet.lexer.JetTokens.*;
  * @see Block for good JavaDoc documentation
  */
 public class JetBlock extends AbstractBlock {
+    private final ASTAlignmentStrategy myAlignmentStrategy;
     private final Indent myIndent;
     private final CodeStyleSettings mySettings;
     private final SpacingBuilder mySpacingBuilder;
@@ -51,13 +53,14 @@ public class JetBlock extends AbstractBlock {
     // private static final List<IndentWhitespaceRule>
 
     public JetBlock(@NotNull ASTNode node,
-            Alignment alignment,
+            ASTAlignmentStrategy alignmentStrategy,
             Indent indent,
             Wrap wrap,
             CodeStyleSettings settings,
             SpacingBuilder spacingBuilder) {
 
-        super(node, wrap, alignment);
+        super(node, wrap, alignmentStrategy.getAlignment(node));
+        myAlignmentStrategy = alignmentStrategy;
         myIndent = indent;
         mySettings = settings;
         mySpacingBuilder = spacingBuilder;
@@ -90,25 +93,24 @@ public class JetBlock extends AbstractBlock {
                 continue;
             }
 
-            Alignment childAlignment = childrenAlignmentStrategy.getAlignment(child);
+            blocks.add(buildSubBlock(child, childrenAlignmentStrategy));
-            blocks.add(buildSubBlock(child, childAlignment));
         }
         return Collections.unmodifiableList(blocks);
     }
 
     @NotNull
-    private Block buildSubBlock(@NotNull ASTNode child, Alignment childAlignment) {
+    private Block buildSubBlock(@NotNull ASTNode child, ASTAlignmentStrategy alignmentStrategy) {
         Wrap wrap = null;
 
         // Affects to spaces around operators...
         if (child.getElementType() == OPERATION_REFERENCE) {
             ASTNode operationNode = child.getFirstChildNode();
             if (operationNode != null) {
-                return new JetBlock(operationNode, childAlignment, Indent.getNoneIndent(), wrap, mySettings, mySpacingBuilder);
+                return new JetBlock(operationNode, alignmentStrategy, Indent.getNoneIndent(), wrap, mySettings, mySpacingBuilder);
             }
         }
 
-        return new JetBlock(child, childAlignment, createChildIndent(child), wrap, mySettings, mySpacingBuilder);
+        return new JetBlock(child, alignmentStrategy, createChildIndent(child), wrap, mySettings, mySpacingBuilder);
     }
 
     private static Indent indentIfNotBrace(@NotNull ASTNode child) {
@@ -221,9 +223,10 @@ public class JetBlock extends AbstractBlock {
 
     private ASTAlignmentStrategy getChildrenAlignmentStrategy() {
         CommonCodeStyleSettings jetCommonSettings = mySettings.getCommonSettings(JetLanguage.INSTANCE);
+        JetCodeStyleSettings jetSettings = mySettings.getCustomSettings(JetCodeStyleSettings.class);
 
         // Prepare default null strategy
-        ASTAlignmentStrategy strategy = ASTAlignmentStrategy.getNullStrategy();
+        ASTAlignmentStrategy strategy = myAlignmentStrategy;
 
         // Redefine list of strategies for some special elements
         IElementType parentType = myNode.getElementType();
@@ -237,6 +240,9 @@ public class JetBlock extends AbstractBlock {
                     jetCommonSettings.ALIGN_MULTILINE_PARAMETERS_IN_CALLS, VALUE_ARGUMENT, COMMA,
                     jetCommonSettings.ALIGN_MULTILINE_METHOD_BRACKETS, LPAR, RPAR);
         }
+        else if (parentType == WHEN) {
+            strategy = getAlignmentForCaseBranch(jetSettings.ALIGN_IN_COLUMNS_CASE_BRANCH);
+        }
         return strategy;
     }
 
@@ -271,6 +277,16 @@ public class JetBlock extends AbstractBlock {
         };
     }
 
+    private static ASTAlignmentStrategy getAlignmentForCaseBranch(boolean shouldAlignInColumns) {
+        if (shouldAlignInColumns) {
+            return ASTAlignmentStrategy
+                    .fromTypes(AlignmentStrategy.createAlignmentPerTypeStrategy(Arrays.asList((IElementType) ARROW), WHEN_ENTRY, true));
+        }
+        else {
+            return ASTAlignmentStrategy.getNullStrategy();
+        }
+    }
+
     static ASTIndentStrategy[] INDENT_RULES = new ASTIndentStrategy[] {
             ASTIndentStrategy.forNode("No indent for braces in blocks")
                     .in(BLOCK, CLASS_BODY, FUNCTION_LITERAL)

idea/src/org/jetbrains/jet/plugin/formatter/JetCodeStyleSettings.java

@@ -32,6 +32,7 @@ public class JetCodeStyleSettings extends CustomCodeStyleSettings {
     public boolean SPACE_AFTER_EXTEND_COLON = true;
 
     public boolean INSERT_WHITESPACES_IN_SIMPLE_ONE_LINE_METHOD = true;
+    public boolean ALIGN_IN_COLUMNS_CASE_BRANCH = false;
 
     public static JetCodeStyleSettings getInstance(Project project) {
         return CodeStyleSettingsManager.getSettings(project).getCustomSettings(JetCodeStyleSettings.class);

idea/src/org/jetbrains/jet/plugin/formatter/JetFormattingModelBuilder.java

@@ -35,7 +35,7 @@ public class JetFormattingModelBuilder implements FormattingModelBuilder {
     @Override
     public FormattingModel createModel(PsiElement element, CodeStyleSettings settings) {
         final JetBlock block = new JetBlock(
-            element.getNode(), null, Indent.getNoneIndent(), null, settings,
+            element.getNode(), ASTAlignmentStrategy.getNullStrategy(), Indent.getNoneIndent(), null, settings,
             createSpacingBuilder(settings));
 
         return FormattingModelProvider.createFormattingModelForPsiFile(
@@ -104,6 +104,7 @@ public class JetFormattingModelBuilder implements FormattingModelBuilder {
                 .afterInside(COLON, TYPE_PARAMETER).spaceIf(jetSettings.SPACE_AFTER_EXTEND_COLON)
 
                 .between(VALUE_ARGUMENT_LIST, FUNCTION_LITERAL_EXPRESSION).spaces(1)
+                .aroundInside(ARROW, WHEN_ENTRY).spaces(1)
                 ;
     }
 

idea/src/org/jetbrains/jet/plugin/formatter/JetLanguageCodeStyleSettingsProvider.java

@@ -41,6 +41,10 @@ public class JetLanguageCodeStyleSettingsProvider extends LanguageCodeStyleSetti
                         "    private fun foo1(i1: Int,\n" +
                         "                     i2: Int,\n" +
                         "                     i3: Int) : Int {\n" +
+                        "        when (i1) {\n" +
+                        "            is Number -> 0\n" +
+                        "            else -> 1\n" +
+                        "        }\n" +
                         "        return 0\n" +
                         "    }\n" +
                         "    private fun foo2():Int {\n" +
@@ -116,6 +120,9 @@ public class JetLanguageCodeStyleSettingsProvider extends LanguageCodeStyleSetti
                         "ALIGN_MULTILINE_PARAMETERS_IN_CALLS",
                         "ALIGN_MULTILINE_METHOD_BRACKETS"
                 );
+                consumer.renameStandardOption(CodeStyleSettingsCustomizable.WRAPPING_SWITCH_STATEMENT, "'when' statements");
+                consumer.showCustomOption(JetCodeStyleSettings.class, "ALIGN_IN_COLUMNS_CASE_BRANCH", "Align in columns 'case' branches",
+                                          CodeStyleSettingsCustomizable.WRAPPING_SWITCH_STATEMENT);
                 break;
             default:
                 consumer.showStandardOptions();

idea/testData/formatter/When.kt

@@ -1,6 +1,8 @@
 fun some(x : Any) {
 when (x) {
-is Int -> 0
+is Number     ->     0
-else -> 1
+else->1
 }
-}
+}
+
+// SET_FALSE: ALIGN_IN_COLUMNS_CASE_BRANCH

idea/testData/formatter/When_after.kt

@@ -1,6 +1,8 @@
 fun some(x: Any) {
     when (x) {
-        is Int -> 0
+        is Number -> 0
         else -> 1
     }
-}
+}
+
+// SET_FALSE: ALIGN_IN_COLUMNS_CASE_BRANCH

idea/testData/formatter/When_after_inv.kt

@@ -0,0 +1,8 @@
+fun some(x: Any) {
+    when (x) {
+        is Number -> 0
+        else      -> 1
+    }
+}
+
+// SET_FALSE: ALIGN_IN_COLUMNS_CASE_BRANCH

idea/tests/org/jetbrains/jet/formatter/JetFormatterTest.java

@@ -121,7 +121,7 @@ public class JetFormatterTest extends AbstractJetFormatterTest {
     }
 
     public void testWhen() throws Exception {
-        doTest();
+        doTestWithInvert();
     }
 
     public void testWhenEntryExpr() throws Exception {