Remove braces intention: fix caret position

#KT-31443 Fixed

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

@@ -22,6 +22,8 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 class RemoveBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class.java, "Remove braces") {
 
@@ -57,13 +59,15 @@ class RemoveBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class
     override fun applyTo(element: KtElement, editor: Editor?) {
         val block = element.findChildBlock() ?: return
         val statement = block.statements.single()
+        val caretOnAfterStatement = if (editor != null) editor.caretModel.offset >= statement.endOffset else false
 
         val container = block.parent
         val construct = container.parent as KtExpression
         handleComments(construct, block)
 
         val newElement = block.replace(statement.copy())
-
+        editor?.caretModel?.moveToOffset(if (caretOnAfterStatement) newElement.endOffset else newElement.startOffset)
+        
         val factory = KtPsiFactory(block)
 
         if (construct is KtDoWhileExpression) {
@@ -75,7 +79,10 @@ class RemoveBracesIntention : SelfTargetingIntention<KtElement>(KtElement::class
             construct.parent is KtExpression &&
             construct.parent !is KtStatementExpression
         ) {
-            construct.replace(factory.createExpressionByPattern("($0)", construct))
+            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)
+            }
         }
     }
 

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

@@ -3,5 +3,5 @@ fun test(b: Boolean) {
     val list1 = mutableListOf(1)
     val list2 = mutableListOf(2)
 
-    (if (b) list1 else list2).add(3)
+    (if (b) list1 else <caret>list2).add(3)
 }

idea/testData/intentions/removeBraces/elseInPlusAssignExpression.kt

@@ -3,5 +3,5 @@ fun test(b: Boolean) {
     val list1 = mutableListOf(1)
     val list2 = mutableListOf(2)
 
-    if (b) {list1} else <caret>{list2} += 3
+    if (b) {list1} else {list2}<caret> += 3
 }

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

@@ -3,5 +3,5 @@ fun test(b: Boolean) {
     val list1 = mutableListOf(1)
     val list2 = mutableListOf(2)
 
-    (if (b) {list1} else list2) += 3
+    (if (b) {list1} else list2<caret>) += 3
 }

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

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

idea/testData/intentions/removeBraces/forWithLocalVariable.kt

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

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

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