KT-1053 Incorrect auto formatting of when expression

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

@@ -84,16 +84,15 @@ public class JetBlock implements ASTBlock {
     private List<Block> buildSubBlocks() {
         List<Block> blocks = new ArrayList<Block>();
         for (ASTNode child = myNode.getFirstChildNode(); child != null; child = child.getTreeNext()) {
+            IElementType childType = child.getElementType();
 
-          IElementType childType = child.getElementType();
+            if (child.getTextRange().getLength() == 0) continue;
 
-          if (child.getTextRange().getLength() == 0) continue;
+            if (childType == TokenType.WHITE_SPACE) {
+                continue;
+            }
 
-          if (childType == TokenType.WHITE_SPACE) {
-            continue;
-          }
-
-          blocks.add(buildSubBlock(child));
+            blocks.add(buildSubBlock(child));
         }
         return Collections.unmodifiableList(blocks);
     }
@@ -106,8 +105,21 @@ public class JetBlock implements ASTBlock {
         if (CODE_BLOCKS.contains(myNode.getElementType())) {
             childIndent = indentIfNotBrace(child);
         }
-        else if (myNode.getElementType() == JetNodeTypes.WHEN) {
-            childIndent = indentIfNotBrace(child);
+        else if (child.getTreeParent() != null && child.getTreeParent().getElementType() == JetNodeTypes.BODY &&
+                 child.getElementType() != JetNodeTypes.BLOCK) {
+            // For a single statement if 'for'
+            childIndent = Indent.getNormalIndent();
+        }
+        else if (child.getElementType() == JetNodeTypes.WHEN_ENTRY) {
+            // For the entry in when
+            // TODO: Add an option for configuration?
+            childIndent = Indent.getNormalIndent();
+        }
+        else if (child.getTreeParent() != null && child.getTreeParent().getElementType() == JetNodeTypes.WHEN_ENTRY) {
+            ASTNode prev = getPrevWithoutWhitespace(child);
+            if (prev != null && prev.getText().equals("->")) {
+                childIndent = indentIfNotBrace(child);
+            }
         }
         else if (STATEMENT_PARTS.contains(myNode.getElementType()) && child.getElementType() != JetNodeTypes.BLOCK) {
             childIndent = Indent.getNormalIndent();
@@ -122,6 +134,15 @@ public class JetBlock implements ASTBlock {
                 : Indent.getNormalIndent();
     }
 
+    private static ASTNode getPrevWithoutWhitespace(ASTNode node) {
+        node = node.getTreePrev();
+        while (node != null && node.getElementType() == TokenType.WHITE_SPACE) {
+            node = node.getTreePrev();
+        }
+
+        return node;
+    }
+
     @Override
     public Spacing getSpacing(Block child1, Block child2) {
         return mySpacingBuilder.getSpacing(this, child1, child2);

idea/testData/formatter/BlockFor.kt

@@ -0,0 +1,5 @@
+fun some() {
+for (var a in 1..12) {
+a += 1
+}
+}

idea/testData/formatter/BlockFor_after.kt

@@ -0,0 +1,5 @@
+fun some() {
+    for (var a in 1..12) {
+        a += 1
+    }
+}

idea/testData/formatter/Class.kt

@@ -0,0 +1,3 @@
+class Some() {
+val a : Int
+}

idea/testData/formatter/Class_after.kt

@@ -0,0 +1,3 @@
+class Some() {
+    val a : Int
+}

idea/testData/formatter/ForNoBraces.kt

@@ -0,0 +1,6 @@
+fun testSome() {
+var a = 12
+for (var a in 1..10)
+a = 11
+a = 10
+}

idea/testData/formatter/ForNoBraces_after.kt

@@ -0,0 +1,6 @@
+fun testSome() {
+    var a = 12
+    for (var a in 1..10)
+        a = 11
+    a = 10
+}

idea/testData/formatter/If.kt

@@ -0,0 +1,6 @@
+fun some() {
+if (3 > 2)
+2
+else
+1
+}

idea/testData/formatter/If_after.kt

@@ -0,0 +1,6 @@
+fun some() {
+    if (3 > 2)
+        2
+    else
+        1
+}

idea/testData/formatter/When.kt

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

idea/testData/formatter/WhenEntryExpr.kt

@@ -0,0 +1,9 @@
+fun some(x : Any) {
+when (x) {
+is Int ->
+0
+else -> {
+1
+}
+}
+}

idea/testData/formatter/WhenEntryExpr_after.kt

@@ -0,0 +1,9 @@
+fun some(x : Any) {
+    when (x) {
+        is Int ->
+            0
+        else -> {
+            1
+        }
+    }
+}

idea/testData/formatter/When_after.kt

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

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

@@ -0,0 +1,191 @@
+/*
+ * Copyright 2000-2010 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.formatter;
+
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.command.CommandProcessor;
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.editor.EditorFactory;
+import com.intellij.openapi.editor.impl.DocumentImpl;
+import com.intellij.openapi.roots.LanguageLevelProjectExtension;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.io.FileUtil;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.pom.java.LanguageLevel;
+import com.intellij.psi.PsiDocumentManager;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.codeStyle.CodeStyleManager;
+import com.intellij.testFramework.LightIdeaTestCase;
+import com.intellij.util.IncorrectOperationException;
+import org.jetbrains.annotations.NonNls;
+import org.jetbrains.jet.plugin.PluginTestCaseBase;
+
+import java.io.File;
+import java.util.EnumMap;
+import java.util.Map;
+
+// Based on from com.intellij.psi.formatter.java.AbstractJavaFormatterTest
+@SuppressWarnings("UnusedDeclaration")
+public abstract class AbstractKotlinFormatterTest extends LightIdeaTestCase {
+
+    protected enum Action {REFORMAT, INDENT}
+
+    private interface TestFormatAction {
+        void run(PsiFile psiFile, int startOffset, int endOffset);
+    }
+
+    private static final Map<Action, TestFormatAction> ACTIONS = new EnumMap<Action, TestFormatAction>(Action.class);
+    static {
+        ACTIONS.put(Action.REFORMAT, new TestFormatAction() {
+            @Override
+            public void run(PsiFile psiFile, int startOffset, int endOffset) {
+                CodeStyleManager.getInstance(getProject()).reformatText(psiFile, startOffset, endOffset);
+            }
+        });
+        ACTIONS.put(Action.INDENT, new TestFormatAction() {
+            @Override
+            public void run(PsiFile psiFile, int startOffset, int endOffset) {
+                CodeStyleManager.getInstance(getProject()).adjustLineIndent(psiFile, startOffset);
+            }
+        });
+    }
+
+    private static final String BASE_PATH =
+            new File(PluginTestCaseBase.getTestDataPathBase(), "/formatter/").getAbsolutePath();
+
+    public TextRange myTextRange;
+    public TextRange myLineRange;
+
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        LanguageLevelProjectExtension.getInstance(getProject()).setLanguageLevel(LanguageLevel.HIGHEST);
+    }
+
+//    public static CommonCodeStyleSettings getSettings() {
+//        CodeStyleSettings rootSettings = CodeStyleSettingsManager.getSettings(getProject());
+//        return rootSettings.getCommonSettings(JavaLanguage.INSTANCE);
+//    }
+//
+//    public static CommonCodeStyleSettings.IndentOptions getIndentOptions() {
+//        return getSettings().getRootSettings().getIndentOptions(StdFileTypes.JAVA);
+//    }
+
+    public void doTest() throws Exception {
+        doTest(getTestName(false) + ".kt", getTestName(false) + "_after.kt");
+    }
+
+    public void doTest(@NonNls String fileNameBefore, @NonNls String fileNameAfter) throws Exception {
+        doTextTest(Action.REFORMAT, loadFile(fileNameBefore), loadFile(fileNameAfter));
+    }
+
+    public void doTextTest(@NonNls final String text, @NonNls String textAfter) throws IncorrectOperationException {
+        doTextTest(Action.REFORMAT, text, textAfter);
+    }
+
+    public void doTextTest(final Action action, final String text, String textAfter) throws IncorrectOperationException {
+        final PsiFile file = createFile("A.kt", text);
+
+        if (myLineRange != null) {
+            final DocumentImpl document = new DocumentImpl(text);
+            myTextRange =
+                    new TextRange(document.getLineStartOffset(myLineRange.getStartOffset()), document.getLineEndOffset(myLineRange.getEndOffset()));
+        }
+
+        final PsiDocumentManager manager = PsiDocumentManager.getInstance(getProject());
+        final Document document = manager.getDocument(file);
+
+        CommandProcessor.getInstance().executeCommand(getProject(), new Runnable() {
+            @Override
+            public void run() {
+                ApplicationManager.getApplication().runWriteAction(new Runnable() {
+                    @Override
+                    public void run() {
+                        document.replaceString(0, document.getTextLength(), text);
+                        manager.commitDocument(document);
+                        try {
+                            TextRange rangeToUse = myTextRange;
+                            if (rangeToUse == null) {
+                                rangeToUse = file.getTextRange();
+                            }
+                            ACTIONS.get(action).run(file, rangeToUse.getStartOffset(), rangeToUse.getEndOffset());
+                        }
+                        catch (IncorrectOperationException e) {
+                            assertTrue(e.getLocalizedMessage(), false);
+                        }
+                    }
+                });
+            }
+        }, "", "");
+
+
+        if (document == null) {
+            fail("Don't expect the document to be null");
+            return;
+        }
+        assertEquals(prepareText(textAfter), prepareText(document.getText()));
+        manager.commitDocument(document);
+        assertEquals(prepareText(textAfter), prepareText(file.getText()));
+
+    }
+
+//    public void doMethodTest(@NonNls final String before, @NonNls final String after) throws Exception {
+//        doTextTest(
+//                Action.REFORMAT,
+//                "class Foo{\n" + "    void foo() {\n" + before + '\n' + "    }\n" + "}",
+//                "class Foo {\n" + "    void foo() {\n" + StringUtil.shiftIndentInside(after, 8, false) + '\n' + "    }\n" + "}"
+//        );
+//    }
+//
+//    public void doClassTest(@NonNls final String before, @NonNls final String after) throws Exception {
+//        doTextTest(
+//                Action.REFORMAT,
+//                "class Foo{\n" + before + '\n' + "}",
+//                "class Foo {\n" + StringUtil.shiftIndentInside(after, 4, false) + '\n' + "}"
+//        );
+//    }
+
+    private static String prepareText(String actual) {
+        if (actual.startsWith("\n")) {
+            actual = actual.substring(1);
+        }
+        if (actual.startsWith("\n")) {
+            actual = actual.substring(1);
+        }
+
+        // Strip trailing spaces
+        final Document doc = EditorFactory.getInstance().createDocument(actual);
+        CommandProcessor.getInstance().executeCommand(getProject(), new Runnable() {
+            @Override
+            public void run() {
+                ApplicationManager.getApplication().runWriteAction(new Runnable() {
+                    @Override
+                    public void run() {
+                        ((DocumentImpl)doc).stripTrailingSpaces();
+                    }
+                });
+            }
+        }, "formatting", null);
+
+        return doc.getText();
+    }
+
+    private static String loadFile(String name) throws Exception {
+        String text = FileUtil.loadFile(new File(BASE_PATH, name));
+        text = StringUtil.convertLineSeparators(text);
+        return text;
+    }
+}

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

@@ -0,0 +1,30 @@
+package org.jetbrains.jet.formatter;
+
+/**
+ * Based on com.intellij.psi.formatter.java.JavaFormatterTest
+ */
+public class KotlinFormatterTest extends AbstractKotlinFormatterTest {
+    public void testBlockFor() throws Exception {
+        doTest();
+    }
+
+    public void testClass() throws Exception {
+        doTest();
+    }
+
+    public void testForNoBraces() throws Exception {
+        doTest();
+    }
+
+    public void testIf() throws Exception {
+        doTest();
+    }
+
+    public void testWhen() throws Exception {
+        doTest();
+    }
+
+    public void testWhenEntryExpr() throws Exception {
+        doTest();
+    }
+}