From 184838ae354f5980478cac35538be89fe29fc63a Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Tue, 16 Mar 2021 15:25:51 +0300
Subject: [PATCH] FIR IDE: Fix import of properties; also fix bug in
 `alreadyHasImport`

---
 ...eadyImportedExtensionPropertyNoImport-1.kt |  8 ++++++
 ...eadyImportedExtensionPropertyNoImport-2.kt |  4 +++
 ...ImportedExtensionPropertyNoImport.kt.after |  8 ++++++
 .../CompletionMultifileHandlerTest.kt         |  4 +++
 .../KotlinFirLookupElementFactory.kt          | 28 +++++++++++--------
 5 files changed, 41 insertions(+), 11 deletions(-)
 create mode 100644 idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport-1.kt
 create mode 100644 idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport-2.kt
 create mode 100644 idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport.kt.after

diff --git a/idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport-1.kt b/idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport-1.kt
new file mode 100644
index 00000000000..37f833aca97
--- /dev/null
+++ b/idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport-1.kt
@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package first
+
+import second.extensionProp
+
+fun foo() {
+    "".extensi<caret>
+}
\ No newline at end of file
diff --git a/idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport-2.kt b/idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport-2.kt
new file mode 100644
index 00000000000..7584c805362
--- /dev/null
+++ b/idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport-2.kt
@@ -0,0 +1,4 @@
+package second
+
+val String.extensionProp: Int
+    get() = 1
\ No newline at end of file
diff --git a/idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport.kt.after b/idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport.kt.after
new file mode 100644
index 00000000000..cb99e91f410
--- /dev/null
+++ b/idea/idea-completion/testData/handlers/multifile/AlreadyImportedExtensionPropertyNoImport.kt.after
@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package first
+
+import second.extensionProp
+
+fun foo() {
+    "".extensionProp<caret>
+}
\ No newline at end of file
diff --git a/idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/handlers/CompletionMultifileHandlerTest.kt b/idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/handlers/CompletionMultifileHandlerTest.kt
index 769865ffa02..76394f93bff 100644
--- a/idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/handlers/CompletionMultifileHandlerTest.kt
+++ b/idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/handlers/CompletionMultifileHandlerTest.kt
@@ -26,6 +26,10 @@ open class CompletionMultiFileHandlerTest : KotlinFixtureCompletionBaseTestCase(
         doTest()
     }
 
+    fun testAlreadyImportedExtensionPropertyNoImport() {
+        doTest()
+    }
+
     fun testExtensionPropertyInSamePackageNoImport() {
         doTest()
     }
diff --git a/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirLookupElementFactory.kt b/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirLookupElementFactory.kt
index 16ac096d07f..c8cee5647f7 100644
--- a/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirLookupElementFactory.kt
+++ b/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirLookupElementFactory.kt
@@ -132,11 +132,15 @@ private class VariableLookupElementFactory {
         if (setterName != null) "$getterName()/$setterName()" else "$getterName()"
 
     private fun detectImportStrategy(symbol: KtVariableLikeSymbol): CallableImportStrategy {
-        if (symbol !is KtKotlinPropertySymbol) return CallableImportStrategy.DoNothing
-        if (symbol.dispatchType != null || symbol.receiverType != null) return CallableImportStrategy.DoNothing
+        if (symbol !is KtKotlinPropertySymbol || symbol.dispatchType != null) return CallableImportStrategy.DoNothing
 
-        return symbol.callableIdIfNonLocal?.let(CallableImportStrategy::InsertFqNameAndShorten)
-            ?: CallableImportStrategy.DoNothing
+        val propertyId = symbol.callableIdIfNonLocal ?: return CallableImportStrategy.DoNothing
+
+        return if (symbol.isExtension) {
+            CallableImportStrategy.AddImport(propertyId)
+        } else {
+            CallableImportStrategy.InsertFqNameAndShorten(propertyId)
+        }
     }
 }
 
@@ -163,12 +167,14 @@ private class FunctionLookupElementFactory {
     }
 
     private fun detectImportStrategy(symbol: KtFunctionSymbol): CallableImportStrategy {
-        val functionFqName = symbol.callableIdIfNonLocal
-        return when {
-            functionFqName == null -> CallableImportStrategy.DoNothing
-            symbol.dispatchType != null -> CallableImportStrategy.DoNothing
-            !symbol.isExtension -> CallableImportStrategy.InsertFqNameAndShorten(functionFqName)
-            else -> CallableImportStrategy.AddImport(functionFqName)
+        if (symbol.dispatchType != null) return CallableImportStrategy.DoNothing
+
+        val functionFqName = symbol.callableIdIfNonLocal ?: return CallableImportStrategy.DoNothing
+
+        return if (symbol.isExtension) {
+            CallableImportStrategy.AddImport(functionFqName)
+        } else {
+            CallableImportStrategy.InsertFqNameAndShorten(functionFqName)
         }
     }
 
@@ -356,7 +362,7 @@ private fun addCallableImportIfRequired(targetFile: KtFile, nameToImport: FqName
 }
 
 private fun alreadyHasImport(file: KtFile, nameToImport: FqName): Boolean {
-    if (file.importDirectives.any { it.importPath?.fqName == nameToImport }) return false
+    if (file.importDirectives.any { it.importPath?.fqName == nameToImport }) return true
 
     withAllowedResolve {
         analyze(file) {