Fix add import quick fix for delegated property and missing extension

^KT-39199 Fixed

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/ReferenceVariantsCollector.kt

@@ -183,7 +183,9 @@ class ReferenceVariantsCollector(
             val extensions = if (runtimeReceiver != null)
                 indicesHelper.getCallableTopLevelExtensions(callTypeAndReceiver, listOf(runtimeReceiver.type), nameFilter)
             else
-                indicesHelper.getCallableTopLevelExtensions(callTypeAndReceiver, nameExpression, bindingContext, nameFilter)
+                indicesHelper.getCallableTopLevelExtensions(
+                    callTypeAndReceiver, nameExpression, bindingContext, receiverTypeFromDiagnostic = null, nameFilter
+                )
 
             val (extensionsVariants, notImportedExtensions) = extensions.partition {
                 importableFqNameClassifier.isImportableDescriptorImported(

idea/idea-core/src/org/jetbrains/kotlin/idea/core/KotlinIndicesHelper.kt

@@ -54,6 +54,7 @@ import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.resolve.scopes.SyntheticScopes
 import org.jetbrains.kotlin.resolve.scopes.collectSyntheticStaticFunctions
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.isError
 import java.lang.reflect.Field
 import java.lang.reflect.Modifier
 import java.util.*
@@ -139,12 +140,21 @@ class KotlinIndicesHelper(
         callTypeAndReceiver: CallTypeAndReceiver<*, *>,
         position: KtExpression,
         bindingContext: BindingContext,
+        receiverTypeFromDiagnostic: KotlinType?,
         nameFilter: (String) -> Boolean
     ): Collection<CallableDescriptor> {
-        val receiverTypes =
+        val receiverTypes = callTypeAndReceiver.receiverTypes(
-            callTypeAndReceiver.receiverTypes(bindingContext, position, moduleDescriptor, resolutionFacade, stableSmartCastsOnly = false)
+            bindingContext, position, moduleDescriptor, resolutionFacade, stableSmartCastsOnly = false
-                ?: return emptyList()
+        )
-        return getCallableTopLevelExtensions(callTypeAndReceiver, receiverTypes, nameFilter)
+
+        return if (receiverTypes == null || receiverTypes.all { it.isError }) {
+            if (receiverTypeFromDiagnostic != null)
+                getCallableTopLevelExtensions(callTypeAndReceiver, listOf(receiverTypeFromDiagnostic), nameFilter)
+            else
+                emptyList()
+        } else {
+            getCallableTopLevelExtensions(callTypeAndReceiver, receiverTypes, nameFilter)
+        }
     }
 
     fun getCallableTopLevelExtensions(

idea/src/org/jetbrains/kotlin/idea/quickfix/AbstractImportFix.kt

@@ -259,7 +259,16 @@ internal abstract class OrdinaryImportFixBase<T : KtExpression>(expression: T, f
             }
         }
 
-        result.addAll(indicesHelper.getCallableTopLevelExtensions(callTypeAndReceiver, expression, bindingContext) { it == name })
+        val receiverFromDiagnostic = if (callTypeAndReceiver.callType == CallType.DELEGATE) getReceiverTypeFromDiagnostic() else null
+        result.addAll(
+            indicesHelper.getCallableTopLevelExtensions(
+                callTypeAndReceiver,
+                expression,
+                bindingContext,
+                receiverFromDiagnostic
+            ) { it == name }
+        )
+
         return result
     }
 }
@@ -494,7 +503,8 @@ internal open class ArrayAccessorImportFix(
 internal class DelegateAccessorsImportFix(
     element: KtExpression,
     override val importNames: Collection<Name>,
-    private val solveSeveralProblems: Boolean
+    private val solveSeveralProblems: Boolean,
+    private val diagnostic: Diagnostic,
 ) : OrdinaryImportFixBase<KtExpression>(element, MyFactory) {
 
     override fun getCallTypeAndReceiver() = CallTypeAndReceiver.DELEGATE(element)
@@ -511,6 +521,9 @@ internal class DelegateAccessorsImportFix(
         return super.createAction(project, editor, element)
     }
 
+    override fun getReceiverTypeFromDiagnostic(): KotlinType? =
+        if (diagnostic.factory === Errors.DELEGATE_SPECIAL_FUNCTION_MISSING) Errors.DELEGATE_SPECIAL_FUNCTION_MISSING.cast(diagnostic).b else null
+
     companion object MyFactory : Factory() {
         private fun importNames(diagnostics: Collection<Diagnostic>): Collection<Name> {
             return diagnostics.map {
@@ -524,14 +537,15 @@ internal class DelegateAccessorsImportFix(
 
         override fun createImportAction(diagnostic: Diagnostic) =
             (diagnostic.psiElement as? KtExpression)?.let {
-                DelegateAccessorsImportFix(it, importNames(listOf(diagnostic)), false)
+                DelegateAccessorsImportFix(it, importNames(listOf(diagnostic)), false, diagnostic)
             }
 
 
         override fun createImportActionsForAllProblems(sameTypeDiagnostics: Collection<Diagnostic>): List<DelegateAccessorsImportFix> {
-            val element = sameTypeDiagnostics.first().psiElement
+            val diagnostic = sameTypeDiagnostics.first()
+            val element = diagnostic.psiElement
             val names = importNames(sameTypeDiagnostics)
-            return listOfNotNull((element as? KtExpression)?.let { DelegateAccessorsImportFix(it, names, true) })
+            return listOfNotNull((element as? KtExpression)?.let { DelegateAccessorsImportFix(it, names, true, diagnostic) })
         }
     }
 }
@@ -641,7 +655,7 @@ internal class ImportForMismatchingArgumentsFix(
         }
 
         indicesHelper
-            .getCallableTopLevelExtensions(callTypeAndReceiver, element, bindingContext) { it == name }
+            .getCallableTopLevelExtensions(callTypeAndReceiver, element, bindingContext, receiverTypeFromDiagnostic = null) { it == name }
             .forEach(::processDescriptor)
 
         if (!isSelectorInQualified(element)) {

idea/testData/quickfix/autoImports/importGetValueExtensionForDelegateWithLambda.after.kt

@@ -0,0 +1,10 @@
+// "Import" "true"
+// WITH_RUNTIME
+// ERROR: Type 'MyDelegate<TypeVariable(T)>' has no method 'getValue(Nothing?, KProperty<*>)' and thus it cannot serve as a delegate
+
+package import
+
+import base.MyDelegate
+import base.getValue
+
+val myVal by MyDelegate { false }

idea/testData/quickfix/autoImports/importGetValueExtensionForDelegateWithLambda.before.Dependency.kt

@@ -0,0 +1,7 @@
+package base
+
+class MyDelegate<T>(init: () -> T) {
+    var value: T = init()
+}
+
+operator fun <T> MyDelegate<T>.getValue(thisObj: Any?, property: kotlin.reflect.KProperty<*>): T = value

idea/testData/quickfix/autoImports/importGetValueExtensionForDelegateWithLambda.before.Main.kt

@@ -0,0 +1,9 @@
+// "Import" "true"
+// WITH_RUNTIME
+// ERROR: Type 'MyDelegate<TypeVariable(T)>' has no method 'getValue(Nothing?, KProperty<*>)' and thus it cannot serve as a delegate
+
+package import
+
+import base.MyDelegate
+
+val myVal by <caret>MyDelegate { false }

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiFileTestGenerated.java

@@ -686,6 +686,11 @@ public class QuickFixMultiFileTestGenerated extends AbstractQuickFixMultiFileTes
             runTest("idea/testData/quickfix/autoImports/importFromRoot.before.Main.kt");
         }
 
+        @TestMetadata("importGetValueExtensionForDelegateWithLambda.before.Main.kt")
+        public void testImportGetValueExtensionForDelegateWithLambda() throws Exception {
+            runTest("idea/testData/quickfix/autoImports/importGetValueExtensionForDelegateWithLambda.before.Main.kt");
+        }
+
         @TestMetadata("importInFirstPartInQualifiedExpression.before.Main.kt")
         public void testImportInFirstPartInQualifiedExpression() throws Exception {
             runTest("idea/testData/quickfix/autoImports/importInFirstPartInQualifiedExpression.before.Main.kt");