Don't auto insert '}' after '{' in unfinished control statements

idea/src/org/jetbrains/jet/plugin/editor/KotlinTypedHandler.java

@@ -18,18 +18,38 @@ package org.jetbrains.jet.plugin.editor;
 
 import com.intellij.codeInsight.CodeInsightSettings;
 import com.intellij.codeInsight.editorActions.TypedHandlerDelegate;
+import com.intellij.codeInsight.highlighting.BraceMatcher;
+import com.intellij.codeInsight.highlighting.BraceMatchingUtil;
+import com.intellij.lang.ASTNode;
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.EditorModificationUtil;
+import com.intellij.openapi.editor.ScrollType;
+import com.intellij.openapi.editor.ex.EditorEx;
+import com.intellij.openapi.editor.highlighter.EditorHighlighter;
+import com.intellij.openapi.editor.highlighter.HighlighterIterator;
 import com.intellij.openapi.fileTypes.FileType;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.TokenType;
+import com.intellij.psi.codeStyle.CodeStyleManager;
+import com.intellij.psi.formatter.FormatterUtil;
 import com.intellij.psi.impl.source.tree.LeafPsiElement;
+import com.intellij.psi.tree.TokenSet;
+import com.intellij.util.text.CharArrayUtil;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.JetNodeTypes;
 import org.jetbrains.jet.lang.psi.JetFile;
 import org.jetbrains.jet.lexer.JetTokens;
 
 public class KotlinTypedHandler extends TypedHandlerDelegate {
+    private final static TokenSet CONTROL_FLOW_EXPRESSIONS = TokenSet.create(
+            JetNodeTypes.IF,
+            JetNodeTypes.FOR,
+            JetNodeTypes.WHILE);
 
     private boolean jetLTTyped;
 
@@ -50,6 +70,42 @@ public class KotlinTypedHandler extends TypedHandlerDelegate {
                         return Result.STOP;
                     }
                 }
+                return Result.CONTINUE;
+            case '{':
+                // Returning Result.CONTINUE will cause inserting "{}" for unmatched '{'
+
+                int offset = editor.getCaretModel().getOffset();
+                if (offset == 0) {
+                    return Result.CONTINUE;
+                }
+
+                HighlighterIterator iterator = ((EditorEx) editor).getHighlighter().createIterator(offset - 1);
+                while (!iterator.atEnd() && iterator.getTokenType() == TokenType.WHITE_SPACE) {
+                    iterator.retreat();
+                }
+
+                if (iterator.atEnd() || iterator.getTokenType() != JetTokens.RPAR) {
+                    return Result.CONTINUE;
+                }
+
+                PsiDocumentManager.getInstance(project).commitDocument(editor.getDocument());
+
+                PsiElement leaf = file.findElementAt(offset);
+                if (leaf != null) {
+                    PsiElement parent = leaf.getParent();
+                    if (parent != null && CONTROL_FLOW_EXPRESSIONS.contains(parent.getNode().getElementType())) {
+                        ASTNode nonWhitespaceSibling = FormatterUtil.getPreviousNonWhitespaceSibling(leaf.getNode());
+                        if (nonWhitespaceSibling != null && nonWhitespaceSibling.getText().equals(")")) {
+                            // Check that ')' belongs to same parent
+
+                            EditorModificationUtil.insertStringAtCaret(editor, "{", false, true);
+                            indentBrace(project, editor, '{');
+
+                            return Result.STOP;
+                        }
+                    }
+                }
+
                 return Result.CONTINUE;
         }
 
@@ -82,4 +138,45 @@ public class KotlinTypedHandler extends TypedHandlerDelegate {
 
         return Result.CONTINUE;
     }
+
+    /**
+     * Copied from
+     * @see com.intellij.codeInsight.editorActions.TypedHandler#indentBrace(Project, Editor, char)
+     */
+    private static void indentBrace(@NotNull final Project project, @NotNull final Editor editor, char braceChar) {
+        final int offset = editor.getCaretModel().getOffset() - 1;
+        Document document = editor.getDocument();
+        CharSequence chars = document.getCharsSequence();
+        if (offset < 0 || chars.charAt(offset) != braceChar) return;
+
+        int spaceStart = CharArrayUtil.shiftBackward(chars, offset - 1, " \t");
+        if (spaceStart < 0 || chars.charAt(spaceStart) == '\n' || chars.charAt(spaceStart) == '\r') {
+            PsiDocumentManager documentManager = PsiDocumentManager.getInstance(project);
+            documentManager.commitDocument(document);
+
+            final PsiFile file = documentManager.getPsiFile(document);
+            if (file == null || !file.isWritable()) return;
+            PsiElement element = file.findElementAt(offset);
+            if (element == null) return;
+
+            EditorHighlighter highlighter = ((EditorEx) editor).getHighlighter();
+            HighlighterIterator iterator = highlighter.createIterator(offset);
+
+            FileType fileType = file.getFileType();
+            BraceMatcher braceMatcher = BraceMatchingUtil.getBraceMatcher(fileType, iterator);
+            boolean isBrace =
+                    braceMatcher.isLBraceToken(iterator, chars, fileType) || braceMatcher.isRBraceToken(iterator, chars, fileType);
+            if (element.getNode() != null && isBrace) {
+                ApplicationManager.getApplication().runWriteAction(new Runnable() {
+                    @Override
+                    public void run() {
+                        int newOffset = CodeStyleManager.getInstance(project).adjustLineIndent(file, offset);
+                        editor.getCaretModel().moveToOffset(newOffset + 1);
+                        editor.getScrollingModel().scrollToCaret(ScrollType.RELATIVE);
+                        editor.getSelectionModel().removeSelection();
+                    }
+                });
+            }
+        }
+    }
 }

idea/tests/org/jetbrains/jet/editor/TypedHandlerTest.java

@@ -91,6 +91,184 @@ public class TypedHandlerTest extends LightCodeInsightTestCase {
         );
     }
 
+    public void testAutoCloseBraceInFunctionDeclaration() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() <caret>",
+
+                "fun foo() {<caret>}"
+        );
+    }
+
+    public void testAutoCloseBraceInLocalFunctionDeclaration() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    fun bar() <caret>\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    fun bar() {<caret>}\n" +
+                "}"
+        );
+    }
+
+    public void testAutoCloseBraceInAssignment() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    val a = <caret>\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    val a = {<caret>}\n" +
+                "}"
+        );
+    }
+
+    public void testDoNotAutoCloseBraceInUnfinishedIfSurroundOnSameLine() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    if() <caret>foo()\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    if() {foo()\n" +
+                "}"
+        );
+    }
+
+    public void testDoNotAutoCloseBraceInUnfinishedWhileSurroundOnSameLine() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    while() <caret>foo()\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    while() {foo()\n" +
+                "}"
+        );
+    }
+
+    public void testDoNotAutoCloseBraceInUnfinishedWhileSurroundOnNewLine() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    while()\n" +
+                "<caret>\n" +
+                "    foo()\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    while()\n" +
+                "    {\n" +
+                "    foo()\n" +
+                "}"
+        );
+    }
+
+    public void testDoNotAutoCloseBraceInUnfinishedIfSurroundOnOtherLine() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    if(true) <caret>\n" +
+                "    foo()\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    if(true) {<caret>\n" +
+                "    foo()\n" +
+                "}"
+        );
+    }
+
+    public void testDoNotAutoCloseBraceInUnfinishedIfSurroundOnNewLine() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    if(true)\n" +
+                "        <caret>\n" +
+                "    foo()\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    if(true)\n" +
+                "    {<caret>\n" +
+                "    foo()\n" +
+                "}"
+        );
+    }
+
+    public void testAutoCloseBraceInsideFor() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    for (elem in some.filter <caret>) {\n" +
+                "    }\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    for (elem in some.filter {<caret>}) {\n" +
+                "    }\n" +
+                "}"
+        );
+    }
+
+    public void testAutoCloseBraceInsideForAfterCloseParen() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    for (elem in some.foo(true) <caret>) {\n" +
+                "    }\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    for (elem in some.foo(true) {<caret>}) {\n" +
+                "    }\n" +
+                "}"
+        );
+    }
+
+    public void testAutoCloseBraceBeforeIf() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    <caret>if (true) {}\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    {<caret>if (true) {}\n" +
+                "}"
+        );
+    }
+
+    public void testAutoCloseBraceInIfCondition() {
+        doCharTypeTest(
+                '{',
+
+                "fun foo() {\n" +
+                "    if (some.hello (12) <caret>)\n" +
+                "}",
+
+                "fun foo() {\n" +
+                "    if (some.hello (12) {<caret>})\n" +
+                "}"
+        );
+    }
+
     public void testTypeLtInFunDeclaration() throws Exception {
         doLtGtTest("fun <caret>");
     }