From fd4aeb75fc07231ffd7b352b0c9b92b9161f9612 Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <Nikolay.Krasko@jetbrains.com>
Date: Tue, 24 Jun 2014 20:21:41 +0400
Subject: [PATCH] Smart enter for one line statements in for, while and
 do-while

 #KT-3600 In Progress
---
 .../plugin/editor/KotlinSmartEnterHandler.kt  |  4 +-
 .../editor/fixers/KotlinDoWhileFixer.kt       |  9 +-
 .../codeInsight/smartEnter/SmartEnterTest.kt  | 92 +++++++++++++++++++
 3 files changed, 102 insertions(+), 3 deletions(-)

diff --git a/idea/src/org/jetbrains/jet/plugin/editor/KotlinSmartEnterHandler.kt b/idea/src/org/jetbrains/jet/plugin/editor/KotlinSmartEnterHandler.kt
index f2ca2c0eaa6..91a10eba4b0 100644
--- a/idea/src/org/jetbrains/jet/plugin/editor/KotlinSmartEnterHandler.kt
+++ b/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
diff --git a/idea/src/org/jetbrains/jet/plugin/editor/fixers/KotlinDoWhileFixer.kt b/idea/src/org/jetbrains/jet/plugin/editor/fixers/KotlinDoWhileFixer.kt
index bba179e8bfe..88bb92c6454 100644
--- a/idea/src/org/jetbrains/jet/plugin/editor/fixers/KotlinDoWhileFixer.kt
+++ b/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()
diff --git a/idea/tests/org/jetbrains/jet/plugin/codeInsight/smartEnter/SmartEnterTest.kt b/idea/tests/org/jetbrains/jet/plugin/codeInsight/smartEnter/SmartEnterTest.kt
index 1b0b73feff2..72e0c45a850 100644
--- a/idea/tests/org/jetbrains/jet/plugin/codeInsight/smartEnter/SmartEnterTest.kt
+++ b/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>()