From 79199260b9d072a859d1431448a32cd33a7380f2 Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Tue, 17 Sep 2019 06:30:26 +0900
Subject: [PATCH] Import quick fix: suggest for operator extension function
 called from name reference

#KT-28049 Fixed
---
 .../kotlin/idea/quickfix/ImportFix.kt         | 12 ++++--
 .../autoImports/invokeExtension2.test         | 42 +++++++++++++++++++
 .../QuickFixMultiFileTestGenerated.java       |  5 +++
 3 files changed, 55 insertions(+), 4 deletions(-)
 create mode 100644 idea/testData/quickfix/autoImports/invokeExtension2.test

diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/ImportFix.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/ImportFix.kt
index b04817654f5..de43a2da4b8 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/ImportFix.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/ImportFix.kt
@@ -50,12 +50,10 @@ import org.jetbrains.kotlin.idea.core.KotlinIndicesHelper
 import org.jetbrains.kotlin.idea.core.isVisible
 import org.jetbrains.kotlin.idea.imports.canBeReferencedViaImport
 import org.jetbrains.kotlin.idea.imports.importableFqName
+import org.jetbrains.kotlin.idea.intentions.getCallableDescriptor
 import org.jetbrains.kotlin.idea.project.TargetPlatformDetector
 import org.jetbrains.kotlin.idea.references.mainReference
-import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
-import org.jetbrains.kotlin.idea.util.ReceiverType
-import org.jetbrains.kotlin.idea.util.getResolutionScope
-import org.jetbrains.kotlin.idea.util.receiverTypesWithIndex
+import org.jetbrains.kotlin.idea.util.*
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -261,6 +259,12 @@ internal abstract class OrdinaryImportFixBase<T : KtExpression>(expression: T, f
 
                 indicesHelper.getTopLevelCallablesByName(name).filterTo(result, filterByCallType)
             }
+            if (callTypeAndReceiver.callType == CallType.OPERATOR) {
+                val type = expression.getCallableDescriptor()?.returnType
+                if (type != null) {
+                    result.addAll(indicesHelper.getCallableTopLevelExtensions(callTypeAndReceiver, listOf(type), { it == name }))
+                }
+            }
         }
 
         result.addAll(indicesHelper.getCallableTopLevelExtensions(callTypeAndReceiver, expression, bindingContext) { it == name })
diff --git a/idea/testData/quickfix/autoImports/invokeExtension2.test b/idea/testData/quickfix/autoImports/invokeExtension2.test
new file mode 100644
index 00000000000..bd5461c50c5
--- /dev/null
+++ b/idea/testData/quickfix/autoImports/invokeExtension2.test
@@ -0,0 +1,42 @@
+// FILE: first.before.kt
+// "Import" "true"
+// ERROR: Expression 'a' of type 'A' cannot be invoked as a function. The function 'invoke()' is not found
+
+package b
+
+import a.A
+
+fun main() {
+    val a = A()
+    <caret>a()
+}
+//-----------------------
+
+
+// FILE: second.kt
+// "Import" "true"
+// ERROR: Expression 'a' of type 'A' cannot be invoked as a function. The function 'invoke()' is not found
+
+package a
+
+class A
+operator fun A.invoke() = 42
+//-----------------------
+
+
+// FILE: first.after.kt
+// "Import" "true"
+// ERROR: Expression 'a' of type 'A' cannot be invoked as a function. The function 'invoke()' is not found
+
+package b
+
+import a.A
+import a.invoke
+
+fun main() {
+    val a = A()
+    <caret>a()
+}
+//-----------------------
+
+
diff --git a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiFileTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiFileTestGenerated.java
index 702ccb96aec..cb2de5d3d6e 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiFileTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiFileTestGenerated.java
@@ -763,6 +763,11 @@ public class QuickFixMultiFileTestGenerated extends AbstractQuickFixMultiFileTes
             runTest("idea/testData/quickfix/autoImports/invokeExtension.test");
         }
 
+        @TestMetadata("invokeExtension2.test")
+        public void testInvokeExtension2() throws Exception {
+            runTest("idea/testData/quickfix/autoImports/invokeExtension2.test");
+        }
+
         @TestMetadata("invokeExtensionNoOperator.test")
         public void testInvokeExtensionNoOperator() throws Exception {
             runTest("idea/testData/quickfix/autoImports/invokeExtensionNoOperator.test");