[kotlin] K2 Introduce variable: enable extracting substring

^KTIJ-27960 Fixed
Existing tests in `K2IntroduceVariableTestGenerated.StringTemplates` don't pass because "Introduce Variable" in K2 doesn't support multiple occurrences at the moment (KTIJ-27861)

GitOrigin-RevId: a4b3817cd3c5c5b673b1c1bc1de8d4f8e65df68b

plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceUtil.kt

@@ -145,52 +145,30 @@ fun selectElementsWithTargetParent(
     selectSingleElement()
 }
 
-fun findExpressionOrStringFragment(file: KtFile, startOffset: Int, endOffset: Int): KtExpression? {
-    val entry1 = file.findElementAt(startOffset)?.getNonStrictParentOfType<KtStringTemplateEntry>() ?: return null
-    val entry2 = file.findElementAt(endOffset - 1)?.getNonStrictParentOfType<KtStringTemplateEntry>() ?: return null
+fun findStringTemplateFragment(file: KtFile, startOffset: Int, endOffset: Int, kind: ElementKind): KtExpression? {
+    if (kind != ElementKind.EXPRESSION) return null
 
-    if (entry1 == entry2 && entry1 is KtStringTemplateEntryWithExpression) return entry1.expression
+    val startEntry = file.findElementAt(startOffset)?.getNonStrictParentOfType<KtStringTemplateEntry>() ?: return null
+    val endEntry = file.findElementAt(endOffset - 1)?.getNonStrictParentOfType<KtStringTemplateEntry>() ?: return null
 
-    val stringTemplate = entry1.parent as? KtStringTemplateExpression ?: return null
-    if (entry2.parent != stringTemplate) return null
+    if (startEntry.parent !is KtStringTemplateExpression || startEntry.parent != endEntry.parent) return null
 
-    val templateOffset = stringTemplate.startOffset
-    if (stringTemplate.getContentRange().equalsToRange(startOffset - templateOffset, endOffset - templateOffset)) return stringTemplate
+    val prefixOffset = startOffset - startEntry.startOffset
+    if (startEntry !is KtLiteralStringTemplateEntry && prefixOffset > 0) return null
 
-    val prefixOffset = startOffset - entry1.startOffset
-    if (entry1 !is KtLiteralStringTemplateEntry && prefixOffset > 0) return null
+    val suffixOffset = endOffset - endEntry.startOffset
+    if (endEntry !is KtLiteralStringTemplateEntry && suffixOffset < endEntry.textLength) return null
 
-    val suffixOffset = endOffset - entry2.startOffset
-    if (entry2 !is KtLiteralStringTemplateEntry && suffixOffset < entry2.textLength) return null
+    val prefix = startEntry.text.substring(0, prefixOffset)
+    val suffix = endEntry.text.substring(suffixOffset)
 
-    val prefix = entry1.text.substring(0, prefixOffset)
-    val suffix = entry2.text.substring(suffixOffset)
-
-    return K1ExtractableSubstringInfo(entry1, entry2, prefix, suffix).createExpression()
+    return K1ExtractableSubstringInfo(startEntry, endEntry, prefix, suffix).createExpression()
 }
 
 fun KotlinPsiRange.getPhysicalTextRange(): TextRange {
     return (elements.singleOrNull() as? KtExpression)?.extractableSubstringInfo?.contentRange ?: textRange
 }
 
-fun ExtractableSubstringInfo.replaceWith(replacement: KtExpression): KtExpression {
-    return with(this) {
-        val psiFactory = KtPsiFactory(replacement.project)
-        val parent = startEntry.parent
-
-        psiFactory.createStringTemplate(prefix).entries.singleOrNull()?.let { parent.addBefore(it, startEntry) }
-
-        val refEntry = psiFactory.createBlockStringTemplateEntry(replacement)
-        val addedRefEntry = parent.addBefore(refEntry, startEntry) as KtStringTemplateEntryWithExpression
-
-        psiFactory.createStringTemplate(suffix).entries.singleOrNull()?.let { parent.addAfter(it, endEntry) }
-
-        parent.deleteChildRange(startEntry, endEntry)
-
-        addedRefEntry.expression!!
-    }
-}
-
 fun isObjectOrNonInnerClass(e: PsiElement): Boolean = e is KtObjectDeclaration || (e is KtClass && !e.isInner())
 
 fun <T : KtDeclaration> insertDeclaration(declaration: T, targetSibling: PsiElement): T {

plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinIntroduceVariableServiceK1Impl.kt

@@ -11,7 +11,8 @@ import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.refactoring.introduce.IntroduceRefactoringException
 import org.jetbrains.kotlin.idea.refactoring.introduce.KotlinIntroduceVariableHelper
 import org.jetbrains.kotlin.idea.refactoring.introduce.KotlinIntroduceVariableService
-import org.jetbrains.kotlin.idea.refactoring.introduce.findExpressionOrStringFragment
+import org.jetbrains.kotlin.idea.refactoring.introduce.findStringTemplateFragment
+import org.jetbrains.kotlin.idea.refactoring.introduce.findStringTemplateOrStringTemplateEntryExpression
 import org.jetbrains.kotlin.idea.util.ElementKind
 import org.jetbrains.kotlin.idea.util.findElement
 import org.jetbrains.kotlin.psi.KtElement
@@ -30,9 +31,8 @@ internal class KotlinIntroduceVariableServiceK1Impl(private val project: Project
         elementKind: ElementKind
     ): PsiElement? {
         var element = findElement(file, startOffset, endOffset, elementKind)
-        if (element == null && elementKind == ElementKind.EXPRESSION) {
-            element = findExpressionOrStringFragment(file, startOffset, endOffset)
-        }
+            ?: findStringTemplateOrStringTemplateEntryExpression(file, startOffset, endOffset, elementKind)
+            ?: findStringTemplateFragment(file, startOffset, endOffset, elementKind)
 
         if (element is KtExpression) {
             val qualifier = element.analyze().get(BindingContext.QUALIFIER, element)

plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/refactoring/introduce/AbstractExtractionTest.kt

@@ -5,8 +5,10 @@ package org.jetbrains.kotlin.idea.refactoring.introduce
 import com.intellij.codeInsight.CodeInsightUtil
 import com.intellij.codeInsight.completion.JavaCompletionUtil
 import com.intellij.codeInsight.template.impl.TemplateManagerImpl
+import com.intellij.codeInsight.template.impl.TemplateState
 import com.intellij.ide.DataManager
 import com.intellij.openapi.application.runReadAction
+import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.io.FileUtil
@@ -90,8 +92,20 @@ abstract class AbstractExtractionTest : KotlinLightCodeInsightFixtureTestCase()
                 DataManager.getInstance().getDataContext(fixture.editor.component)
             )
 
+            val inplaceVariableName = InTextDirectivesUtils.findStringWithPrefixes(file.text, "// INPLACE_VARIABLE_NAME:")
             val templateState = TemplateManagerImpl.getTemplateState(editor)
-            if (templateState?.isFinished() == false) {
+
+            if (inplaceVariableName != null) {
+                templateState as TemplateState
+
+                WriteCommandAction.runWriteCommandAction(project) {
+                    val range = templateState.currentVariableRange ?: error("No variable range was found")
+                    templateState.editor.document.replaceString(range.startOffset, range.endOffset, inplaceVariableName)
+                    templateState.update()
+                }
+            }
+
+            if (templateState?.isFinished == false) {
                 project.executeCommand("") { templateState.gotoEnd(false) }
             }
         }
@@ -195,8 +209,8 @@ abstract class AbstractExtractionTest : KotlinLightCodeInsightFixtureTestCase()
                 true
             )
             val suggestedNames = JavaNameSuggestionUtil.appendUnresolvedExprName(
-              JavaCompletionUtil.completeVariableNameForRefactoring(codeStyleManager, type, VariableKind.LOCAL_VARIABLE, info),
-              initializer
+                JavaCompletionUtil.completeVariableNameForRefactoring(codeStyleManager, type, VariableKind.LOCAL_VARIABLE, info),
+                initializer
             )
 
             IntroduceParameterProcessor(
@@ -366,7 +380,7 @@ abstract class AbstractExtractionTest : KotlinLightCodeInsightFixtureTestCase()
     protected fun doExtractInterfaceTest(path: String) = doExtractSuperTest(path, true)
 
     protected fun doTestIfNotDisabledByFileDirective(action: (PsiFile) -> Unit) {
-        val disableTestDirective = if (isFirPlugin) IgnoreTests.DIRECTIVES.IGNORE_K2 else  IgnoreTests.DIRECTIVES.IGNORE_K1
+        val disableTestDirective = if (isFirPlugin) IgnoreTests.DIRECTIVES.IGNORE_K2 else IgnoreTests.DIRECTIVES.IGNORE_K1
 
         IgnoreTests.runTestIfNotDisabledByFileDirective(
             dataFilePath(),
@@ -553,7 +567,7 @@ fun doExtractFunction(fixture: CodeInsightTestFixture, file: KtFile) {
                     descriptor
                 }
 
-                fun afterFinish(extraction: ExtractionResult){
+                fun afterFinish(extraction: ExtractionResult) {
                     processDuplicates(extraction.duplicateReplacers, project, editor)
                     onFinish(extraction)
                 }

plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/refactoring/introduce/ExtractionTestGenerated.java

@@ -304,6 +304,11 @@ public abstract class ExtractionTestGenerated extends AbstractExtractionTest {
             public void testSingleEntrySuffix() throws Exception {
                 runTest("testData/refactoring/introduceVariable/stringTemplates/singleEntrySuffix.kt");
             }
+
+            @TestMetadata("singleOccurence.kt")
+            public void testSingleOccurence() throws Exception {
+                runTest("testData/refactoring/introduceVariable/stringTemplates/singleOccurence.kt");
+            }
         }
 
         @RunWith(JUnit3RunnerWithInners.class)

plugins/kotlin/idea/tests/testData/refactoring/introduceVariable/WhenAddBlock.kt

@@ -1,5 +1,7 @@
+// INPLACE_VARIABLE_NAME: xxx
 fun a() {
-    when (1) {
+    return when (1) {
         is 1 -> <selection>2</selection>
+        else -> 3
     }
 }

plugins/kotlin/idea/tests/testData/refactoring/introduceVariable/WhenAddBlock.kt.after

@@ -1,7 +1,10 @@
+// INPLACE_VARIABLE_NAME: xxx
 fun a() {
-    when (1) {
+    return when (1) {
         is 1 -> {
-            val i = 2
+            val xxx = 2
+            xxx
         }
+        else -> 3
     }
 }

plugins/kotlin/idea/tests/testData/refactoring/introduceVariable/stringTemplates/singleOccurence.kt

@@ -0,0 +1,4 @@
+// INPLACE_VARIABLE_NAME: s
+fun foo(): String {
+    return "abc<selection>def</selection>"
+}

plugins/kotlin/idea/tests/testData/refactoring/introduceVariable/stringTemplates/singleOccurence.kt.after

@@ -0,0 +1,5 @@
+// INPLACE_VARIABLE_NAME: s
+fun foo(): String {
+    val s = "def"
+    return "abc$s"
+}

plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceUtils.kt

@@ -6,12 +6,9 @@ import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.idea.base.psi.dropCurlyBracketsIfPossible
-import org.jetbrains.kotlin.psi.KtBinaryExpression
-import org.jetbrains.kotlin.psi.KtBlockStringTemplateEntry
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
-import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
-import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
+import org.jetbrains.kotlin.idea.util.ElementKind
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.*
 
 fun KtExpression.removeTemplateEntryBracesIfPossible(): KtExpression {
     val parent = parent as? KtBlockStringTemplateEntry ?: return this
@@ -41,4 +38,39 @@ fun PsiElement.findExpressionsByCopyableDataAndClearIt(key: Key<Boolean>): List<
     val results = collectDescendantsOfType<KtExpression> { it.getCopyableUserData(key) != null }
     results.forEach { it.putCopyableUserData(key, null) }
     return results
+}
+
+fun ExtractableSubstringInfo.replaceWith(replacement: KtExpression): KtExpression {
+    return with(this) {
+        val psiFactory = KtPsiFactory(replacement.project)
+        val parent = startEntry.parent
+
+        psiFactory.createStringTemplate(prefix).entries.singleOrNull()?.let { parent.addBefore(it, startEntry) }
+
+        val refEntry = psiFactory.createBlockStringTemplateEntry(replacement)
+        val addedRefEntry = parent.addBefore(refEntry, startEntry) as KtStringTemplateEntryWithExpression
+
+        psiFactory.createStringTemplate(suffix).entries.singleOrNull()?.let { parent.addAfter(it, endEntry) }
+
+        parent.deleteChildRange(startEntry, endEntry)
+
+        addedRefEntry.expression!!
+    }
+}
+
+fun findStringTemplateOrStringTemplateEntryExpression(file: KtFile, startOffset: Int, endOffset: Int, kind: ElementKind): KtExpression? {
+    if (kind != ElementKind.EXPRESSION) return null
+
+    val startEntry = file.findElementAt(startOffset)?.getNonStrictParentOfType<KtStringTemplateEntry>() ?: return null
+    val endEntry = file.findElementAt(endOffset - 1)?.getNonStrictParentOfType<KtStringTemplateEntry>() ?: return null
+
+    if (startEntry == endEntry && startEntry is KtStringTemplateEntryWithExpression) return startEntry.expression
+
+    val stringTemplate = startEntry.parent as? KtStringTemplateExpression ?: return null
+    if (endEntry.parent != stringTemplate) return null
+
+    val templateOffset = stringTemplate.startOffset
+    if (stringTemplate.getContentRange().equalsToRange(startOffset - templateOffset, endOffset - templateOffset)) return stringTemplate
+
+    return null
 }

plugins/kotlin/refactorings/kotlin.refactorings.introduce.k2/src/org/jetbrains/kotlin/idea/k2/refactoring/introduce/K2ExtractableSubstringInfo.kt

@@ -0,0 +1,19 @@
+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.kotlin.idea.k2.refactoring.introduce
+
+import org.jetbrains.kotlin.idea.refactoring.introduce.ExtractableSubstringInfo
+import org.jetbrains.kotlin.psi.KtStringTemplateEntry
+
+class K2ExtractableSubstringInfo(
+    startEntry: KtStringTemplateEntry,
+    endEntry: KtStringTemplateEntry,
+    prefix: String,
+    suffix: String,
+) : ExtractableSubstringInfo(startEntry, endEntry, prefix, suffix) {
+    override val isString: Boolean = true   // TODO: KTIJ-28404
+
+    override fun copy(
+        newStartEntry: KtStringTemplateEntry,
+        newEndEntry: KtStringTemplateEntry
+    ): ExtractableSubstringInfo = K2ExtractableSubstringInfo(newStartEntry, newEndEntry, prefix, suffix)
+}

plugins/kotlin/refactorings/kotlin.refactorings.introduce.k2/src/org/jetbrains/kotlin/idea/k2/refactoring/introduce/introduceVariable/K2IntroduceVariableHandler.kt

@@ -36,9 +36,12 @@ import org.jetbrains.kotlin.idea.codeinsight.utils.getRenderedTypeArguments
 import org.jetbrains.kotlin.idea.refactoring.KotlinCommonRefactoringSettings
 import org.jetbrains.kotlin.idea.refactoring.introduce.KotlinIntroduceVariableHandler
 import org.jetbrains.kotlin.idea.refactoring.introduce.KotlinIntroduceVariableHelper.Containers
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractableSubstringInfo
 import org.jetbrains.kotlin.idea.refactoring.introduce.findExpressionByCopyableDataAndClearIt
 import org.jetbrains.kotlin.idea.refactoring.introduce.mustBeParenthesizedInInitializerPosition
 import org.jetbrains.kotlin.idea.refactoring.introduce.removeTemplateEntryBracesIfPossible
+import org.jetbrains.kotlin.idea.refactoring.introduce.replaceWith
+import org.jetbrains.kotlin.idea.refactoring.introduce.substringContextOrThis
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.util.application.executeCommand
 import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
@@ -81,6 +84,7 @@ object K2IntroduceVariableHandler : KotlinIntroduceVariableHandler() {
             val isActualExpression = expression == expressionToReplace
 
             val replacement = psiFactory.createExpression(nameSuggestions.single().first())
+            val substringInfo = expressionToReplace.extractableSubstringInfo
             var result = when {
                 expressionToReplace.isLambdaOutsideParentheses() -> {
                     val functionLiteralArgument = expressionToReplace.getStrictParentOfType<KtLambdaArgument>()!!
@@ -88,6 +92,7 @@ object K2IntroduceVariableHandler : KotlinIntroduceVariableHandler() {
                     newCallExpression.valueArguments.last().getArgumentExpression()!!
                 }
 
+                substringInfo != null -> substringInfo.replaceWith(replacement)
                 else -> expressionToReplace.replace(replacement) as KtExpression
             }
 
@@ -104,7 +109,7 @@ object K2IntroduceVariableHandler : KotlinIntroduceVariableHandler() {
             return result
         }
 
-        private fun createBasicPropertyOrDestructuringDeclaration(isVar: Boolean): KtDeclaration {
+        private fun createBasicPropertyOrDestructuringDeclaration(expression: KtExpression, isVar: Boolean): KtDeclaration {
             val initializer = (expression as? KtParenthesizedExpression)?.expression ?: expression
             val initializerText = if (initializer.mustBeParenthesizedInInitializerPosition()) "(${initializer.text})" else initializer.text
 
@@ -313,8 +318,9 @@ object K2IntroduceVariableHandler : KotlinIntroduceVariableHandler() {
     }
 
     private fun calculateAnchor(expression: PsiElement, container: PsiElement): PsiElement? {
-        if (expression != container) return expression.parentsWithSelf.firstOrNull { it.parent == container }
-        val startOffset = min(container.endOffset, expression.startOffset)
+        val physicalExpression = expression.substringContextOrThis
+        if (physicalExpression != container) return physicalExpression.parentsWithSelf.firstOrNull { it.parent == container }
+        val startOffset = min(container.endOffset, physicalExpression.startOffset)
         return container.allChildren.lastOrNull { it.textRange.contains(startOffset) }
     }
 

plugins/kotlin/refactorings/kotlin.refactorings.introduce.k2/src/org/jetbrains/kotlin/idea/k2/refactoring/introduce/introduceVariable/KotlinIntroduceVariableServiceK2Impl.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.idea.k2.refactoring.introduce.introduceVariable.K2In
 import org.jetbrains.kotlin.idea.refactoring.introduce.IntroduceRefactoringException
 import org.jetbrains.kotlin.idea.refactoring.introduce.KotlinIntroduceVariableHelper
 import org.jetbrains.kotlin.idea.refactoring.introduce.KotlinIntroduceVariableService
+import org.jetbrains.kotlin.idea.refactoring.introduce.findStringTemplateOrStringTemplateEntryExpression
 import org.jetbrains.kotlin.idea.util.ElementKind
 import org.jetbrains.kotlin.idea.util.findElement
 import org.jetbrains.kotlin.psi.KtElement
@@ -25,6 +26,8 @@ internal class KotlinIntroduceVariableServiceK2Impl(private val project: Project
         elementKind: ElementKind
     ): PsiElement? {
         val element = findElement(file, startOffset, endOffset, elementKind)
+            ?: findStringTemplateOrStringTemplateEntryExpression(file, startOffset, endOffset, elementKind)
+            ?: findStringTemplateFragment(file, startOffset, endOffset, elementKind)
 
         if (element == null) {
             if (failOnNoExpression) {

plugins/kotlin/refactorings/kotlin.refactorings.introduce.k2/src/org/jetbrains/kotlin/idea/k2/refactoring/introduce/introduceVariable/introduceVariableUtil.kt

@@ -7,8 +7,16 @@ import com.intellij.openapi.ui.popup.JBPopupFactory
 import org.jetbrains.kotlin.idea.base.analysis.api.utils.analyzeInModalWindow
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.codeinsight.utils.getParameterNames
+import org.jetbrains.kotlin.idea.k2.refactoring.introduce.K2ExtractableSubstringInfo
+import org.jetbrains.kotlin.idea.util.ElementKind
 import org.jetbrains.kotlin.idea.util.application.isUnitTestMode
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtLiteralStringTemplateEntry
+import org.jetbrains.kotlin.psi.KtStringTemplateEntry
+import org.jetbrains.kotlin.psi.KtStringTemplateExpression
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 internal fun chooseApplicableComponentNames(
     contextExpression: KtExpression,
@@ -40,3 +48,24 @@ internal fun chooseApplicableComponentNames(
         .createPopup()
         .showInBestPositionFor(editor)
 }
+
+// a copy-paste of `org.jetbrains.kotlin.idea.refactoring.introduce.findStringTemplateFragment` from `idea.kotlin` module
+internal fun findStringTemplateFragment(file: KtFile, startOffset: Int, endOffset: Int, kind: ElementKind): KtExpression? {
+    if (kind != ElementKind.EXPRESSION) return null
+
+    val startEntry = file.findElementAt(startOffset)?.getNonStrictParentOfType<KtStringTemplateEntry>() ?: return null
+    val endEntry = file.findElementAt(endOffset - 1)?.getNonStrictParentOfType<KtStringTemplateEntry>() ?: return null
+
+    if (startEntry.parent !is KtStringTemplateExpression || startEntry.parent != endEntry.parent) return null
+
+    val prefixOffset = startOffset - startEntry.startOffset
+    if (startEntry !is KtLiteralStringTemplateEntry && prefixOffset > 0) return null
+
+    val suffixOffset = endOffset - endEntry.startOffset
+    if (endEntry !is KtLiteralStringTemplateEntry && suffixOffset < endEntry.textLength) return null
+
+    val prefix = startEntry.text.substring(0, prefixOffset)
+    val suffix = endEntry.text.substring(suffixOffset)
+
+    return K2ExtractableSubstringInfo(startEntry, endEntry, prefix, suffix).createExpression()
+}

plugins/kotlin/refactorings/kotlin.refactorings.introduce.k2/test/org/jetbrains/kotlin/idea/k2/refactoring/introduce/introduceVariable/K2IntroduceVariableTestGenerated.java

@@ -302,6 +302,11 @@ public abstract class K2IntroduceVariableTestGenerated extends AbstractK2Introdu
         public void testSingleEntrySuffix() throws Exception {
             runTest("../../idea/tests/testData/refactoring/introduceVariable/stringTemplates/singleEntrySuffix.kt");
         }
+
+        @TestMetadata("singleOccurence.kt")
+        public void testSingleOccurence() throws Exception {
+            runTest("../../idea/tests/testData/refactoring/introduceVariable/stringTemplates/singleOccurence.kt");
+        }
     }
 
     @RunWith(JUnit3RunnerWithInners.class)