Fix "Introduce import alias" on extensions

#KT-30214 Fixed

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceImportAlias/KotlinIntroduceImportAliasHandler.kt

@@ -40,8 +40,10 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElement
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.resolve.PropertyImportedFromObject
+import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.utils.findClassifier
 import org.jetbrains.kotlin.resolve.scopes.utils.findFunction
@@ -58,12 +60,14 @@ object KotlinIntroduceImportAliasHandler : RefactoringActionHandler {
         val file = element.containingKtFile
         val declarationDescriptors = file.resolveImportReference(fqName)
 
-        @Suppress("UNCHECKED_CAST")
-        val usages = declarationDescriptors
-            .flatMap { findPsiElements(project, file, it) }
-            .flatMap {
-                ReferencesSearch.search(it, file.useScope).findAll() as List<KtSimpleNameReference>
+        val usages = declarationDescriptors.flatMap { descriptor ->
+            val isExtension = descriptor.isExtension
+            findPsiElements(project, file, descriptor).flatMap {
+                ReferencesSearch.search(it, file.useScope)
+                    .findAll()
+                    .map { reference -> UsageContext(reference as KtSimpleNameReference, isExtension = isExtension) }
             }
+        }
 
         val suggestedName = suggestedName(element.mainReference.value, file.getResolutionScope())
         ImportInsertHelperImpl.addImport(project, file, fqName, false, Name.identifier(suggestedName))
@@ -87,6 +91,8 @@ object KotlinIntroduceImportAliasHandler : RefactoringActionHandler {
     }
 }
 
+private data class UsageContext(val reference: KtSimpleNameReference, val isExtension: Boolean)
+
 private fun cleanImport(file: KtFile, fqName: FqName) {
     file.importDirectives.find { it.alias == null && fqName == it.importedFqName }?.delete()
 }
@@ -125,11 +131,16 @@ private fun invokeRename(project: Project, editor: Editor, file: KtFile) {
     KotlinRenameDispatcherHandler().invoke(project, editor, file, dataContext)
 }
 
-private fun replaceUsages(usages: List<KtSimpleNameReference>, newName: String) {
-    usages.filter { !it.isImportUsage() }
+private fun replaceUsages(usages: List<UsageContext>, newName: String) {
+    usages.filter { !it.reference.isImportUsage() }
         .reversed() // case: inner element
-        .map {
-            val newExpression = it.handleElementRename(newName) as KtNameReferenceExpression
+        .forEach {
+            val newExpression = it.reference.handleElementRename(newName) as KtNameReferenceExpression
+            if (it.isExtension) {
+                newExpression.getQualifiedElementSelector()?.replace(newExpression)
+                return@forEach
+            }
+
             val qualifiedElement = newExpression.getQualifiedElement()
             if (qualifiedElement != newExpression) {
                 val parent = newExpression.parent

idea/testData/intentions/introduceImportAlias/withReceiver.kt

@@ -0,0 +1,11 @@
+package my.sample
+
+class A
+
+fun A.check() {}
+
+fun test() {
+    val a = A()
+    a.check<caret>()
+    A().check()
+}

idea/testData/intentions/introduceImportAlias/withReceiver.kt.after

@@ -0,0 +1,13 @@
+package my.sample
+
+import my.sample.check as check1
+
+class A
+
+fun A.check() {}
+
+fun test() {
+    val a = A()
+    a.check1()
+    A().check1()
+}

idea/testData/intentions/introduceImportAlias/withReceiver2.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+package my.sample
+
+val Any.foo: Any get() = this
+val Any.bar: Any get() = this
+
+fun test() {
+    1.foo.foo.bar<caret>.foo.bar.toString()
+}

idea/testData/intentions/introduceImportAlias/withReceiver2.kt.after

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+package my.sample
+
+import my.sample.bar as bar1
+
+val Any.foo: Any get() = this
+val Any.bar: Any get() = this
+
+fun test() {
+    1.foo.foo.bar1.foo.bar1.toString()
+}

idea/testData/intentions/introduceImportAlias/withReceiver3.kt

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+package my.sample
+
+fun test() {
+    0.let { true }.let<caret> { println(it) }
+}

idea/testData/intentions/introduceImportAlias/withReceiver3.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+package my.sample
+
+import kotlin.let as let1
+
+fun test() {
+    0.let1 { true }.let1 { println(it) }
+}

idea/testData/intentions/introduceImportAlias/withReceiver4.kt

@@ -0,0 +1,9 @@
+package my.sample
+
+class A
+
+fun A.check() {}
+
+fun test() {
+    (my.sample.A::check<caret>)(A())
+}

idea/testData/intentions/introduceImportAlias/withReceiver4.kt.after

@@ -0,0 +1,11 @@
+package my.sample
+
+import my.sample.check as check1
+
+class A
+
+fun A.check() {}
+
+fun test() {
+    (my.sample.A::check1)(A())
+}

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

@@ -10074,6 +10074,26 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
         public void testVariable() throws Exception {
             runTest("idea/testData/intentions/introduceImportAlias/variable.kt");
         }
+
+        @TestMetadata("withReceiver.kt")
+        public void testWithReceiver() throws Exception {
+            runTest("idea/testData/intentions/introduceImportAlias/withReceiver.kt");
+        }
+
+        @TestMetadata("withReceiver2.kt")
+        public void testWithReceiver2() throws Exception {
+            runTest("idea/testData/intentions/introduceImportAlias/withReceiver2.kt");
+        }
+
+        @TestMetadata("withReceiver3.kt")
+        public void testWithReceiver3() throws Exception {
+            runTest("idea/testData/intentions/introduceImportAlias/withReceiver3.kt");
+        }
+
+        @TestMetadata("withReceiver4.kt")
+        public void testWithReceiver4() throws Exception {
+            runTest("idea/testData/intentions/introduceImportAlias/withReceiver4.kt");
+        }
     }
 
     @TestMetadata("idea/testData/intentions/introduceVariable")