Smart enter for if statement

#KT-3600 In Progress

annotations/com/intellij/lang/annotations.xml

@@ -2,4 +2,8 @@
   <item name='com.intellij.lang.ASTNode java.lang.String getText()'>
     <annotation name='org.jetbrains.annotations.NotNull'/>
   </item>
+  <item
+      name='com.intellij.lang.SmartEnterProcessorWithFixers.FixEnterProcessor boolean doEnter(com.intellij.psi.PsiElement, com.intellij.psi.PsiFile, com.intellij.openapi.editor.Editor, boolean) 0'>
+    <annotation name='org.jetbrains.annotations.NotNull'/>
+  </item>
 </root>

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2013 JetBrains s.r.o.
+ * Copyright 2010-2014 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.
@@ -17,9 +17,11 @@
 package org.jetbrains.jet.lang.psi;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.JetNodeTypes;
+import org.jetbrains.jet.lexer.JetTokens;
 
 public class JetIfExpression extends JetExpressionImpl {
     public JetIfExpression(@NotNull ASTNode node) {
@@ -36,6 +38,16 @@ public class JetIfExpression extends JetExpressionImpl {
         return findExpressionUnder(JetNodeTypes.CONDITION);
     }
 
+    @Nullable @IfNotParsed
+    public PsiElement getLeftParenthesis() {
+        return findChildByType(JetTokens.LPAR);
+    }
+
+    @Nullable @IfNotParsed
+    public PsiElement getRightParenthesis() {
+        return findChildByType(JetTokens.RPAR);
+    }
+
     @Nullable
     public JetExpression getThen() {
         return findExpressionUnder(JetNodeTypes.THEN);
@@ -45,4 +57,9 @@ public class JetIfExpression extends JetExpressionImpl {
     public JetExpression getElse() {
         return findExpressionUnder(JetNodeTypes.ELSE);
     }
+
+    @Nullable
+    public PsiElement getElseKeyword() {
+        return findChildByType(JetTokens.ELSE_KEYWORD);
+    }
 }

idea/src/META-INF/plugin.xml

@@ -290,6 +290,7 @@
     <typedHandler implementation="org.jetbrains.jet.plugin.editor.KotlinTypedHandler"/>
     <enterHandlerDelegate implementation="org.jetbrains.jet.plugin.editor.KotlinEnterHandler"
                           id="KotlinEnterHandler" order="before EnterBetweenBracesHandler"/>
+    <lang.smartEnterProcessor language="jet" implementationClass="org.jetbrains.jet.plugin.editor.KotlinSmartEnterHandler"/>
     <backspaceHandlerDelegate implementation="org.jetbrains.jet.plugin.editor.KotlinBackspaceHandler"/>
 
     <copyPastePostProcessor implementation="org.jetbrains.jet.plugin.conversion.copy.ConvertJavaCopyPastePostProcessor"/>

idea/src/org/jetbrains/jet/plugin/codeInsight/CodeInsightUtils.java

@@ -97,7 +97,7 @@ public class CodeInsightUtils {
         }
         if (endOffset != element2.getTextRange().getEndOffset()) return PsiElement.EMPTY_ARRAY;
 
-        ArrayList<PsiElement> array = new ArrayList<PsiElement>();
+        List<PsiElement> array = new ArrayList<PsiElement>();
         PsiElement stopElement = element2.getNextSibling();
         for (PsiElement currentElement = element1; currentElement != stopElement; currentElement = currentElement.getNextSibling()) {
             if (!(currentElement instanceof PsiWhiteSpace)) {

idea/src/org/jetbrains/jet/plugin/editor/KotlinSmartEnterHandler.kt

@@ -0,0 +1,138 @@
+/*
+ * Copyright 2010-2014 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.plugin.editor
+
+import com.intellij.lang.SmartEnterProcessorWithFixers
+import org.jetbrains.jet.plugin.editor.fixers.KotlinIfConditionFixer
+import org.jetbrains.jet.plugin.editor.fixers.KotlinMissingIfBranchFixer
+import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiFile
+import com.intellij.util.text.CharArrayUtil
+import com.intellij.psi.codeStyle.CodeStyleSettingsManager
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.jet.lang.psi.JetDeclaration
+import org.jetbrains.jet.lang.psi.JetBlockExpression
+import org.jetbrains.jet.lang.psi.JetExpression
+import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrIfStatement
+import org.jetbrains.jet.lang.psi.JetDeclarationWithBody
+import org.jetbrains.jet.lang.psi.JetIfExpression
+import com.intellij.psi.tree.TokenSet
+import org.jetbrains.jet.lexer.JetTokens
+import org.jetbrains.jet.JetNodeTypes
+
+public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
+    {
+        addFixers(
+                KotlinMissingIfBranchFixer,
+                KotlinIfConditionFixer
+        )
+
+        addEnterProcessors(KotlinPlainEnterProcessor)
+    }
+
+    override fun getStatementAtCaret(editor: Editor?, psiFile: PsiFile?): PsiElement? {
+        var atCaret = super.getStatementAtCaret(editor, psiFile)
+
+        if (atCaret is PsiWhiteSpace) return null
+
+        while (atCaret != null) {
+            if (atCaret is JetDeclaration || atCaret?.isJetStatement() == true) {
+                return atCaret
+            }
+
+            atCaret = atCaret?.getParent()
+        }
+
+        return null
+    }
+
+    override fun moveCaretInsideBracesIfAny(editor: Editor, file: PsiFile) {
+        var caretOffset = editor.getCaretModel().getOffset()
+        val chars = editor.getDocument().getCharsSequence()
+
+        if (CharArrayUtil.regionMatches(chars, caretOffset, "{}")) {
+            caretOffset += 2
+        }
+        else {
+            if (CharArrayUtil.regionMatches(chars, caretOffset, "{\n}")) {
+                caretOffset += 3
+            }
+        }
+
+        caretOffset = CharArrayUtil.shiftBackward(chars, caretOffset - 1, " \t") + 1
+
+        if (CharArrayUtil.regionMatches(chars, caretOffset - "{}".length(), "{}") ||
+                CharArrayUtil.regionMatches(chars, caretOffset - "{\n}".length(), "{\n}")) {
+            commit(editor)
+            val settings = CodeStyleSettingsManager.getSettings(file.getProject())
+            val old = settings.KEEP_SIMPLE_BLOCKS_IN_ONE_LINE
+            settings.KEEP_SIMPLE_BLOCKS_IN_ONE_LINE = false
+            val elt = PsiTreeUtil.getParentOfType(file.findElementAt(caretOffset - 1), javaClass<JetBlockExpression>())
+            reformat(elt)
+            settings.KEEP_SIMPLE_BLOCKS_IN_ONE_LINE = old
+            editor.getCaretModel().moveToOffset(caretOffset - 1)
+        }
+    }
+
+    public fun registerUnresolvedError(offset: Int) {
+        if (myFirstErrorOffset > offset) {
+            myFirstErrorOffset = offset
+        }
+    }
+
+    private fun PsiElement.isJetStatement() =
+        getParent() is JetBlockExpression || (getParent()?.getNode()?.getElementType() in IF_BRANCHES_CONTAINERS)
+
+    object KotlinPlainEnterProcessor : SmartEnterProcessorWithFixers.FixEnterProcessor() {
+        private fun getControlStatementBlock(caret: Int, element: PsiElement): JetBlockExpression? {
+            if (element is JetDeclarationWithBody) return element.getBodyExpression() as? JetBlockExpression
+            if (element is JetIfExpression) {
+                if (element.getThen()?.getTextRange()?.contains(caret) == true) {
+                    return element.getThen() as? JetBlockExpression
+                }
+
+                if (element.getElse()?.getTextRange()?.contains(caret) == true) {
+                    return element.getElse() as? JetBlockExpression
+                }
+            }
+
+            return null
+        }
+
+        override fun doEnter(atCaret: PsiElement, file: PsiFile?, editor: Editor, modified: Boolean): Boolean {
+            val block = getControlStatementBlock(editor.getCaretModel().getOffset(), atCaret)
+            if (block != null) {
+                val firstElement = block.getFirstChild()?.getNextSibling()
+
+                val offset = if (firstElement != null) {
+                    firstElement.getTextRange()!!.getStartOffset() - 1
+                } else {
+                    block.getTextRange()!!.getEndOffset()
+                }
+
+                editor.getCaretModel().moveToOffset(offset)
+            }
+
+            plainEnter(editor)
+            return true
+        }
+    }
+}
+
+private val IF_BRANCHES_CONTAINERS = TokenSet.create(JetNodeTypes.THEN, JetNodeTypes.ELSE)

idea/src/org/jetbrains/jet/plugin/editor/fixers/KotlinIfConditionFixer.kt

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2010-2014 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.plugin.editor.fixers
+
+import org.jetbrains.jet.plugin.editor.KotlinSmartEnterHandler
+import com.intellij.lang.SmartEnterProcessorWithFixers
+import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
+import org.jetbrains.jet.lang.psi.JetIfExpression
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.openapi.util.TextRange
+
+object KotlinIfConditionFixer : SmartEnterProcessorWithFixers.Fixer<KotlinSmartEnterHandler>() {
+    override fun apply(editor: Editor, processor: KotlinSmartEnterHandler, element: PsiElement) {
+        if (element !is JetIfExpression) return
+        val ifExpression = element as JetIfExpression
+
+        val doc = editor.getDocument()
+        val lParen = ifExpression.getLeftParenthesis()
+        val rParen = ifExpression.getRightParenthesis()
+        val condition = ifExpression.getCondition()
+
+        if (condition == null) {
+            if (lParen == null || rParen == null) {
+                var stopOffset = doc.getLineEndOffset(doc.getLineNumber(ifExpression.range.start))
+                val then = ifExpression.getThen()
+                if (then != null) {
+                    stopOffset = Math.min(stopOffset, then.range.start)
+                }
+
+                stopOffset = Math.min(stopOffset, ifExpression.range.end)
+
+                doc.replaceString(ifExpression.range.start, stopOffset, "if ()")
+                processor.registerUnresolvedError(ifExpression.range.start + "if (".length())
+            }
+            else {
+                processor.registerUnresolvedError(lParen.range.end)
+            }
+        }
+        else {
+            if (rParen == null) {
+                doc.insertString(condition.range.end, ")")
+            }
+        }
+    }
+}

idea/src/org/jetbrains/jet/plugin/editor/fixers/KotlinMissingIfBranchFixer.kt

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2010-2014 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.plugin.editor.fixers
+
+import com.intellij.lang.SmartEnterProcessorWithFixers
+import org.jetbrains.jet.plugin.editor.KotlinSmartEnterHandler
+import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.jet.lang.psi.JetIfExpression
+import org.jetbrains.jet.plugin.formatter.JetBlock
+import org.jetbrains.jet.lang.psi.JetBlockExpression
+
+object KotlinMissingIfBranchFixer : SmartEnterProcessorWithFixers.Fixer<KotlinSmartEnterHandler>() {
+    override fun apply(editor: Editor, processor: KotlinSmartEnterHandler, element: PsiElement) {
+        if (element !is JetIfExpression) return
+        val ifExpression = element as JetIfExpression
+
+        val document = editor.getDocument()
+        val elseBranch = ifExpression.getElse()
+        val elseKeyword = ifExpression.getElseKeyword()
+
+        if (elseKeyword != null) {
+            if (elseBranch == null || elseBranch !is JetBlockExpression && elseBranch.startLine(editor.getDocument()) > elseKeyword.startLine(editor.getDocument())) {
+                document.insertString(elseKeyword.range.end, "{}")
+                return
+            }
+        }
+
+        val thenBranch = ifExpression.getThen()
+        if (thenBranch is JetBlockExpression) return
+
+        val rParen = ifExpression.getRightParenthesis()
+        if (rParen == null) return
+
+        var transformingOneLiner = false
+        if (thenBranch != null && thenBranch.startLine(editor.getDocument()) == ifExpression.startLine(editor.getDocument())) {
+            if (ifExpression.getCondition() != null) return
+            transformingOneLiner = true
+        }
+
+        val probablyNextStatementParsedAsThen = elseKeyword == null && elseBranch == null && !transformingOneLiner
+
+        if (thenBranch == null || probablyNextStatementParsedAsThen) {
+            document.insertString(rParen.range.end, "{}")
+        }
+        else {
+            document.insertString(rParen.range.end, "{")
+            document.insertString(thenBranch.range.end + 1, "}")
+        }
+    }
+}

idea/src/org/jetbrains/jet/plugin/editor/fixers/fixersUtil.kt

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2014 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.plugin.editor.fixers
+
+import com.intellij.psi.PsiElement
+import com.intellij.openapi.util.TextRange
+import com.intellij.openapi.editor.Document
+
+val PsiElement.range: TextRange get() = getTextRange()!!
+val TextRange.start: Int get() = getStartOffset()
+val TextRange.end: Int get() = getEndOffset()
+
+fun PsiElement.startLine(doc: Document): Int = doc.getLineNumber(range.start)

idea/tests/org/jetbrains/jet/plugin/codeInsight/smartEnter/SmartEnterTest.kt

@@ -0,0 +1,287 @@
+/*
+ * Copyright 2010-2014 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.plugin.codeInsight.smartEnter
+
+import org.jetbrains.jet.plugin.JetLightCodeInsightFixtureTestCase
+import com.intellij.openapi.actionSystem.IdeActions
+import org.jetbrains.jet.plugin.JetFileType
+import org.jetbrains.jet.test.util.trimIndent
+import com.intellij.testFramework.LightProjectDescriptor
+import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase
+
+class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
+    fun testIfCondition() = doFunTest(
+            """
+            if <caret>
+            """
+            ,
+            """
+            if (<caret>) {
+            }
+            """
+    )
+
+    fun testIfCondition2() = doFunTest(
+            """
+            if<caret>
+            """
+            ,
+            """
+            if (<caret>) {
+            }
+            """
+    )
+
+    fun testIfWithFollowingCode() = doFunTest(
+            """
+            if<caret>
+
+            return true
+            """
+            ,
+            """
+            if (<caret>) {
+            }
+
+            return true
+            """
+    )
+
+    fun testIfCondition3() = doFunTest(
+            """
+            if (<caret>
+            """
+            ,
+            """
+            if (<caret>) {
+            }
+            """
+    )
+
+    fun testIfCondition4() = doFunTest(
+            """
+            if (true<caret>) {
+            }
+            """
+            ,
+            """
+            if (true) {
+                <caret>
+            }
+            """
+    )
+
+    fun testIfCondition5() = doFunTest(
+            """
+            if (true) {<caret>
+            """
+            ,
+            """
+            if (true) {
+                <caret>
+            }
+            """
+    )
+
+    fun testIfCondition6() = doFunTest(
+            """
+            if (true<caret>) {
+                println()
+            }
+            """
+            ,
+            """
+            if (true) {
+                <caret>
+                println()
+            }
+            """
+    )
+
+    fun testIfThenOneLine1() = doFunTest(
+            """
+            if (true) println()<caret>
+            """
+            ,
+            """
+            if (true) println()
+            <caret>
+            """
+    )
+
+    fun testIfThenOneLine2() = doFunTest(
+            """
+            if (true) <caret>println()
+            """
+            ,
+            """
+            if (true) println()
+            <caret>
+            """
+    )
+
+    fun testIfThenMultiLine1() = doFunTest(
+            """
+            if (true)
+                println()<caret>
+            """
+            ,
+            """
+            if (true)
+                println()
+            <caret>
+            """
+    )
+
+    fun testIfThenMultiLine2() = doFunTest(
+            """
+            if (true)
+                println()<caret>
+            """
+            ,
+            """
+            if (true)
+                println()
+            <caret>
+            """
+    )
+
+    // TODO: indent for println
+    fun testIfThenMultiLine3() = doFunTest(
+            """
+            if (true<caret>)
+                println()
+            """
+            ,
+            """
+            if (true) {
+                <caret>
+            }
+                println()
+            """
+    )
+
+    fun testIfWithReformat() = doFunTest(
+            """
+            if     (true<caret>) {
+            }
+            """,
+            """
+            if (true) {
+                <caret>
+            }
+            """
+    )
+
+    fun testElse() = doFunTest(
+            """
+            if (true) {
+            } else<caret>
+            """,
+            """
+            if (true) {
+            } else {
+                <caret>
+            }
+            """
+    )
+
+    fun testElseOneLine1() = doFunTest(
+            """
+            if (true) {
+            } else println()<caret>
+            """,
+            """
+            if (true) {
+            } else println()
+            <caret>
+            """
+    )
+
+    fun testElseOneLine2() = doFunTest(
+            """
+            if (true) {
+            } else <caret>println()
+            """,
+            """
+            if (true) {
+            } else println()
+            <caret>
+            """
+    )
+
+    fun testElseTwoLines1() = doFunTest(
+            """
+            if (true) {
+            } else
+                <caret>println()
+            """,
+            """
+            if (true) {
+            } else
+                println()
+            <caret>
+            """
+    )
+
+    fun testElseTwoLines2() = doFunTest(
+            """
+            if (true) {
+            } else
+                println()<caret>
+            """,
+            """
+            if (true) {
+            } else
+                println()
+            <caret>
+            """
+    )
+
+    // TODO: remove space in expected data
+    fun testElseWithSpace() = doFunTest(
+            """
+            if (true) {
+            } else <caret>
+            """,
+            """
+            if (true) {
+            } else {
+                <caret>
+            }${' '}
+            """
+    )
+
+    fun doFunTest(before: String, after: String) {
+        fun String.withFunContext(): String {
+            val bodyText = "//----${this.trimIndent()}//----"
+            val withIndent = bodyText.split("\n").map { "    $it" }.joinToString(separator = "\n")
+
+            return "fun method() {\n${withIndent}\n}"
+        }
+
+        doTest(before.withFunContext(), after.withFunContext())
+    }
+
+    fun doTest(before: String, after: String) {
+        myFixture.configureByText(JetFileType.INSTANCE, before)
+        myFixture.performEditorAction(IdeActions.ACTION_EDITOR_COMPLETE_STATEMENT)
+        myFixture.checkResult(after)
+    }
+
+    override fun getProjectDescriptor(): LightProjectDescriptor = LightCodeInsightFixtureTestCase.JAVA_LATEST
+}