KtStringTemplateExpressionManipulator: properly works with interpolations (KT-24958)

compiler/frontend/src/org/jetbrains/kotlin/psi/psiUtil/KtStringTemplateExpressionManipulator.kt

@@ -16,12 +16,17 @@
 
 package org.jetbrains.kotlin.psi.psiUtil
 
+import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.AbstractElementManipulator
+import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.psi.KtStringTemplateEntryWithExpression
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
 
+private val LOG = Logger.getInstance(KtStringTemplateExpressionManipulator::class.java)
+
 class KtStringTemplateExpressionManipulator : AbstractElementManipulator<KtStringTemplateExpression>() {
     override fun handleContentChange(
         element: KtStringTemplateExpression,
@@ -29,11 +34,34 @@ class KtStringTemplateExpressionManipulator : AbstractElementManipulator<KtStrin
         newContent: String
     ): KtStringTemplateExpression? {
         val node = element.node
-        val content = if (element.isSingleQuoted()) StringUtil.escapeStringCharacters(newContent) else newContent
         val oldText = node.text
-        val newText = oldText.substring(0, range.startOffset) + content + oldText.substring(range.endOffset)
-        val expression = KtPsiFactory(element.project).createExpression(newText)
-        node.replaceAllChildrenToChildrenOf(expression.node)
+
+        fun wrapAsInOld(content: String) = oldText.substring(0, range.startOffset) + content + oldText.substring(range.endOffset)
+
+        fun makeKtExpressionFromText(text: String): KtExpression {
+            val ktExpression = KtPsiFactory(element.project).createExpression(text)
+            if (ktExpression !is KtStringTemplateExpression) {
+                LOG.error("can't create a `KtStringTemplateExpression` from '$text'")
+            }
+            return ktExpression
+        }
+
+        val newContentPreprocessed: String =
+            if (element.isSingleQuoted()) {
+                val expressionFromText = makeKtExpressionFromText("\"\"\"$newContent\"\"\"")
+                if (expressionFromText is KtStringTemplateExpression) {
+                    expressionFromText.entries.joinToString("") { entry ->
+                        when (entry) {
+                            is KtStringTemplateEntryWithExpression -> entry.text
+                            else -> StringUtil.escapeStringCharacters(entry.text)
+                        }
+                    }
+                } else newContent
+            } else newContent
+
+        val newKtExpression = makeKtExpressionFromText(wrapAsInOld(newContentPreprocessed))
+        node.replaceAllChildrenToChildrenOf(newKtExpression.node)
+
         return node.getPsi(KtStringTemplateExpression::class.java)
     }