Smart enter for 'do-while' expression

#KT-3600 In Progress

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetDoWhileExpression.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,7 +17,10 @@
 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.lexer.JetTokens;
 
 public class JetDoWhileExpression extends JetWhileExpressionBase {
     public JetDoWhileExpression(@NotNull ASTNode node) {
@@ -28,4 +31,11 @@ public class JetDoWhileExpression extends JetWhileExpressionBase {
     public <R, D> R accept(@NotNull JetVisitor<R, D> visitor, D data) {
         return visitor.visitDoWhileExpression(this, data);
     }
+
+    @Nullable
+    @IfNotParsed
+    public PsiElement getWhileKeywordElement() {
+        //noinspection ConstantConditions
+        return findChildByType(JetTokens.WHILE_KEYWORD);
+    }
 }

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

@@ -30,10 +30,10 @@ 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.lang.psi.JetWhileExpression
 import org.jetbrains.jet.lang.psi.JetForExpression
 import com.intellij.psi.tree.TokenSet
 import org.jetbrains.jet.JetNodeTypes
+import org.jetbrains.jet.lang.psi.JetLoopExpression
 
 public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
     {
@@ -46,7 +46,9 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
                 KotlinMissingForOrWhileBodyFixer(),
 
                 KotlinWhenSubjectCaretFixer(),
-                KotlinMissingWhenBodyFixer()
+                KotlinMissingWhenBodyFixer(),
+
+                KotlinDoWhileFixer()
         )
 
         addEnterProcessors(KotlinPlainEnterProcessor())
@@ -115,9 +117,7 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
                     if (element.getThen().isWithCaret(caret)) return element.getThen()
                     if (element.getElse().isWithCaret(caret)) return element.getElse()
                 }
-                is JetWhileExpression -> {
+                is JetLoopExpression -> return element.getBody()
-                    if (element.getBody()?.getTextRange()?.contains(caret) == true) return element.getBody()
-                }
             }
 
             return null

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

@@ -0,0 +1,58 @@
+/*
+ * 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.JetDoWhileExpression
+import org.jetbrains.jet.lang.psi.JetBlockExpression
+
+public class KotlinDoWhileFixer : SmartEnterProcessorWithFixers.Fixer<KotlinSmartEnterHandler>() {
+    override fun apply(editor: Editor, processor: KotlinSmartEnterHandler, psiElement: PsiElement) {
+        if (psiElement !is JetDoWhileExpression) return
+
+        val doc = editor.getDocument()
+        val stmt = psiElement as JetDoWhileExpression
+        val start = stmt.range.start
+
+        val whileKeyword = stmt.getWhileKeywordElement()
+        if (stmt.getBody() == null) {
+            if (whileKeyword == null) {
+                doc.replaceString(start, start + "do".length(), "do {} while()")
+            }
+            else {
+                doc.insertString(start + "do".length(), "{}")
+            }
+            return
+        }
+
+        if (stmt.getCondition() == null) {
+            val lParen = stmt.getLeftParenthesis()
+            val rParen = stmt.getRightParenthesis()
+
+            when {
+                whileKeyword == null -> doc.insertString(stmt.range.end, "while()")
+                lParen == null && rParen == null -> {
+                    doc.replaceString(whileKeyword.range.start, whileKeyword.range.end, "while()")
+                }
+                lParen != null -> processor.registerUnresolvedError(lParen.range.end)
+            }
+        }
+    }
+}

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

@@ -457,8 +457,6 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             """
     )
 
-
-
     fun testWhen() = doFunTest(
             """
             when <caret>
@@ -560,6 +558,171 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             """
     )
 
+    fun testDoWhile() = doFunTest(
+            """
+            do <caret>
+            """
+            ,
+            """
+            do {
+            } while (<caret>)${' '}
+            """
+    )
+
+    fun testDoWhile2() = doFunTest(
+            """
+            do<caret>
+            """
+            ,
+            """
+            do {
+            } while (<caret>)
+            """
+    )
+
+    fun testDoWhile3() = doFunTest(
+            """
+            do<caret> {
+                println(hi)
+            }
+            """
+            ,
+            """
+            do {
+                println(hi)
+            } while (<caret>)
+            """
+    )
+
+    fun testDoWhile5() = doFunTest(
+            """
+            do<caret> {
+            } while ()
+            """
+            ,
+            """
+            do {
+            } while (<caret>)
+            """
+    )
+
+    fun testDoWhile6() = doFunTest(
+            """
+            do<caret> {
+            } while (true)
+            """
+            ,
+            """
+            do {
+                <caret>
+            } while (true)
+            """
+    )
+
+    fun testDoWhile7() = doFunTest(
+            """
+            do {
+            } <caret>while (true)
+            """
+            ,
+            """
+            do {
+            } while (true)
+            <caret>
+            """
+    )
+
+    fun testDoWhile8() = doFunTest(
+            """
+            do {
+            } while (<caret>true)
+            """
+            ,
+            """
+            do {
+                <caret>
+            } while (true)
+            """
+    )
+
+    fun testDoWhile9() = doFunTest(
+            """
+            do while<caret>
+            """
+            ,
+            """
+            do {
+            } while (<caret>)
+            """
+    )
+
+    fun testDoWhile10() = doFunTest(
+            """
+            do while (true<caret>)
+            """
+            ,
+            """
+            do {
+                <caret>
+            } while (true)
+            """
+    )
+
+    fun testDoWhile11() = doFunTest(
+            """
+            do {
+                println("some")
+            } while<caret>
+            """
+            ,
+            """
+            do {
+                println("some")
+            } while (<caret>)
+            """
+    )
+
+    fun testDoWhile12() = doFunTest(
+            """
+            do {
+                println("some")
+            } while (true<caret>
+            """
+            ,
+            """
+            do {
+                <caret>
+                println("some")
+            } while (true
+            """
+    )
+
+    fun testDoWhileOneLine1() = doFunTest(
+            """
+            do println("some") while (true<caret>)
+            println("hi")
+            """
+            ,
+            """
+            do println("some") while (true)
+            <caret>
+            println("hi")
+            """
+    )
+
+    fun testDoWhileOneLine2() = doFunTest(
+            """
+            do <caret>println("some") while (true)
+            println("hi")
+            """
+            ,
+            """
+            do println("some") while (true)
+            <caret>
+            println("hi")
+            """
+    )
+
     fun doFunTest(before: String, after: String) {
         fun String.withFunContext(): String {
             val bodyText = "//----${this.trimIndent()}//----"