Smaller availability range for RemoveBracesIntention + refactored it and AddBracesIntention

idea/idea-analysis/src/org/jetbrains/kotlin/idea/JetBundle.properties

@@ -315,10 +315,6 @@ move.lambda.outside.parentheses=Move lambda expression out of parentheses
 move.lambda.outside.parentheses.family=Move Lambda Expression out of Parentheses
 replace.it.with.explicit.function.literal.param=Replace 'it' with explicit parameter
 replace.it.with.explicit.function.literal.param.family=Replace 'it' with Explicit Parameter
-remove.braces=Remove braces
-remove.braces.family=Remove Braces
-add.braces=Add braces
-add.braces.family=Add Braces
 convert.negated.boolean.sequence=Replace negated sequence with DeMorgan equivalent
 convert.negated.boolean.sequence.family=Replace Negated Sequence with DeMorgan Equivalent
 convert.negated.expression.with.demorgans.law=DeMorgan Law

idea/src/org/jetbrains/kotlin/idea/intentions/AddBracesIntention.kt

@@ -16,49 +16,53 @@
 
 package org.jetbrains.kotlin.idea.intentions
 
-import org.jetbrains.kotlin.psi.JetExpressionImpl
-import org.jetbrains.kotlin.psi.JetPsiFactory
 import com.intellij.openapi.editor.Editor
-import com.intellij.lang.ASTNode
-import org.jetbrains.kotlin.JetNodeTypes
 import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.psi.*
 
-public class AddBracesIntention : JetSelfTargetingIntention<JetExpressionImpl>("add.braces", javaClass()) {
-    override fun isApplicableTo(element: JetExpressionImpl, caretOffset: Int): Boolean {
-        val expressionKind = element.getExpressionKind(caretOffset) ?: return false
+public class AddBracesIntention : JetSelfTargetingIntention<JetExpression>(javaClass(), "Add braces") {
+    override fun isApplicableTo(element: JetExpression, caretOffset: Int): Boolean {
+        val expression = element.getTargetExpression(caretOffset) ?: return false
+        if (expression is JetBlockExpression) return false
 
-        if (element.findBlockInExpression(expressionKind) != null) return false
-
-        setText("Add braces to '${expressionKind.text}' statement")
+        val description = (expression.getParent() as JetContainerNode).description()
+        setText("Add braces to '$description' statement")
         return true
     }
 
-    override fun applyTo(element: JetExpressionImpl, editor: Editor) {
-        val expressionKind = element.getExpressionKind(editor.getCaretModel().getOffset())!!
-        val bodyNode = when (expressionKind) {
-            ExpressionKind.ELSE -> element.getNode().findChildByType(JetNodeTypes.ELSE)
-            ExpressionKind.IF -> element.getNode().findChildByType(JetNodeTypes.THEN)
-            else -> element.getNode().findChildByType(JetNodeTypes.BODY)
-        }
-        generateCleanOutput(element, bodyNode, expressionKind)
-    }
+    override fun applyTo(element: JetExpression, editor: Editor) {
+        val expression = element.getTargetExpression(editor.getCaretModel().getOffset())!!
 
-    fun generateCleanOutput(element: JetExpressionImpl, bodyNode: ASTNode?, expressionKind: ExpressionKind) {
         if (element.getNextSibling()?.getText() == ";") {
             element.getNextSibling()!!.delete()
         }
-        val psiFactory = JetPsiFactory(element)
-        val newElement = bodyNode!!.getPsi()!!.replace(psiFactory.createFunctionBody(bodyNode.getText()))
 
-        //handles the case of the block statement being on a new line
-        if (newElement.getPrevSibling() is PsiWhiteSpace) {
-            newElement.getPrevSibling()!!.replace(psiFactory.createWhiteSpace())
-        } else {
-            //handles the case of no space between condition and statement
-            newElement.addBefore(psiFactory.createWhiteSpace(), newElement.getFirstChild())
+        val psiFactory = JetPsiFactory(element)
+        expression.replace(psiFactory.createFunctionBody(expression.getText()))
+
+        if (element is JetDoWhileExpression) { // remove new line between '}' and while
+            (element.getBody()!!.getParent().getNextSibling() as? PsiWhiteSpace)?.delete()
         }
-        if (expressionKind == ExpressionKind.DOWHILE) {
-            newElement.getNextSibling()?.delete()
+    }
+
+    private fun JetExpression.getTargetExpression(caretLocation: Int): JetExpression? {
+        when (this) {
+            is JetIfExpression -> {
+                val thenExpr = getThen() ?: return null
+                val elseExpr = getElse()
+                if (elseExpr != null && caretLocation >= getElseKeyword()!!.getTextRange().getStartOffset()) {
+                    return elseExpr
+                }
+                return thenExpr
+            }
+
+            is JetWhileExpression -> return getBody()
+
+            is JetDoWhileExpression -> return getBody()
+
+            is JetForExpression -> return getBody()
+
+            else -> return null
         }
     }
 }

idea/src/org/jetbrains/kotlin/idea/intentions/RemoveBracesIntention.kt

@@ -16,53 +16,52 @@
 
 package org.jetbrains.kotlin.idea.intentions
 
-import org.jetbrains.kotlin.psi.JetBlockExpression
-import org.jetbrains.kotlin.psi.JetPsiFactory
-import org.jetbrains.kotlin.psi.JetExpressionImpl
 import com.intellij.openapi.editor.Editor
-import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 
-public class RemoveBracesIntention : JetSelfTargetingIntention<JetExpressionImpl>("remove.braces", javaClass()) {
-    override fun isApplicableTo(element: JetExpressionImpl, caretOffset: Int): Boolean {
-        val expressionKind = element.getExpressionKind(caretOffset) ?: return false
+public class RemoveBracesIntention : JetSelfTargetingIntention<JetBlockExpression>(javaClass(), "Remove braces") {
+    override fun isApplicableTo(element: JetBlockExpression, caretOffset: Int): Boolean {
+        if (element.getStatements().size() != 1) return false
 
-        val jetBlockElement = element.findBlockInExpression(expressionKind) ?: return false
+        val containerNode = element.getParent() as? JetContainerNode ?: return false
 
-        if (jetBlockElement.getStatements().size == 1) {
-            setText("Remove braces from '${expressionKind.text}' statement")
-            return true
-        }
-        return false
+        val lBrace = element.getLBrace() ?: return false
+        val rBrace = element.getRBrace() ?: return false
+        if (!lBrace.getTextRange().containsOffset(caretOffset) && !rBrace.getTextRange().containsOffset(caretOffset)) return false
+
+        setText("Remove braces from '${containerNode.description()}' statement")
+        return true
     }
 
-    override fun applyTo(element: JetExpressionImpl, editor: Editor) {
-        val expressionKind = element.getExpressionKind(editor.getCaretModel().getOffset())!!
+    override fun applyTo(element: JetBlockExpression, editor: Editor) {
+        val statement = element.getStatements().single()
 
-        val jetBlockElement = element.findBlockInExpression(expressionKind)
-        val firstStatement = jetBlockElement!!.getStatements().first()
+        val containerNode = element.getParent() as JetContainerNode
+        val construct = containerNode.getParent() as JetExpression
+        handleComments(construct, element)
 
-        handleComments(element, jetBlockElement)
+        val newElement = element.replace(statement.copy())
 
-        val newElement = jetBlockElement.replace(firstStatement.copy())
-
-        if (expressionKind == ExpressionKind.DOWHILE) {
+        if (construct is JetDoWhileExpression) {
             newElement.getParent()!!.addAfter(JetPsiFactory(element).createNewLine(), newElement)
         }
     }
 
-    fun handleComments(element: JetExpressionImpl, blockElement: JetBlockExpression) {
-        var sibling = blockElement.getFirstChild()?.getNextSibling()
+    private fun handleComments(construct: JetExpression, block: JetBlockExpression) {
+        var sibling = block.getFirstChild()?.getNextSibling()
 
         while (sibling != null) {
             if (sibling is PsiComment) {
                 //cleans up extra whitespace
-                val psiFactory = JetPsiFactory(element)
-                if (element.getPrevSibling() is PsiWhiteSpace) {
-                    element.getPrevSibling()!!.replace(psiFactory.createNewLine())
+                val psiFactory = JetPsiFactory(construct)
+                if (construct.getPrevSibling() is PsiWhiteSpace) {
+                    construct.getPrevSibling()!!.replace(psiFactory.createNewLine())
                 }
-                val commentElement = element.getParent()!!.addBefore(sibling as PsiComment, element.getPrevSibling())
-                element.getParent()!!.addBefore(psiFactory.createNewLine(), commentElement)
+                val commentElement = construct.getParent()!!.addBefore(sibling as PsiComment, construct.getPrevSibling())
+                construct.getParent()!!.addBefore(psiFactory.createNewLine(), commentElement)
             }
             sibling = sibling!!.getNextSibling()
         }

idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt

@@ -16,15 +16,17 @@
 
 package org.jetbrains.kotlin.idea.intentions
 
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.types.JetType
-import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.idea.util.ShortenReferences
 import org.jetbrains.kotlin.JetNodeTypes
-import org.jetbrains.kotlin.lexer.JetTokens
-import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
+import org.jetbrains.kotlin.idea.util.ShortenReferences
+import org.jetbrains.kotlin.lexer.JetTokens
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.parents
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.types.JetType
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
 fun specifyTypeExplicitly(declaration: JetNamedFunction, typeText: String) {
     specifyTypeExplicitly(declaration, JetPsiFactory(declaration).createType(typeText))
@@ -55,43 +57,17 @@ fun functionReturnType(function: JetNamedFunction): JetType? {
     return (descriptor as FunctionDescriptor).getReturnType()
 }
 
-enum class ExpressionKind(val text: String) {
-    IF : ExpressionKind("if")
-    ELSE : ExpressionKind("else")
-    WHILE : ExpressionKind("while")
-    DOWHILE : ExpressionKind("do...while")
-    FOR : ExpressionKind("for")
-}
-
-fun JetExpressionImpl.findBlockInExpression(expressionKind: ExpressionKind?): JetBlockExpression? {
-    val bodyNode = when (expressionKind) {
-        ExpressionKind.IF -> this.getNode().findChildByType(JetNodeTypes.THEN)
-        ExpressionKind.ELSE -> this.getNode().findChildByType(JetNodeTypes.ELSE)
-        else -> this.getNode().findChildByType(JetNodeTypes.BODY)
-    }
-    return bodyNode!!.getPsi()!!.getFirstChild() as? JetBlockExpression
-}
-
-fun JetExpressionImpl.getExpressionKind(caretLocation: Int): ExpressionKind? {
-    when (this) {
-        is JetIfExpression -> {
-            if (this.getElse() != null) {
-                val elseLocation = this.getNode().findChildByType(JetTokens.ELSE_KEYWORD)!!.getPsi()!!.getTextOffset()
-                if (caretLocation >= elseLocation) return ExpressionKind.ELSE
+fun JetContainerNode.description(): String? {
+    when (getNode().getElementType()) {
+        JetNodeTypes.THEN -> return "if"
+        JetNodeTypes.ELSE -> return "else"
+        JetNodeTypes.BODY -> {
+            when (getParent()) {
+                is JetWhileExpression -> return "while"
+                is JetDoWhileExpression -> return "do...while"
+                is JetForExpression -> return "for"
             }
-            return ExpressionKind.IF
-        }
-        is JetWhileExpression -> {
-            return ExpressionKind.WHILE
-        }
-        is JetDoWhileExpression -> {
-            return ExpressionKind.DOWHILE
-        }
-        is JetForExpression -> {
-            return ExpressionKind.FOR
-        }
-        else -> {
-            return null
         }
     }
+    return null
 }

idea/testData/intentions/removeBraces/else.kt

@@ -3,7 +3,7 @@ fun doSomething<T>(a: T) {}
 fun foo() {
     if (true) {
         doSomething("test")
-    } <caret>else {
+    } else <caret>{
         doSomething("test2")
     }
 }

idea/testData/intentions/removeBraces/for.kt

@@ -1,7 +1,7 @@
 fun doSomething<T>(a: T) {}
 
 fun foo() {
-    for<caret> (i in 1..4) {
+    for (i in 1..4) {<caret>
         doSomething("test")
     }
 }

idea/testData/intentions/removeBraces/forWithLocalVariable.kt

@@ -1,5 +1,5 @@
 fun foo(a: List<Int>) {
-    for<caret> (x in a) {
+    for (x in a) <caret>{
         val y = x
     }
 }

idea/testData/intentions/removeBraces/forWithLocalVariable.kt.after

@@ -1,3 +1,3 @@
 fun foo(a: List<Int>) {
-    for<caret> (x in a) val y = x
+    for (x in a) <caret>val y = x
 }

idea/testData/intentions/removeBraces/if.kt

@@ -1,7 +1,7 @@
 fun doSomething<T>(a: T) {}
 
 fun foo() {
-    <caret>if (true) {
+    if (true) {
         doSomething("test")
-    }
+    }<caret>
 }

idea/testData/intentions/removeBraces/ifWithComment.kt

@@ -1,7 +1,7 @@
 fun doSomething<T>(a: T) {}
 
 fun foo() {
-    <caret>if (true) {
+    if (true) <caret>{
         //comment
         doSomething("test")
     }

idea/testData/intentions/removeBraces/ifWithNoStatement.kt

@@ -1,4 +1,4 @@
 // IS_APPLICABLE: false
 fun foo() {
-    if (true) { }
+    if (true) <caret>{ }
 }

idea/testData/intentions/removeBraces/ifWithSemicolon.kt

@@ -1,7 +1,7 @@
 fun doSomething<T>(a: T) {}
 
 fun foo() {
-    <caret>if (true) {
+    if (true) {
         doSomething("test");
-    }
+    <caret>}
 }

idea/testData/intentions/removeBraces/ifWithTwoStatements.kt

@@ -5,5 +5,5 @@ fun foo() {
     if (true) {
         doSomething("test")
         doSomething("test2")
-        <caret>}
+    <caret>}
 }

idea/testData/intentions/removeBraces/whileWithTwoStatements.kt

@@ -5,5 +5,5 @@ fun foo() {
     while (true) {
         doSomething("test")
         doSomething("test2")
-        <caret>}
+    <caret>}
 }