From 7b86c6013e604b305b5e37b3fbaf072de129f427 Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Mon, 13 May 2024 21:18:35 +0200
Subject: [PATCH] [kotlin] extract function as local: add filtering on
 resulting call expression

otherwise, when added before the call site, RangeMarker would include the newly extracted function. And if the resulted function contains call expressions, they would be preferred instead of the extracted call expression

^KTIJ-29862 fixed

GitOrigin-RevId: 5249b385ee27d3d46b70663416951788a998e6ad
---
 .../ExtractLocal.after.kt                     |  7 +++++
 .../extractFunctionInplace/ExtractLocal.kt    |  3 +++
 .../AbstractInplaceExtractionHelper.kt        | 26 ++++++++++++++-----
 .../FirInplaceExtractFunctionTest.kt          |  6 ++++-
 4 files changed, 35 insertions(+), 7 deletions(-)
 create mode 100644 plugins/kotlin/idea/tests/testData/refactoring/extractFunctionInplace/ExtractLocal.after.kt
 create mode 100644 plugins/kotlin/idea/tests/testData/refactoring/extractFunctionInplace/ExtractLocal.kt

diff --git a/plugins/kotlin/idea/tests/testData/refactoring/extractFunctionInplace/ExtractLocal.after.kt b/plugins/kotlin/idea/tests/testData/refactoring/extractFunctionInplace/ExtractLocal.after.kt
new file mode 100644
index 000000000000..f307e8da46f9
--- /dev/null
+++ b/plugins/kotlin/idea/tests/testData/refactoring/extractFunctionInplace/ExtractLocal.after.kt
@@ -0,0 +1,7 @@
+fun main() {
+    fun extracted() {
+        val a = "".takeIf { true }
+    }
+
+    extracted<caret>()
+}
\ No newline at end of file
diff --git a/plugins/kotlin/idea/tests/testData/refactoring/extractFunctionInplace/ExtractLocal.kt b/plugins/kotlin/idea/tests/testData/refactoring/extractFunctionInplace/ExtractLocal.kt
new file mode 100644
index 000000000000..4973da33d736
--- /dev/null
+++ b/plugins/kotlin/idea/tests/testData/refactoring/extractFunctionInplace/ExtractLocal.kt
@@ -0,0 +1,3 @@
+fun main() {
+    <selection>val a = "".takeIf { true }</selection>
+}
\ No newline at end of file
diff --git a/plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractFunction/AbstractInplaceExtractionHelper.kt b/plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractFunction/AbstractInplaceExtractionHelper.kt
index 51a5dcb7553d..a08e6412ed76 100644
--- a/plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractFunction/AbstractInplaceExtractionHelper.kt
+++ b/plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractFunction/AbstractInplaceExtractionHelper.kt
@@ -10,6 +10,7 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.keymap.KeymapUtil
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.NlsSafe
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
@@ -29,6 +30,8 @@ import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.processD
 import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtNameReferenceExpression
+import org.jetbrains.kotlin.psi.psiUtil.getCallNameExpression
 import kotlin.math.max
 import kotlin.math.min
 
@@ -82,7 +85,7 @@ interface AbstractInplaceExtractionHelper<KotlinType,
                 return
             }
             val callRange: TextRange = callTextRange.textRange
-            val callIdentifier = findCallExpressionInRange(file, callRange)?.calleeExpression ?: throw IllegalStateException()
+            val callIdentifier = findCallExpressionInRange(file, callRange, extraction.declaration.name)?.calleeExpression ?: throw IllegalStateException()
             val methodIdentifier = extraction.declaration.nameIdentifier ?: throw IllegalStateException()
             val methodRange = extraction.declaration.textRange
             val methodOffset = extraction.declaration.navigationElement.textRange.endOffset
@@ -140,12 +143,23 @@ interface AbstractInplaceExtractionHelper<KotlinType,
         return ::findRange
     }
 
-
-    fun findCallExpressionInRange(file: KtFile, range: TextRange?): KtCallExpression? {
+    /**
+     * Finds the first occurrence of a call expression within the given range.
+     * If [name] is specified, first occurrence with the given reference name would be taken.
+     *
+     * Range is created from the previous element of the extraction, and thus when the function is inserted before the selection,
+     * e.g., when extracted function is added as local function, range contains extracted function and the call to it.
+     * So we need additional ad hock filtering to prevent broken code on rename template start.
+     */
+    fun findCallExpressionInRange(file: KtFile, range: TextRange?, name: @NlsSafe String?): KtCallExpression? {
         if (range == null) return null
-        val container = PsiTreeUtil.findCommonParent(file.findElementAt(range.startOffset),
-                                                     file.findElementAt(min(file.textLength - 1, range.endOffset)))
+        val container = PsiTreeUtil.findCommonParent(
+            file.findElementAt(range.startOffset), file.findElementAt(min(file.textLength - 1, range.endOffset))
+        )
         val callExpressions = PsiTreeUtil.findChildrenOfType(container, KtCallExpression::class.java)
-        return callExpressions.firstOrNull { it.textRange in range }
+        val callExpressionsInRange = callExpressions.filter { it.textRange in range }
+        if (name == null) return callExpressionsInRange.firstOrNull()
+        return callExpressionsInRange.find { (it.calleeExpression as? KtNameReferenceExpression)?.getReferencedName() == name }
+            ?: callExpressionsInRange.firstOrNull()
     }
 }
\ No newline at end of file
diff --git a/plugins/kotlin/refactorings/kotlin.refactorings.tests.k2/test/org/jetbrains/kotlin/idea/k2/refactoring/introduce/FirInplaceExtractFunctionTest.kt b/plugins/kotlin/refactorings/kotlin.refactorings.tests.k2/test/org/jetbrains/kotlin/idea/k2/refactoring/introduce/FirInplaceExtractFunctionTest.kt
index 0cbc6bfdf10f..4f999e4ace7c 100644
--- a/plugins/kotlin/refactorings/kotlin.refactorings.tests.k2/test/org/jetbrains/kotlin/idea/k2/refactoring/introduce/FirInplaceExtractFunctionTest.kt
+++ b/plugins/kotlin/refactorings/kotlin.refactorings.tests.k2/test/org/jetbrains/kotlin/idea/k2/refactoring/introduce/FirInplaceExtractFunctionTest.kt
@@ -54,6 +54,10 @@ class FirInplaceExtractFunctionTest : KotlinLightCodeInsightFixtureTestCase() {
         doTest()
     }
 
+    fun testExtractLocal() {
+        doTest()
+    }
+
     fun testStringTemplateWithNameConflict() {
         doTest(changedName = "substring")
     }
@@ -74,7 +78,7 @@ class FirInplaceExtractFunctionTest : KotlinLightCodeInsightFixtureTestCase() {
     fun doTest(changedName: String? = null, checkResult: Boolean = true) {
         TemplateManagerImpl.setTemplateTesting(testRootDisposable)
         myFixture.configureByFile("${getTestName(false)}.kt")
-        KotlinFirExtractFunctionHandler().invoke(myFixture.project, myFixture.editor, myFixture.file, null)
+        KotlinFirExtractFunctionHandler(getTestName(true).contains("Local")).invoke(myFixture.project, myFixture.editor, myFixture.file, null)
         NonBlockingReadActionImpl.waitForAsyncTaskCompletion()
         val template = getActiveTemplate()
         require(template != null) { "Failed to start refactoring" }