Better caret position after Insert curly braces around variable (KT-28739)

Introduce helper class and move caret to it's original position.

 #KT-28739 Fixed

idea/idea-core/src/org/jetbrains/kotlin/idea/core/generateUtil.kt

@@ -44,6 +44,66 @@ fun moveCaretIntoGeneratedElement(editor: Editor, element: PsiElement) {
     pointer.element?.let { moveCaretIntoGeneratedElementDocumentUnblocked(editor, it) }
 }
 
+class RestoreCaret<T: PsiElement>(beforeElement: T, val editor: Editor?) {
+    private val relativeOffset: Int
+    private val beforeElementTextLength: Int = beforeElement.textLength
+
+    init {
+        relativeOffset = findRelativeOffset(beforeElement, editor)
+    }
+
+    fun restoreCaret(afterElement: T, defaultOffset: ((element: T) -> Int)? = null) {
+        if (editor == null) return
+        val project = editor.project ?: return
+
+        val pointer = SmartPointerManager.getInstance(project).createSmartPsiElementPointer(afterElement)
+
+        val document = editor.document
+        PsiDocumentManager.getInstance(project).commitDocument(document)
+        PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(document)
+
+        val afterElementChanged = pointer.element ?: return
+
+        val offset = if (relativeOffset != -1 && afterElementChanged.textLength == beforeElementTextLength) {
+            afterElementChanged.startOffset + relativeOffset
+        } else {
+            if (defaultOffset != null) {
+                defaultOffset(afterElementChanged)
+            } else {
+                -1
+            }
+        }
+
+        if (offset == -1) {
+            return
+        }
+
+        if (document.textLength > offset) {
+            editor.caretModel.moveToOffset(offset)
+        }
+    }
+
+    companion object {
+        fun findRelativeOffset(element: PsiElement, editor: Editor?): Int {
+            if (editor != null) {
+                val singleCaret = editor.caretModel.allCarets.singleOrNull()
+                if (singleCaret != null) {
+                    val caretOffset = singleCaret.offset
+                    val textRange = element.textRange
+                    if (textRange.startOffset <= caretOffset && caretOffset <= textRange.endOffset) {
+                        val relative = caretOffset - element.startOffset
+                        if (relative >= 0) {
+                            return relative
+                        }
+                    }
+                }
+            }
+
+            return -1
+        }
+    }
+}
+
 private fun moveCaretIntoGeneratedElementDocumentUnblocked(editor: Editor, element: PsiElement): Boolean {
     // Inspired by GenerateMembersUtils.positionCaret()
 

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

@@ -18,16 +18,28 @@ package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.codeInsight.intention.LowPriorityAction
 import com.intellij.openapi.editor.Editor
+import org.jetbrains.kotlin.idea.core.RestoreCaret
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.KtSimpleNameStringTemplateEntry
+import org.jetbrains.kotlin.psi.KtStringTemplateEntryWithExpression
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
 
-class InsertCurlyBracesToTemplateIntention : SelfTargetingOffsetIndependentIntention<KtSimpleNameStringTemplateEntry> (
+class InsertCurlyBracesToTemplateIntention :
-        KtSimpleNameStringTemplateEntry::class.java, "Insert curly braces around variable"), LowPriorityAction {
+    SelfTargetingOffsetIndependentIntention<KtSimpleNameStringTemplateEntry>(
+        KtSimpleNameStringTemplateEntry::class.java, "Insert curly braces around variable"
+    ),
+    LowPriorityAction {
 
     override fun isApplicableTo(element: KtSimpleNameStringTemplateEntry): Boolean = true
 
     override fun applyTo(element: KtSimpleNameStringTemplateEntry, editor: Editor?) {
         val expression = element.expression ?: return
-        element.replace(KtPsiFactory(element).createBlockStringTemplateEntry(expression))
+
+        with(RestoreCaret(expression, editor)) {
+            val wrapped = element.replace(KtPsiFactory(element).createBlockStringTemplateEntry(expression))
+            val afterExpression = (wrapped as? KtStringTemplateEntryWithExpression)?.expression ?: return
+
+            restoreCaret(afterExpression, defaultOffset = { it.endOffset })
+        }
     }
 }

idea/testData/intentions/insertCurlyBracesToTemplate/insertBrackets1.kt.after

@@ -1,5 +1,5 @@
 // IS_APPLICABLE: true
 fun foo() {
     val x = 4
-    val y = "${x}"
+    val y = "${<caret>x}"
 }

idea/testData/intentions/insertCurlyBracesToTemplate/insertBrackets2.kt.after

@@ -1,5 +1,5 @@
 // IS_APPLICABLE: true
 fun foo() {
     val x = 4
-    val y = "something${x}.somethingelse"
+    val y = "something${<caret>x}.somethingelse"
 }

idea/testData/intentions/insertCurlyBracesToTemplate/insertBrackets4.kt.after

@@ -1,5 +1,5 @@
 // IS_APPLICABLE: true
 fun foo() {
     val x = 4
-    val y = "$x${x}$x"
+    val y = "$x${<caret>x}$x"
 }

idea/testData/intentions/insertCurlyBracesToTemplate/insertBrackets5.kt.after

@@ -1,5 +1,5 @@
 // IS_APPLICABLE: true
 fun foo() {
     val x = 4
-    val y = """$x${x}$x"""
+    val y = """$x${<caret>x}$x"""
 }

idea/testData/intentions/insertCurlyBracesToTemplate/insertBracketsEnd.kt

@@ -0,0 +1,2 @@
+val test = "some"
+val other = "text $test<caret> text"

idea/testData/intentions/insertCurlyBracesToTemplate/insertBracketsEnd.kt.after

@@ -0,0 +1,2 @@
+val test = "some"
+val other = "text ${test<caret>} text"

idea/testData/intentions/insertCurlyBracesToTemplate/insertBracketsMiddle.kt

@@ -0,0 +1,2 @@
+val test = "some"
+val other = "text $te<caret>st text"

idea/testData/intentions/insertCurlyBracesToTemplate/insertBracketsMiddle.kt.after

@@ -0,0 +1,2 @@
+val test = "some"
+val other = "text ${te<caret>st} text"

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

@@ -9466,6 +9466,16 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
         public void testInsertBrackets5() throws Exception {
             runTest("idea/testData/intentions/insertCurlyBracesToTemplate/insertBrackets5.kt");
         }
+
+        @TestMetadata("insertBracketsEnd.kt")
+        public void testInsertBracketsEnd() throws Exception {
+            runTest("idea/testData/intentions/insertCurlyBracesToTemplate/insertBracketsEnd.kt");
+        }
+
+        @TestMetadata("insertBracketsMiddle.kt")
+        public void testInsertBracketsMiddle() throws Exception {
+            runTest("idea/testData/intentions/insertCurlyBracesToTemplate/insertBracketsMiddle.kt");
+        }
     }
 
     @TestMetadata("idea/testData/intentions/insertExplicitTypeArguments")