KT-19154: Fix auto-import to check receiver for members properly

 #KT-19154 fixed

Author: semoro

idea/ide-common/src/org/jetbrains/kotlin/idea/util/CallType.kt

@@ -20,12 +20,14 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
 import org.jetbrains.kotlin.psi.psiUtil.isImportDirectiveExpression
 import org.jetbrains.kotlin.psi.psiUtil.isPackageDirectiveExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfoBefore
+import org.jetbrains.kotlin.resolve.calls.checkers.DslScopeViolationCallChecker.extractDslMarkerFqNames
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.calls.smartcasts.SmartCastManager
@@ -328,3 +330,20 @@ private fun receiverValueTypes(
         listOf(receiverValue.type)
     }
 }
+
+
+fun Collection<ReceiverType>.shadowedByDslMarkers(): Set<ReceiverType> {
+    val typesByDslScopes = LinkedHashMap<FqName, MutableList<ReceiverType>>()
+
+    this
+            .mapNotNull { receiver ->
+                val dslMarkers = receiver.type.extractDslMarkerFqNames()
+                (receiver to dslMarkers).takeIf { dslMarkers.isNotEmpty() }
+            }
+            .forEach { (v, dslMarkers) -> dslMarkers.forEach { typesByDslScopes.getOrPut(it, { mutableListOf() }) += v } }
+
+    val shadowedDslReceivers = mutableSetOf<ReceiverType>()
+    typesByDslScopes.flatMapTo(shadowedDslReceivers) { (_, v) -> v.asSequence().drop(1).asIterable() }
+
+    return shadowedDslReceivers
+}

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

@@ -403,20 +403,7 @@ abstract class CompletionSession(
         }
 
         if (receiverTypes != null && nameExpression.languageVersionSettings.supportsFeature(LanguageFeature.DslMarkersSupport)) {
-
-            val typesByDslScopes = LinkedHashMap<FqName, MutableList<ReceiverType>>()
-
-            receiverTypes
-                    .mapNotNull { receiver ->
-                        val dslMarkers = receiver.type.extractDslMarkerFqNames()
-                        (receiver to dslMarkers).takeIf { dslMarkers.isNotEmpty() }
-                    }
-                    .forEach { (v, dslMarkers) -> dslMarkers.forEach { typesByDslScopes.getOrPut(it, { mutableListOf() }) += v } }
-
-            val shadowedDslReceivers = mutableSetOf<ReceiverType>()
-            typesByDslScopes.flatMapTo(shadowedDslReceivers) { (_, v) -> v.asSequence().drop(1).asIterable() }
-
-            receiverTypes -= shadowedDslReceivers
+            receiverTypes -= receiverTypes.shadowedByDslMarkers()
         }
 
         return receiverTypes

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

@@ -31,6 +31,7 @@ import com.intellij.psi.PsiErrorElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiModifier
 import com.intellij.psi.util.PsiModificationTracker
+import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.DiagnosticFactory
@@ -51,10 +52,9 @@ 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.project.TargetPlatformDetector
+import org.jetbrains.kotlin.idea.project.languageVersionSettings
 import org.jetbrains.kotlin.idea.references.mainReference
-import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
-import org.jetbrains.kotlin.idea.util.getResolutionScope
-import org.jetbrains.kotlin.idea.util.receiverTypes
+import org.jetbrains.kotlin.idea.util.*
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.js.resolve.JsPlatform
 import org.jetbrains.kotlin.name.FqName
@@ -74,6 +74,7 @@ import org.jetbrains.kotlin.resolve.scopes.utils.addImportingScope
 import org.jetbrains.kotlin.resolve.scopes.utils.collectFunctions
 import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.ifEmpty
 import java.util.*
 
 /**
@@ -300,17 +301,27 @@ internal class ImportFix(expression: KtSimpleNameExpression) : OrdinaryImportFix
         indicesHelper.getKotlinEnumsByName(name).filterTo(result, filterByCallType)
 
         val resolutionFacade = element.getResolutionFacade()
-        val actualReceiverTypes =
-                callTypeAndReceiver.receiverTypes(bindingContext, element, resolutionFacade.moduleDescriptor, resolutionFacade, false).orEmpty()
+        var actualReceiverTypes = callTypeAndReceiver
+                .receiverTypesWithIndex(bindingContext, element,
+                                        resolutionFacade.moduleDescriptor, resolutionFacade,
+                                        stableSmartCastsOnly = false,
+                                        withImplicitReceiversWhenExplicitPresent = true).orEmpty()
+
+        if (element.languageVersionSettings.supportsFeature(LanguageFeature.DslMarkersSupport)) {
+            actualReceiverTypes -= actualReceiverTypes.shadowedByDslMarkers()
+        }
 
-        val processor = { descriptor: CallableDescriptor ->
-            if (descriptor.canBeReferencedViaImport() && filterByCallType(descriptor)) {
-                val extensionReceiverType = descriptor.extensionReceiverParameter?.type
+        val explicitReceiverTypes = actualReceiverTypes.filterNot { it.implicit }
 
-                if ((actualReceiverTypes.isEmpty() && extensionReceiverType == null) ||
-                    (extensionReceiverType != null && actualReceiverTypes.any { it.isSubtypeOf(extensionReceiverType) })) {
-                    result.add(descriptor)
-                }
+        val checkDispatchReceiver = when(callTypeAndReceiver) {
+            is CallTypeAndReceiver.OPERATOR, is CallTypeAndReceiver.INFIX -> true
+            else -> false
+        }
+
+        val processor = { descriptor: CallableDescriptor ->
+            if (descriptor.canBeReferencedViaImport() && filterByCallType(descriptor)
+                && descriptor.isValidByReceiversFor(explicitReceiverTypes, actualReceiverTypes, checkDispatchReceiver)) {
+                result.add(descriptor)
             }
         }
 
@@ -331,6 +342,18 @@ internal class ImportFix(expression: KtSimpleNameExpression) : OrdinaryImportFix
     }
 
 
+    private fun CallableDescriptor.isValidByReceiversFor(explicitReceiverTypes: Collection<ReceiverType>,
+                                                         allReceiverTypes: Collection<ReceiverType>,
+                                                         checkDispatchReceiver: Boolean): Boolean {
+        val bothReceivers = listOfNotNull(extensionReceiverParameter, dispatchReceiverParameter.takeIf { checkDispatchReceiver })
+
+        val receiverTypesPerReceiver = generateSequence(explicitReceiverTypes.ifEmpty { allReceiverTypes }) { allReceiverTypes }
+
+        return bothReceivers
+                .zip(receiverTypesPerReceiver.asIterable())
+                .all { (receiver, possibleTypes) -> possibleTypes.any { it.type.isSubtypeOf(receiver.type) } }
+    }
+
     override fun fillCandidates(
             name: String,
             callTypeAndReceiver: CallTypeAndReceiver<*, *>,

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

@@ -0,0 +1,13 @@
+import BExtSpace.aaa
+
+// "Import" "true"
+// WITH_RUNTIME
+// ERROR: Unresolved reference: aaa
+
+fun test() {
+    AAA().apply {
+        sub {
+            aaa<caret>()
+        }
+    }
+}

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

@@ -0,0 +1,25 @@
+
+@DslMarker
+annotation class DSL
+
+@DSL
+class AAA {
+    fun sub(l: BBB.() -> Unit) {
+        l(BBB())
+    }
+}
+
+@DSL
+class BBB
+
+object AExtSpace {
+    fun AAA.aaa() {
+
+    }
+}
+
+object BExtSpace {
+    fun BBB.aaa() {
+
+    }
+}

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

@@ -0,0 +1,11 @@
+// "Import" "true"
+// WITH_RUNTIME
+// ERROR: Unresolved reference: aaa
+
+fun test() {
+    AAA().apply {
+        sub {
+            aaa<caret>()
+        }
+    }
+}

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

@@ -0,0 +1,11 @@
+// "Import" "true"
+// ERROR: Unresolved reference: foobar
+package p2
+
+import p1.Some.foobar
+
+class A {
+    fun some() {
+        foobar<caret>()
+    }
+}

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

@@ -0,0 +1,5 @@
+package p1
+
+object Some {
+    fun foobar() {}
+}

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

@@ -0,0 +1,9 @@
+// "Import" "true"
+// ERROR: Unresolved reference: foobar
+package p2
+
+class A {
+    fun some() {
+        foobar<caret>()
+    }
+}

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

@@ -131,6 +131,12 @@ public void testDivOperator() throws Exception {
             doTestWithExtraFile(fileName);
         }
 
+        @TestMetadata("dslMarkers.before.Main.kt")
+        public void testDslMarkers() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/autoImports/dslMarkers.before.Main.kt");
+            doTestWithExtraFile(fileName);
+        }
+
         @TestMetadata("extensionFunctionImport.before.Main.kt")
         public void testExtensionFunctionImport() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/autoImports/extensionFunctionImport.before.Main.kt");
@@ -545,6 +551,12 @@ public void testObjectImport() throws Exception {
             doTestWithExtraFile(fileName);
         }
 
+        @TestMetadata("objectMemberFunctionImportWhenReceiverPresent.before.Main.kt")
+        public void testObjectMemberFunctionImportWhenReceiverPresent() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/autoImports/objectMemberFunctionImportWhenReceiverPresent.before.Main.kt");
+            doTestWithExtraFile(fileName);
+        }
+
         @TestMetadata("operatorAssignPlus.test")
         public void testOperatorAssignPlus() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/autoImports/operatorAssignPlus.test");