Smart enter for do-while fixed

- Do not brake code if some expression parsed as body for bare "do"
 - Fix isStatement() check

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

@@ -36,6 +36,7 @@ import org.jetbrains.jet.lang.psi.JetFunctionLiteral
 import com.intellij.psi.tree.TokenSet
 import org.jetbrains.jet.JetNodeTypes
 import org.jetbrains.jet.lang.psi.JetLoopExpression
+import org.jetbrains.jet.lexer.JetTokens
 
 public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
     {
@@ -120,8 +121,11 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
         }
     }
 
-    private fun PsiElement.isJetStatement() =
-        getParent() is JetBlockExpression || (getParent()?.getNode()?.getElementType() in BRANCH_CONTAINERS)
+    private fun PsiElement.isJetStatement() = when {
+        getParent() is JetBlockExpression && getNode()?.getElementType() !in BRACES -> true
+        getParent()?.getNode()?.getElementType() in BRANCH_CONTAINERS && this !is JetBlockExpression -> true
+        else -> false
+    }
 
     class KotlinPlainEnterProcessor : SmartEnterProcessorWithFixers.FixEnterProcessor() {
         private fun getControlStatementBlock(caret: Int, element: PsiElement): JetExpression? {
@@ -158,4 +162,5 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
 }
 
 private val BRANCH_CONTAINERS = TokenSet.create(JetNodeTypes.THEN, JetNodeTypes.ELSE, JetNodeTypes.BODY)
+private val BRACES = TokenSet.create(JetTokens.RBRACE, JetTokens.LBRACE)
 private fun JetParameter.isInLambdaExpression() = this.getParent()?.getParent() is JetFunctionLiteral

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

@@ -43,8 +43,8 @@ public class KotlinDoWhileFixer : SmartEnterProcessorWithFixers.Fixer<KotlinSmar
             return
         }
         else if (whileKeyword != null && body !is JetBlockExpression && body.startLine(doc) > stmt.startLine(doc)) {
+            doc.insertString(whileKeyword.range.start, "}")
             doc.insertString(start + "do".length(), "{")
-            doc.insertString(whileKeyword.range.start - 1, "}")
 
             return
         }

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

@@ -677,8 +677,8 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             ,
             """
             do {
+                <caret>
             } while (true)
-            <caret>
             """
     )
 
@@ -747,6 +747,32 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             """
     )
 
+    fun testDoWhile13() = doFunTest(
+            """
+            do<caret>
+            println("some")
+            """
+            ,
+            """
+            do {
+                println("some")
+            } while (<caret>)
+            """
+    )
+
+    fun testDoWhile14() = doFunTest(
+            """
+            do <caret>
+            println("some")
+            """
+            ,
+            """
+            do {
+                println("some")
+            } while (<caret>)
+            """
+    )
+
     fun testDoWhileOneLine1() = doFunTest(
             """
             do println("some") while (true<caret>)