Smart enter for 'for' expression

#KT-3600 In Progress

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetLoopExpression.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,9 +17,11 @@
 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.JetNodeTypes;
+import org.jetbrains.jet.lexer.JetTokens;
 
 public abstract class JetLoopExpression extends JetExpressionImpl implements JetStatementExpression {
     public JetLoopExpression(@NotNull ASTNode node) {
@@ -30,4 +32,15 @@ public abstract class JetLoopExpression extends JetExpressionImpl implements Jet
     public JetExpression getBody() {
         return findExpressionUnder(JetNodeTypes.BODY);
     }
+
+    @Nullable
+    @IfNotParsed
+    public PsiElement getLeftParenthesis() {
+        return findChildByType(JetTokens.LPAR);
+    }
+
+    @Nullable @IfNotParsed
+    public PsiElement getRightParenthesis() {
+        return findChildByType(JetTokens.RPAR);
+    }
 }

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetWhileExpression.java

@@ -17,10 +17,7 @@
 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 JetWhileExpression extends JetWhileExpressionBase {
     public JetWhileExpression(@NotNull ASTNode node) {
@@ -31,15 +28,4 @@ public class JetWhileExpression extends JetWhileExpressionBase {
     public <R, D> R accept(@NotNull JetVisitor<R, D> visitor, D data) {
         return visitor.visitWhileExpression(this, data);
     }
-
-    @Nullable
-    @IfNotParsed
-    public PsiElement getLeftParenthesis() {
-        return findChildByType(JetTokens.LPAR);
-    }
-
-    @Nullable @IfNotParsed
-    public PsiElement getRightParenthesis() {
-        return findChildByType(JetTokens.RPAR);
-    }
 }

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

@@ -32,9 +32,11 @@ 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.plugin.editor.fixers.KotlinMissingWhileBodyFixer
 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
 
@@ -45,7 +47,8 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
                 KotlinMissingIfBranchFixer(),
 
                 KotlinWhileConditionFixer(),
-                KotlinMissingWhileBodyFixer()
+                KotlinForConditionFixer(),
+                KotlinMissingForOrWhileBodyFixer()
         )
 
         addEnterProcessors(KotlinPlainEnterProcessor())
@@ -57,7 +60,9 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
         if (atCaret is PsiWhiteSpace) return null
 
         while (atCaret != null) {
-            if (atCaret is JetDeclaration || atCaret?.isJetStatement() == true) {
+            if (atCaret?.isJetStatement() == true) return atCaret
+
+            if (atCaret is JetDeclaration && (atCaret?.getParent() !is JetForExpression)) {
                 return atCaret
             }
 

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

@@ -0,0 +1,31 @@
+/*
+ * 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 org.jetbrains.jet.lang.psi.JetWhileExpression
+import com.intellij.psi.PsiElement
+import org.jetbrains.jet.lang.psi.JetForExpression
+
+public class KotlinForConditionFixer: MissingConditionFixer<JetForExpression>() {
+    override val keyword = "for"
+    override fun getElement(element: PsiElement?) = element as? JetForExpression
+    override fun getCondition(element: JetForExpression) =
+            element.getLoopRange() ?: element.getLoopParameter() ?: element.getMultiParameter()
+    override fun getLeftParenthesis(element: JetForExpression) = element.getLeftParenthesis()
+    override fun getRightParenthesis(element: JetForExpression) = element.getRightParenthesis()
+    override fun getBody(element: JetForExpression) = element.getBody()
+}

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

@@ -22,22 +22,29 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.psi.PsiElement
 import org.jetbrains.jet.lang.psi.JetWhileExpression
 import org.jetbrains.jet.lang.psi.JetBlockExpression
+import org.jetbrains.jet.lang.psi.JetLoopExpression
+import org.jetbrains.jet.lang.psi.JetForExpression
 
-public class KotlinMissingWhileBodyFixer: SmartEnterProcessorWithFixers.Fixer<KotlinSmartEnterHandler>() {
+public class KotlinMissingForOrWhileBodyFixer : SmartEnterProcessorWithFixers.Fixer<KotlinSmartEnterHandler>() {
     override fun apply(editor: Editor, processor: KotlinSmartEnterHandler, element: PsiElement) {
-        if (element !is JetWhileExpression) return
+        if (!(element is JetForExpression || element is JetWhileExpression)) return
-        val whileStatement = element as JetWhileExpression
+        val loopExpression = element as JetLoopExpression
 
         val doc = editor.getDocument()
 
-        val body = whileStatement.getBody()
+        val body = loopExpression.getBody()
         if (body is JetBlockExpression) return
-        if (body != null && body.startLine(doc) == whileStatement.startLine(doc) && whileStatement.getCondition() != null) return
 
-        val rParenth = whileStatement.getRightParenthesis()
+        if (!loopExpression.isValidLoopCondition()) return
-        if (rParenth == null) return
 
-        doc.insertString(rParenth.range.end, "{}")
+        if (body != null && body.startLine(doc) == loopExpression.startLine(doc)) return
+
+        val rParen = loopExpression.getRightParenthesis()
+        if (rParen == null) return
+
+        doc.insertString(rParen.range.end, "{}")
     }
+
+    fun JetLoopExpression.isValidLoopCondition() = getLeftParenthesis() != null && getRightParenthesis() != null
 }
 

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

@@ -360,6 +360,103 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             """
     )
 
+    fun testForStatement() = doFunTest(
+            """
+            for <caret>
+            """
+            ,
+            """
+            for (<caret>) {
+            }
+            """
+    )
+
+    fun testForStatement2() = doFunTest(
+            """
+            for<caret>
+            """
+            ,
+            """
+            for (<caret>) {
+            }
+            """
+    )
+
+    fun testForStatement4() = doFunTest(
+            """
+            for (i in 1..10<caret>) {
+            }
+            """
+            ,
+            """
+            for (i in 1..10) {
+                <caret>
+            }
+            """
+    )
+
+    fun testForStatement5() = doFunTest(
+            """
+            for (i in 1..10) {<caret>
+            """
+            ,
+            """
+            for (i in 1..10) {
+                <caret>
+            }
+            """
+    )
+
+    fun testForStatement6() = doFunTest(
+            """
+            for (i in 1..10<caret>) {
+                println()
+            }
+            """
+            ,
+            """
+            for (i in 1..10) {
+                <caret>
+                println()
+            }
+            """
+    )
+
+    fun testForStatementSingle() = doFunTest(
+            """
+            for (i in 1..10<caret>) println()
+            """
+            ,
+            """
+            for (i in 1..10) println()
+            <caret>
+            """
+    )
+
+    fun testForStatementSingleEmpty() = doFunTest(
+            """
+            for (<caret>) println()
+            """
+            ,
+            """
+            for (<caret>) println()
+            """
+    )
+
+    fun testForStatementOnLoopParameter() = doFunTest(
+            """
+            for (som<caret>e)
+            println()
+            """
+            ,
+            """
+            for (some) {
+                <caret>
+            }
+            println()
+            """
+    )
+
     fun doFunTest(before: String, after: String) {
         fun String.withFunContext(): String {
             val bodyText = "//----${this.trimIndent()}//----"