Minor: remove braces is now applicable for the whole loop or when entry

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

@@ -18,21 +18,27 @@ package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
 import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.psi.*
 
-class RemoveBracesIntention : SelfTargetingIntention<KtBlockExpression>(KtBlockExpression::class.java, "Remove braces") {
+class RemoveBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.java, "Remove braces") {
-    override fun isApplicableTo(element: KtBlockExpression, caretOffset: Int): Boolean {
+
-        val singleStatement = element.statements.singleOrNull() ?: return false
+    private fun KtElement.findChildBlock() = when (this) {
-        val container = element.parent
+        is KtBlockExpression -> this
+        is KtLoopExpression -> body as? KtBlockExpression
+        is KtWhenEntry -> expression as? KtBlockExpression
+        else -> null
+    }
+
+    override fun isApplicableTo(element: KtElement, caretOffset: Int): Boolean {
+        val block = element.findChildBlock() ?: return false
+        val singleStatement = block.statements.singleOrNull() ?: return false
+        val container = block.parent
         when (container) {
             is KtContainerNode -> {
                 if (singleStatement is KtIfExpression && container.parent is KtIfExpression) return false
 
-                val lBrace = element.lBrace ?: return false
-                val rBrace = element.rBrace ?: return false
-                if (!lBrace.textRange.containsOffset(caretOffset) && !rBrace.textRange.containsOffset(caretOffset)) return false
-
                 val description = container.description() ?: return false
                 text = "Remove braces from '$description' statement"
                 return true
@@ -45,17 +51,18 @@ class RemoveBracesIntention : SelfTargetingIntention<KtBlockExpression>(KtBlockE
         }
     }
 
-    override fun applyTo(element: KtBlockExpression, editor: Editor?) {
+    override fun applyTo(element: KtElement, editor: Editor?) {
-        val statement = element.statements.single()
+        val block = element.findChildBlock() ?: return
+        val statement = block.statements.single()
 
-        val container = element.parent!!
+        val container = block.parent!!
         val construct = container.parent as KtExpression
-        handleComments(construct, element)
+        handleComments(construct, block)
 
-        val newElement = element.replace(statement.copy())
+        val newElement = block.replace(statement.copy())
 
         if (construct is KtDoWhileExpression) {
-            newElement.parent!!.addAfter(KtPsiFactory(element).createNewLine(), newElement)
+            newElement.parent!!.addAfter(KtPsiFactory(block).createNewLine(), newElement)
         }
     }
 
@@ -75,4 +82,6 @@ class RemoveBracesIntention : SelfTargetingIntention<KtBlockExpression>(KtBlockE
             sibling = sibling.nextSibling
         }
     }
+
+    override fun allowCaretInsideElement(element: PsiElement) = element !is KtBlockExpression || element.parent is KtWhenEntry
 }

idea/testData/intentions/removeBraces/whenSimpleOutsideBlock.kt

@@ -0,0 +1,7 @@
+fun foo() {
+    when (1) {
+        else<caret> -> {
+            foo()
+        }
+    }
+}

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

@@ -0,0 +1,5 @@
+fun foo() {
+    when (1) {
+        else -> foo()
+    }
+}

idea/testData/intentions/removeBraces/whileOutsideBlock.kt

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

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

@@ -0,0 +1,5 @@
+fun <T> doSomething(a: T) {}
+
+fun foo() {
+    while (true) doSomething("test")
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -9665,6 +9665,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("whenSimpleOutsideBlock.kt")
+        public void testWhenSimpleOutsideBlock() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/removeBraces/whenSimpleOutsideBlock.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("whenStatement.kt")
         public void testWhenStatement() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/removeBraces/whenStatement.kt");
@@ -9677,6 +9683,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("whileOutsideBlock.kt")
+        public void testWhileOutsideBlock() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/removeBraces/whileOutsideBlock.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("whileWithTwoStatements.kt")
         public void testWhileWithTwoStatements() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/removeBraces/whileWithTwoStatements.kt");