Smart enter for 'when'

#KT-3600 In Progress

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetWhenExpression.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.
@@ -47,13 +47,28 @@ public class JetWhenExpression extends JetExpressionImpl {
 
     @NotNull
     public PsiElement getWhenKeywordElement() {
+        //noinspection ConstantConditions
         return findChildByType(JetTokens.WHEN_KEYWORD);
     }
 
     @Nullable
-    public PsiElement getCloseBraceNode() {
+    public PsiElement getCloseBrace() {
-        ASTNode openBraceNode = getNode().findChildByType(JetTokens.RBRACE);
+        return findChildByType(JetTokens.RBRACE);
-        return openBraceNode != null ? openBraceNode.getPsi() : null;
+    }
+
+    @Nullable
+    public PsiElement getOpenBrace() {
+        return findChildByType(JetTokens.LBRACE);
+    }
+
+    @Nullable
+    public PsiElement getLeftParenthesis() {
+        return findChildByType(JetTokens.LPAR);
+    }
+
+    @Nullable
+    public PsiElement getRightParenthesis() {
+        return findChildByType(JetTokens.RPAR);
     }
 
     @Nullable

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

@@ -16,9 +16,8 @@
 
 package org.jetbrains.jet.plugin.editor
 
+import org.jetbrains.jet.plugin.editor.fixers.*
 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
@@ -31,11 +30,7 @@ import org.jetbrains.jet.lang.psi.JetBlockExpression
 import org.jetbrains.jet.lang.psi.JetExpression
 import org.jetbrains.jet.lang.psi.JetDeclarationWithBody
 import org.jetbrains.jet.lang.psi.JetIfExpression
-import org.jetbrains.jet.plugin.editor.fixers.KotlinWhileConditionFixer
 import org.jetbrains.jet.lang.psi.JetWhileExpression
-import org.jetbrains.jet.plugin.editor.fixers.isWithCaret
-import org.jetbrains.jet.plugin.editor.fixers.KotlinForConditionFixer
-import org.jetbrains.jet.plugin.editor.fixers.KotlinMissingForOrWhileBodyFixer
 import org.jetbrains.jet.lang.psi.JetForExpression
 import com.intellij.psi.tree.TokenSet
 import org.jetbrains.jet.JetNodeTypes
@@ -48,7 +43,10 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
 
                 KotlinWhileConditionFixer(),
                 KotlinForConditionFixer(),
-                KotlinMissingForOrWhileBodyFixer()
+                KotlinMissingForOrWhileBodyFixer(),
+
+                KotlinWhenSubjectCaretFixer(),
+                KotlinMissingWhenBodyFixer()
         )
 
         addEnterProcessors(KotlinPlainEnterProcessor())

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

@@ -0,0 +1,49 @@
+/*
+ * 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 org.jetbrains.jet.lang.psi.JetWhenExpression
+
+public class KotlinMissingWhenBodyFixer : SmartEnterProcessorWithFixers.Fixer<KotlinSmartEnterHandler>() {
+    override fun apply(editor: Editor, processor: KotlinSmartEnterHandler, element: PsiElement) {
+        if (element !is JetWhenExpression) return
+        val whenExpression = element as JetWhenExpression
+
+        val doc = editor.getDocument()
+
+        val openBrace = whenExpression.getOpenBrace()
+        val closeBrace = whenExpression.getCloseBrace()
+
+        if (openBrace == null && closeBrace == null && whenExpression.getEntries().isEmpty()) {
+            val openBraceAfter = whenExpression.insertOpenBraceAfter()
+            if (openBraceAfter != null) {
+                doc.insertString(openBraceAfter.range.end, "{}")
+            }
+        }
+    }
+
+    fun JetWhenExpression.insertOpenBraceAfter(): PsiElement? = when {
+        getRightParenthesis() != null -> getRightParenthesis()
+        getSubjectExpression() != null -> null
+        getLeftParenthesis() != null -> null
+        else -> getWhenKeywordElement()
+    }
+}

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

@@ -0,0 +1,37 @@
+/*
+ * 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.lang.SmartEnterProcessorWithFixers
+import org.jetbrains.jet.plugin.editor.KotlinSmartEnterHandler
+import com.intellij.openapi.editor.Editor
+import org.jetbrains.jet.lang.psi.JetWhenExpression
+
+public class KotlinWhenSubjectCaretFixer : SmartEnterProcessorWithFixers.Fixer<KotlinSmartEnterHandler>() {
+    override fun apply(editor: Editor, processor: KotlinSmartEnterHandler, element: PsiElement) {
+        if (element !is JetWhenExpression) return
+
+        val lParen = element.getLeftParenthesis()
+        val rParen = element.getRightParenthesis()
+        val subject = element.getSubjectExpression()
+
+        if (subject == null && lParen != null && rParen != null) {
+            processor.registerUnresolvedError(lParen.range.end)
+        }
+    }
+}

idea/src/org/jetbrains/jet/plugin/quickfix/AddWhenElseBranchFix.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.
@@ -55,12 +55,12 @@ public class AddWhenElseBranchFix extends JetIntentionAction<JetWhenExpression>
 
     @Override
     public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
-        return super.isAvailable(project, editor, file) && element.getCloseBraceNode() != null;
+        return super.isAvailable(project, editor, file) && element.getCloseBrace() != null;
     }
 
     @Override
     public void invoke(@NotNull Project project, Editor editor, JetFile file) throws IncorrectOperationException {
-        PsiElement whenCloseBrace = element.getCloseBraceNode();
+        PsiElement whenCloseBrace = element.getCloseBrace();
         assert (whenCloseBrace != null) : "isAvailable should check if close brace exist";
 
         JetWhenEntry entry = JetPsiFactory.createWhenEntry(project, ELSE_ENTRY_TEXT);

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

@@ -457,6 +457,109 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             """
     )
 
+
+
+    fun testWhen() = doFunTest(
+            """
+            when <caret>
+            """
+            ,
+            """
+            when {
+                <caret>
+            }
+            """
+    )
+
+    fun testWhen1() = doFunTest(
+            """
+            when<caret>
+            """
+            ,
+            """
+            when {
+                <caret>
+            }
+            """
+    )
+
+    fun testWhen2() = doFunTest(
+            """
+            when (true<caret>) {
+            }
+            """
+            ,
+            """
+            when (true) {
+                <caret>
+            }
+            """
+    )
+
+    fun testWhen3() = doFunTest(
+            """
+            when (true) {<caret>
+            """
+            ,
+            """
+            when (true) {
+                <caret>
+            }
+            """
+    )
+
+    fun testWhen4() = doFunTest(
+            """
+            when (true<caret>) {
+                false -> println("false")
+            }
+            """
+            ,
+            """
+            when (true) {
+                <caret>
+                false -> println("false")
+            }
+            """
+    )
+
+    fun testWhen5() = doFunTest(
+            """
+            when (<caret>)
+            """
+            ,
+            """
+            when (<caret>) {
+            }
+            """
+    )
+
+    fun testWhen6() = doFunTest(
+            """
+            when (true<caret>)
+            """
+            ,
+            """
+            when (true) {
+                <caret>
+            }
+            """
+    )
+
+    // Check that no addition {} inserted
+    fun testWhenBadParsed() = doFunTest(
+            """
+            when ( {<caret>
+            }
+            """
+            ,
+            """
+            when ( {
+            }
+                <caret>
+            """
+    )
+
     fun doFunTest(before: String, after: String) {
         fun String.withFunContext(): String {
             val bodyText = "//----${this.trimIndent()}//----"