[kotlin] Migrate Add/Remove BracesIntention to ModCommand

#KTIJ-35337

GitOrigin-RevId: 850cb03915b972e739afb2893460af76cb26592d

plugins/kotlin/code-insight/api/src/org/jetbrains/kotlin/idea/codeinsight/api/applicable/intentions/KotlinPsiUpdateModCommandAction.kt

@@ -35,7 +35,8 @@ sealed class KotlinPsiUpdateModCommandAction<E : PsiElement, C : Any>(
         context: ActionContext,
         element: E,
     ): ModCommand = try {
-        ModCommand.psiUpdate(element) { e, updater ->
+        ModCommand.psiUpdate(context) { updater ->
+            val e = updater.getWritable(element)
             val elementContext = getElementContext(context, e)
                                  ?: throw NoContextException()
             invoke(context, e, elementContext, updater)

plugins/kotlin/code-insight/intentions-shared/src/org/jetbrains/kotlin/idea/codeInsight/intentions/shared/AddBracesIntention.kt

@@ -2,11 +2,16 @@
 
 package org.jetbrains.kotlin.idea.codeInsight.intentions.shared
 
-import com.intellij.openapi.editor.Editor
+import com.intellij.codeInspection.util.IntentionFamilyName
+import com.intellij.modcommand.ActionContext
+import com.intellij.modcommand.ModPsiUpdater
+import com.intellij.modcommand.Presentation
 import com.intellij.openapi.util.IntellijInternalApi
+import com.intellij.psi.PsiElement
 import org.jetbrains.annotations.ApiStatus.Internal
+import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
-import org.jetbrains.kotlin.idea.codeinsight.api.classic.intentions.SelfTargetingIntention
+import org.jetbrains.kotlin.idea.codeinsight.api.applicable.intentions.KotlinPsiUpdateModCommandAction
 import org.jetbrains.kotlin.idea.codeinsight.utils.AddBracesUtils
 import org.jetbrains.kotlin.idea.codeinsight.utils.getControlFlowElementDescription
 import org.jetbrains.kotlin.psi.*
@@ -14,30 +19,46 @@ import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 @Internal
 @IntellijInternalApi
-class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.java, KotlinBundle.lazyMessage("add.braces")) {
-    override fun isApplicableTo(element: KtElement, caretOffset: Int): Boolean {
-        val expression = element.getTargetExpression(caretOffset) ?: return false
-        if (expression is KtBlockExpression) return false
+class AddBracesIntention : KotlinPsiUpdateModCommandAction.ClassBased<KtElement, Unit>(KtElement::class) {
+    override fun getFamilyName(): @IntentionFamilyName String =
+        KotlinBundle.message("add.braces")
+
+    override fun isElementApplicable(element: KtElement, context: ActionContext): Boolean {
+        val expression = element.getTargetExpression(context.offset) ?: return false
+        // it stops further element look-up at `com.intellij.modcommand.PsiBasedModCommandAction.getElement`
+        if (expression is KtBlockExpression) return true
+
+        return when (val parent = expression.parent) {
+            is KtContainerNode -> parent.getControlFlowElementDescription() != null
+            is KtWhenEntry -> true
+            else -> false
+        }
+    }
+
+    override fun stopSearchAt(element: PsiElement, context: ActionContext): Boolean {
+        return super.stopSearchAt(element, context)
+    }
+
+    override fun getPresentation(context: ActionContext, element: KtElement): Presentation? {
+        val expression = element.getTargetExpression(context.offset) ?: return null
+        if (expression is KtBlockExpression) return null
 
         return when (val parent = expression.parent) {
             is KtContainerNode -> {
-                val description = parent.getControlFlowElementDescription() ?: return false
-                setTextGetter(KotlinBundle.lazyMessage("add.braces.to.0.statement", description))
-                true
+                val description = parent.getControlFlowElementDescription() ?: return null
+                Presentation.of(KotlinBundle.message("add.braces.to.0.statement", description))
             }
             is KtWhenEntry -> {
-                setTextGetter(KotlinBundle.lazyMessage("add.braces.to.when.entry"))
-                true
+                Presentation.of(KotlinBundle.message("add.braces.to.when.entry"))
             }
             else -> {
-                false
+                null
             }
         }
     }
 
-    override fun applyTo(element: KtElement, editor: Editor?) {
-        if (editor == null) throw IllegalArgumentException("This intention requires an editor")
-        val expression = element.getTargetExpression(editor.caretModel.offset) ?: return
+    override fun invoke(actionContext: ActionContext, element: KtElement, elementContext: Unit, updater: ModPsiUpdater) {
+        val expression = element.getTargetExpression(updater.caretOffset) ?: return
         AddBracesUtils.addBraces(element, expression)
     }
 
@@ -46,7 +67,7 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
             is KtIfExpression -> {
                 val thenExpr = then ?: return null
                 val elseExpr = `else`
-                if (elseExpr != null && caretLocation >= (elseKeyword?.startOffset ?: return null)) {
+                if (elseExpr != null && (caretLocation >= (elseKeyword?.startOffset ?: return null))) {
                     elseExpr
                 } else {
                     thenExpr
@@ -58,4 +79,6 @@ class AddBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.ja
             else -> null
         }
     }
+
+    override fun KaSession.prepareContext(element: KtElement) {}
 }

plugins/kotlin/code-insight/intentions-shared/src/org/jetbrains/kotlin/idea/codeInsight/intentions/shared/AddBracesToAllBranchesIntention.kt

@@ -2,14 +2,18 @@
 
 package org.jetbrains.kotlin.idea.codeInsight.intentions.shared
 
-import com.intellij.openapi.editor.Editor
+import com.intellij.codeInspection.util.IntentionFamilyName
+import com.intellij.modcommand.ActionContext
+import com.intellij.modcommand.ModPsiUpdater
+import com.intellij.modcommand.Presentation
 import com.intellij.psi.PsiElement
 import com.intellij.psi.util.elementType
 import com.intellij.psi.util.endOffset
 import com.intellij.util.containers.addIfNotNull
+import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.codeInsight.intentions.shared.AddBracesToAllBranchesIntention.Util.allBranchExpressions
-import org.jetbrains.kotlin.idea.codeinsight.api.classic.intentions.SelfTargetingIntention
+import org.jetbrains.kotlin.idea.codeinsight.api.applicable.intentions.KotlinPsiUpdateModCommandAction
 import org.jetbrains.kotlin.idea.codeinsight.utils.AddBracesUtils
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
@@ -19,25 +23,35 @@ import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.util.match
 
-internal class AddBracesToAllBranchesIntention : SelfTargetingIntention<KtExpression>(
-    KtExpression::class.java,
-    KotlinBundle.lazyMessage("add.braces.to.all.branches")
-) {
-    override fun isApplicableTo(element: KtExpression, caretOffset: Int): Boolean {
+internal class AddBracesToAllBranchesIntention : KotlinPsiUpdateModCommandAction.ClassBased<KtExpression, Unit>(KtExpression::class) {
+    override fun getFamilyName(): @IntentionFamilyName String =
+        KotlinBundle.message("add.braces.to.all.branches")
+
+    override fun isElementApplicable(element: KtExpression, context: ActionContext): Boolean {
         val targetIfOrWhenExpression = Util.targetIfOrWhenExpression(element) ?: return false
 
         val targetBranchExpressions = targetIfOrWhenExpression.targetBranchExpressions()
         if (targetBranchExpressions.isEmpty()) return false
-        if (caretIsOnSingleTargetBranch(targetIfOrWhenExpression, targetBranchExpressions, caretOffset)) return false
+        if (caretIsOnSingleTargetBranch(targetIfOrWhenExpression, targetBranchExpressions, context.offset)) return false
 
-        when (targetIfOrWhenExpression) {
-            is KtIfExpression -> setTextGetter(KotlinBundle.lazyMessage("add.braces.to.if.all.statements"))
-            is KtWhenExpression -> setTextGetter(KotlinBundle.lazyMessage("add.braces.to.when.all.entries"))
-        }
-        return true
+        return targetIfOrWhenExpression is KtIfExpression || targetIfOrWhenExpression is KtWhenExpression
     }
 
-    override fun applyTo(element: KtExpression, editor: Editor?) {
+    override fun getPresentation(context: ActionContext, element: KtExpression): Presentation? {
+        val targetIfOrWhenExpression = Util.targetIfOrWhenExpression(element) ?: return null
+
+        val targetBranchExpressions = targetIfOrWhenExpression.targetBranchExpressions()
+        if (targetBranchExpressions.isEmpty()) return null
+        if (caretIsOnSingleTargetBranch(targetIfOrWhenExpression, targetBranchExpressions, context.offset)) return null
+
+        return when (targetIfOrWhenExpression) {
+            is KtIfExpression -> Presentation.of(KotlinBundle.message("add.braces.to.if.all.statements"))
+            is KtWhenExpression -> Presentation.of(KotlinBundle.message("add.braces.to.when.all.entries"))
+            else -> null
+        }
+    }
+
+    override fun invoke(actionContext: ActionContext, element: KtExpression, elementContext: Unit, updater: ModPsiUpdater) {
         val targetIfOrWhenExpression = Util.targetIfOrWhenExpression(element) ?: return
         val branches = targetIfOrWhenExpression.targetBranchExpressions()
         for (branch in branches) {
@@ -105,4 +119,7 @@ internal class AddBracesToAllBranchesIntention : SelfTargetingIntention<KtExpres
             else -> emptyList()
         }
     }
+
+    override fun KaSession.prepareContext(element: KtExpression) {
+    }
 }

plugins/kotlin/code-insight/intentions-shared/src/org/jetbrains/kotlin/idea/codeInsight/intentions/shared/RemoveBracesFromAllBranchesIntention.kt

@@ -2,37 +2,51 @@
 
 package org.jetbrains.kotlin.idea.codeInsight.intentions.shared
 
-import com.intellij.openapi.editor.Editor
+import com.intellij.codeInspection.util.IntentionFamilyName
+import com.intellij.modcommand.ActionContext
+import com.intellij.modcommand.ModPsiUpdater
+import com.intellij.modcommand.Presentation
+import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.codeInsight.intentions.shared.AddBracesToAllBranchesIntention.Util.allBranchExpressions
 import org.jetbrains.kotlin.idea.codeInsight.intentions.shared.AddBracesToAllBranchesIntention.Util.targetIfOrWhenExpression
-import org.jetbrains.kotlin.idea.codeinsight.api.classic.intentions.SelfTargetingIntention
+import org.jetbrains.kotlin.idea.codeinsight.api.applicable.intentions.KotlinPsiUpdateModCommandAction
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
-internal class RemoveBracesFromAllBranchesIntention : SelfTargetingIntention<KtExpression>(
-    KtExpression::class.java,
-    KotlinBundle.lazyMessage("remove.braces.from.all.branches")
-) {
-    override fun isApplicableTo(element: KtExpression, caretOffset: Int): Boolean {
+internal class RemoveBracesFromAllBranchesIntention  : KotlinPsiUpdateModCommandAction.ClassBased<KtExpression, Unit>(KtExpression::class) {
+    override fun getFamilyName(): @IntentionFamilyName String =
+        KotlinBundle.message("remove.braces.from.all.branches")
+
+    override fun isElementApplicable(element: KtExpression, context: ActionContext): Boolean {
         val targetIfOrWhenExpression = targetIfOrWhenExpression(element) ?: return false
 
         val targetBranchExpressions = targetIfOrWhenExpression.targetBranchExpressions()
         if (targetBranchExpressions.isEmpty() || targetBranchExpressions.any { !RemoveBracesIntention.Holder.isApplicableTo(it) }) return false
-        if (caretIsOnSingleTargetBranch(targetIfOrWhenExpression, targetBranchExpressions, caretOffset)) return false
+        if (caretIsOnSingleTargetBranch(targetIfOrWhenExpression, targetBranchExpressions, context.offset)) return false
 
-        when (targetIfOrWhenExpression) {
-            is KtIfExpression -> setTextGetter(KotlinBundle.lazyMessage("remove.braces.from.if.all.statements"))
-            is KtWhenExpression -> setTextGetter(KotlinBundle.lazyMessage("remove.braces.from.when.all.entries"))
-        }
-        return true
+        return targetIfOrWhenExpression is KtIfExpression || targetIfOrWhenExpression is KtWhenExpression
     }
 
-    override fun applyTo(element: KtExpression, editor: Editor?) {
+    override fun getPresentation(context: ActionContext, element: KtExpression): Presentation? {
+        val targetIfOrWhenExpression = targetIfOrWhenExpression(element) ?: return null
+
+        val targetBranchExpressions = targetIfOrWhenExpression.targetBranchExpressions()
+        if (targetBranchExpressions.isEmpty() || targetBranchExpressions.any { !RemoveBracesIntention.Holder.isApplicableTo(it) }) return null
+        if (caretIsOnSingleTargetBranch(targetIfOrWhenExpression, targetBranchExpressions, context.offset)) return null
+
+        return when (targetIfOrWhenExpression) {
+            is KtIfExpression -> Presentation.of(KotlinBundle.message("remove.braces.from.if.all.statements"))
+            is KtWhenExpression -> Presentation.of(KotlinBundle.message("remove.braces.from.when.all.entries"))
+            else -> null
+        }
+    }
+
+    override fun invoke(actionContext: ActionContext, element: KtExpression, elementContext: Unit, updater: ModPsiUpdater) {
         val targetIfOrWhenExpression = targetIfOrWhenExpression(element) ?: return
         targetIfOrWhenExpression.targetBranchExpressions().forEach {
-            RemoveBracesIntention.Holder.removeBraces(targetIfOrWhenExpression, it)
+            RemoveBracesIntention.Holder.removeBraces(actionContext, targetIfOrWhenExpression, it, updater)
         }
     }
 
@@ -55,4 +69,8 @@ internal class RemoveBracesFromAllBranchesIntention : SelfTargetingIntention<KtE
         }?.startOffset ?: return false
         return caretOffset in startOffset..singleBranchExpression.endOffset
     }
+
+
+    override fun KaSession.prepareContext(element: KtExpression) {
+    }
 }

plugins/kotlin/code-insight/intentions-shared/src/org/jetbrains/kotlin/idea/codeInsight/intentions/shared/RemoveBracesIntention.kt

@@ -2,45 +2,62 @@
 
 package org.jetbrains.kotlin.idea.codeInsight.intentions.shared
 
-import com.intellij.openapi.editor.Editor
+import com.intellij.application.options.CodeStyle
+import com.intellij.codeInspection.util.IntentionFamilyName
+import com.intellij.modcommand.ActionContext
+import com.intellij.modcommand.ModPsiUpdater
+import com.intellij.modcommand.Presentation
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
-import com.intellij.psi.util.PsiTreeUtil
-import com.intellij.psi.util.elementType
 import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
-import org.jetbrains.kotlin.idea.codeinsight.api.classic.intentions.SelfTargetingIntention
+import org.jetbrains.kotlin.idea.codeinsight.api.applicable.intentions.KotlinApplicableModCommandAction
 import org.jetbrains.kotlin.idea.codeinsight.utils.getControlFlowElementDescription
-import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
-internal class RemoveBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.java, KotlinBundle.lazyMessage("remove.braces")) {
-    override fun isApplicableTo(element: KtElement, caretOffset: Int): Boolean {
+internal class RemoveBracesIntention: KotlinApplicableModCommandAction<KtElement, Unit>(KtElement::class) {
+    override fun getFamilyName(): @IntentionFamilyName String =
+        KotlinBundle.message("remove.braces")
+
+    override fun isApplicableByPsi(element: KtElement): Boolean {
         val block = element.findChildBlock() ?: return false
         if (!Holder.isApplicableTo(block)) return false
-        when (val container = block.parent) {
+        return when(block.parent) {
+            is KtContainerNode, is KtWhenEntry -> true
+            else -> false
+        }
+    }
+
+    override fun getPresentation(context: ActionContext, element: KtElement): Presentation? {
+        val block = element.findChildBlock() ?: return null
+        if (!Holder.isApplicableTo(block)) return null
+
+        return when (val container = block.parent) {
             is KtContainerNode -> {
-                val description = container.getControlFlowElementDescription() ?: return false
-                setTextGetter(KotlinBundle.lazyMessage("remove.braces.from.0.statement", description))
+                val description = container.getControlFlowElementDescription() ?: return null
+                Presentation.of(KotlinBundle.message("remove.braces.from.0.statement", description))
             }
             is KtWhenEntry -> {
-                setTextGetter(KotlinBundle.lazyMessage("remove.braces.from.when.entry"))
+                Presentation.of(KotlinBundle.message("remove.braces.from.when.entry"))
             }
+
+            else -> null
         }
-        return true
     }
 
-    override fun applyTo(element: KtElement, editor: Editor?) {
+    override fun invoke(actionContext: ActionContext, element: KtElement, elementContext: Unit, updater: ModPsiUpdater) {
         val block = element.findChildBlock() ?: return
-        Holder.removeBraces(element, block, editor)
+        Holder.removeBraces(actionContext, element, block, updater)
     }
 
-    override fun skipProcessingFurtherElementsAfter(element: PsiElement): Boolean =
-        element is KtBlockExpression && element.parent !is KtWhenEntry
+    override fun stopSearchAt(element: PsiElement, context: ActionContext): Boolean {
+        return element is KtBlockExpression && element.parent !is KtWhenEntry
+    }
 
     private fun KtElement.findChildBlock(): KtBlockExpression? = when (this) {
         is KtBlockExpression -> this
@@ -69,25 +86,28 @@ internal class RemoveBracesIntention : SelfTargetingIntention<KtElement>(KtEleme
             }
         }
 
-        fun removeBraces(element: KtElement, block: KtBlockExpression, editor: Editor? = null) {
-            val factory = KtPsiFactory(element.project)
+        fun removeBraces(actionContext: ActionContext, element: KtElement, block: KtBlockExpression, updater: ModPsiUpdater) {
+            val project = element.project
+            val factory = KtPsiFactory(project)
             val statement = block.statements.single()
-            val caretOnAfterStatement = if (editor != null) editor.caretModel.offset >= statement.endOffset else false
+            val caretOnAfterStatement = updater.caretOffset >= statement.endOffset
 
             val container = block.parent
             val construct = container.parent as KtExpression
-            statement.handleComments(block, factory)
+            statement.handleComments(block)
 
             val newElement = block.replace(statement.copy())
-            editor?.caretModel?.moveToOffset(if (caretOnAfterStatement) newElement.endOffset else newElement.startOffset)
+            updater.moveCaretTo(if (caretOnAfterStatement) newElement.endOffset else newElement.startOffset)
 
             if (construct is KtDoWhileExpression) {
                 newElement.parent!!.addAfter(factory.createNewLine(), newElement)
-            } else if (editor != null) {
-                val document = editor.document
+            } else {
+                val document = actionContext.file.fileDocument
+                val rightMargin = CodeStyle.getSettings(project).getRightMargin(element.language)
                 val line = document.getLineNumber(newElement.startOffset)
-                val rightMargin = editor.settings.getRightMargin(editor.project)
-                if (document.getLineEndOffset(line) - document.getLineStartOffset(line) >= rightMargin) {
+                val lineStartOffset = document.getLineStartOffset(line)
+                val lineEndOffset = document.getLineEndOffset(line) + newElement.textLength
+                if (lineEndOffset - lineStartOffset >= rightMargin) {
                     newElement.parent.addBefore(factory.createNewLine(), newElement)
                 }
             }
@@ -99,12 +119,12 @@ internal class RemoveBracesIntention : SelfTargetingIntention<KtElement>(KtEleme
             ) {
                 val replaced = construct.replace(factory.createExpressionByPattern("($0)", construct))
                 (replaced.children[0] as? KtIfExpression)?.`else`?.let {
-                    editor?.caretModel?.moveToOffset(if (caretOnAfterStatement) it.endOffset else it.startOffset)
+                    updater.moveCaretTo(if (caretOnAfterStatement) it.endOffset else it.startOffset)
                 }
             }
         }
 
-        private fun KtExpression.handleComments(block: KtBlockExpression, factory: KtPsiFactory) {
+        private fun KtExpression.handleComments(block: KtBlockExpression) {
             val nextComments = comments(forward = true)
             val prevComments = comments(forward = false).reversed()
             val blockParent = block.parent
@@ -123,7 +143,7 @@ internal class RemoveBracesIntention : SelfTargetingIntention<KtElement>(KtEleme
             return if (elements.any { it is PsiComment }) elements else emptyList()
         }
 
-        private fun List<PsiElement>.hasLineBreak(): Boolean =
-            any { it is PsiWhiteSpace && it.textContains('\n') }
     }
+
+    override fun KaSession.prepareContext(element: KtElement) {}
 }

plugins/kotlin/code-insight/intentions-shared/tests/testData/intentions/addBraces/elseWithComment.kt

@@ -1,6 +1,7 @@
-fun foo() {}
+fun foo1() {}
+fun foo2() {}
 
 fun test(b: Boolean) {
-    if (b) foo() <caret>else
-        /* aaa */ foo() // bbb
+    if (b) foo1() <caret>else
+        /* aaa */ foo2() // bbb
 }

plugins/kotlin/code-insight/intentions-shared/tests/testData/intentions/addBraces/elseWithComment.kt.after

@@ -1,7 +1,8 @@
-fun foo() {}
+fun foo1() {}
+fun foo2() {}
 
 fun test(b: Boolean) {
-    if (b) foo() else {
-        /* aaa */ foo() // bbb 
+    if (b) foo1() else {
+        /* aaa */ foo2() // bbb 
     }
 }

plugins/kotlin/code-insight/utils/src/org/jetbrains/kotlin/idea/codeinsight/utils/AddBracesUtils.kt

@@ -6,20 +6,8 @@ import com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.idea.base.psi.getLineNumber
 import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.psi.KtDoWhileExpression
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtForExpression
-import org.jetbrains.kotlin.psi.KtIfExpression
-import org.jetbrains.kotlin.psi.KtPsiFactory
-import org.jetbrains.kotlin.psi.KtWhenEntry
-import org.jetbrains.kotlin.psi.KtWhileExpression
-import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
-import org.jetbrains.kotlin.psi.psiUtil.allChildren
-import org.jetbrains.kotlin.psi.psiUtil.getNextSiblingIgnoringWhitespace
-import org.jetbrains.kotlin.psi.psiUtil.getNextSiblingIgnoringWhitespaceAndComments
-import org.jetbrains.kotlin.psi.psiUtil.getPrevSiblingIgnoringWhitespaceAndComments
-import org.jetbrains.kotlin.psi.psiUtil.siblings
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.*
 
 object AddBracesUtils {
     fun addBraces(element: KtElement, expression: KtExpression) {
@@ -72,15 +60,18 @@ object AddBracesUtils {
         }
 
         val result = expression.replace(psiFactory.createSingleStatementBlock(expression, nextComment = nextComment?.text))
-        when (element) {
-            is KtDoWhileExpression -> {
-                // remove new line between '}' and while
-                (element.body?.parent?.nextSibling as? PsiWhiteSpace)?.delete()
-            }
 
-            is KtIfExpression -> {
-                (result?.parent?.nextSibling as? PsiWhiteSpace)?.delete()
-            }
+        val parent = when (element) {
+            is KtDoWhileExpression -> element.body?.parent
+            is KtIfExpression -> result.parent
+            is KtWhenEntry -> result
+            else -> null
+        }
+        parent?.let {
+            // remove new line between '{' and if/when
+            (it.prevSibling as? PsiWhiteSpace)?.delete()
+            // remove new line between '}' and while
+            (it.nextSibling as? PsiWhiteSpace)?.delete()
         }
         saver?.restore(result, forceAdjustIndent = false)
     }

plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/intentions/AbstractIntentionTestBase.kt

@@ -215,12 +215,13 @@ abstract class AbstractIntentionTestBase : KotlinLightCodeInsightFixtureTestCase
         val isApplicableOnPooled: Boolean = project.computeOnBackground {
             runReadAction { intentionAction.isAvailable(project, editor, file) }
         }
+
+        val modCommandAction: ModCommandAction? = intentionAction.asModCommandAction()
         Assert.assertTrue(
-            "isAvailable() for " + intentionAction.javaClass + " should return " + isApplicableExpected,
+            "isAvailable() for " + (modCommandAction ?: intentionAction).javaClass + " should return " + isApplicableExpected,
             isApplicableExpected == isApplicableOnPooled
         )
 
-        val modCommandAction: ModCommandAction? = intentionAction.asModCommandAction()
         if (modCommandAction == null) {
             val isApplicableOnEdt = intentionAction.isAvailable(project, editor, file)