Smart enter for one line statements in for, while and do-while

#KT-3600 In Progress

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

@@ -121,7 +121,7 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
     }
 
     private fun PsiElement.isJetStatement() =
-        getParent() is JetBlockExpression || (getParent()?.getNode()?.getElementType() in IF_BRANCHES_CONTAINERS)
+        getParent() is JetBlockExpression || (getParent()?.getNode()?.getElementType() in BRANCH_CONTAINERS)
 
     class KotlinPlainEnterProcessor : SmartEnterProcessorWithFixers.FixEnterProcessor() {
         private fun getControlStatementBlock(caret: Int, element: PsiElement): JetExpression? {
@@ -157,5 +157,5 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
     }
 }
 
-private val IF_BRANCHES_CONTAINERS = TokenSet.create(JetNodeTypes.THEN, JetNodeTypes.ELSE)
+private val BRANCH_CONTAINERS = TokenSet.create(JetNodeTypes.THEN, JetNodeTypes.ELSE, JetNodeTypes.BODY)
 private fun JetParameter.isInLambdaExpression() = this.getParent()?.getParent() is JetFunctionLiteral

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

@@ -30,9 +30,10 @@ public class KotlinDoWhileFixer : SmartEnterProcessorWithFixers.Fixer<KotlinSmar
         val doc = editor.getDocument()
         val stmt = psiElement as JetDoWhileExpression
         val start = stmt.range.start
+        val body = stmt.getBody()
 
         val whileKeyword = stmt.getWhileKeywordElement()
-        if (stmt.getBody() == null) {
+        if (body == null) {
             if (whileKeyword == null) {
                 doc.replaceString(start, start + "do".length(), "do {} while()")
             }
@@ -41,6 +42,12 @@ public class KotlinDoWhileFixer : SmartEnterProcessorWithFixers.Fixer<KotlinSmar
             }
             return
         }
+        else if (whileKeyword != null && body !is JetBlockExpression && body.startLine(doc) > stmt.startLine(doc)) {
+            doc.insertString(start + "do".length(), "{")
+            doc.insertString(whileKeyword.range.start - 1, "}")
+
+            return
+        }
 
         if (stmt.getCondition() == null) {
             val lParen = stmt.getLeftParenthesis()

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

@@ -360,6 +360,31 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             """
     )
 
+    fun testWhileMultiLine1() = doFunTest(
+            """
+            while (true)
+                println()<caret>
+            """,
+            """
+            while (true)
+                println()
+            <caret>
+            """
+    )
+
+    fun testWhileMultiLine2() = doFunTest(
+            """
+            while (<caret>true)
+                println()
+            """,
+            """
+            while (true) {
+                <caret>
+            }
+                println()
+            """
+    )
+
     fun testForStatement() = doFunTest(
             """
             for <caret>
@@ -457,6 +482,31 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             """
     )
 
+    fun testForMultiLine1() = doFunTest(
+            """
+            for (i in 1..10<caret>)
+                println()
+            """,
+            """
+            for (i in 1..10) {
+                <caret>
+            }
+                println()
+            """
+    )
+
+    fun testForMultiLine2() = doFunTest(
+            """
+            for (i in 1..10)
+                println()<caret>
+            """,
+            """
+            for (i in 1..10)
+                println()
+            <caret>
+            """
+    )
+
     fun testWhen() = doFunTest(
             """
             when <caret>
@@ -723,6 +773,48 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             """
     )
 
+    fun testDoWhileMultiLine1() = doFunTest(
+            """
+            do
+                println()<caret>
+            while (true)
+            """,
+            """
+            do
+                println()
+                <caret>
+            while (true)
+            """
+    )
+
+    fun testDoWhileMultiLine2() = doFunTest(
+            """
+            do<caret>
+                println()
+            while (true)
+            """,
+            """
+            do {
+                <caret>
+                println()
+            } while (true)
+            """
+    )
+
+    fun testDoWhileMultiLine3() = doFunTest(
+            """
+            do
+                println()
+            while <caret>(true)
+            """,
+            """
+            do {
+                <caret>
+                println()
+            } while (true)
+            """
+    )
+
     fun testFunBody() = doFileTest(
             """
             fun test<caret>()