Fixed completion for extension function after removing members for implicit receiver from scope.

idea/ide-common/src/org/jetbrains/kotlin/idea/codeInsight/ReferenceVariantsHelper.kt

@@ -29,12 +29,9 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
-import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.SmartCastUtils
-import org.jetbrains.kotlin.resolve.scopes.DescriptorKindExclude
-import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
-import org.jetbrains.kotlin.resolve.scopes.JetScope
-import org.jetbrains.kotlin.resolve.scopes.getDescriptorsFiltered
+import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
+import org.jetbrains.kotlin.resolve.scopes.*
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
@@ -110,6 +107,10 @@ public class ReferenceVariantsHelper(
 
         val dataFlowInfo = context.getDataFlowInfo(expression)
 
+        val implicitReceiverTypes = resolutionScope.getImplicitReceiversWithInstance().flatMap {
+            SmartCastUtils.getSmartCastVariantsWithLessSpecificExcluded(it.value, context, containingDeclaration, dataFlowInfo)
+        }.toSet()
+
         val pair = getExplicitReceiverData(expression)
         if (pair != null) {
             val (receiverExpression, callType) = pair
@@ -126,29 +127,22 @@ public class ReferenceVariantsHelper(
                                         context.getType(receiverExpression)
             if (expressionType != null && !expressionType.isError()) {
                 val receiverValue = ExpressionReceiver(receiverExpression, expressionType)
-                descriptors.addMembersFromReceiverAndSyntheticExtensions(receiverValue, callType, kindFilter, nameFilter, resolutionScope, dataFlowInfo)
+                val explicitReceiverTypes = SmartCastUtils.getSmartCastVariantsWithLessSpecificExcluded(receiverValue, context, containingDeclaration, dataFlowInfo)
 
-                descriptors.addCallableExtensions(resolutionScope, receiverValue, dataFlowInfo, callType, kindFilter, nameFilter)
+                descriptors.processAll(implicitReceiverTypes, explicitReceiverTypes, resolutionScope, callType, kindFilter, nameFilter)
             }
         }
         else {
-            // process instance members that can be called via implicit receiver's instances
-            val receivers = resolutionScope.getImplicitReceiversWithInstance()
-            val receiverValues = receivers.map { it.getValue() }
-            for (receiverValue in receiverValues) {
-                descriptors.addMembersFromReceiverAndSyntheticExtensions(receiverValue, CallType.NORMAL, kindFilter, nameFilter, resolutionScope, dataFlowInfo)
-            }
+            descriptors.processAll(implicitReceiverTypes, implicitReceiverTypes, resolutionScope, CallType.NORMAL, kindFilter, nameFilter)
 
-            // process extensions and non-instance members
+            // process non-instance members
             for (descriptor in resolutionScope.getDescriptorsFiltered(kindFilter, nameFilter)) {
-                if (descriptor is CallableDescriptor && descriptor.getExtensionReceiverParameter() != null) {
-                    val dispatchReceiver = descriptor.getDispatchReceiverParameter()
-                    if (dispatchReceiver == null || dispatchReceiver in receivers) {
-                        descriptors.addAll(descriptor.substituteExtensionIfCallable(receiverValues, context, dataFlowInfo, CallType.NORMAL, containingDeclaration))
+                if (descriptor is CallableDescriptor) {
+                    assert(descriptor.dispatchReceiverParameter == null) {
+                        "Resolution scope with member descriptor: $descriptor. Scope structure: ${JetScopeUtils.printStructure(resolutionScope)}"
                     }
                 }
-                else {
-                    if (descriptor is CallableDescriptor && descriptor.getDispatchReceiverParameter() != null) continue // should already be processed via implicit receivers
+                if (!descriptor.isExtension) {
                     descriptors.add(descriptor)
                 }
             }
@@ -157,25 +151,47 @@ public class ReferenceVariantsHelper(
         return descriptors
     }
 
-    private fun MutableSet<DeclarationDescriptor>.addMembersFromReceiverAndSyntheticExtensions(
-            receiverValue: ReceiverValue,
+    private fun MutableSet<DeclarationDescriptor>.processAll(
+            implicitReceiverTypes: Collection<JetType>,
+            receiverTypes: Collection<JetType>,
+            resolutionScope: JetScope,
             callType: CallType,
             kindFilter: DescriptorKindFilter,
-            nameFilter: (Name) -> Boolean,
-            resolutionScope: JetScope,
-            dataFlowInfo: DataFlowInfo
+            nameFilter: (Name) -> Boolean
+    ) {
+        addNonExtensionMembers(receiverTypes, callType, kindFilter, nameFilter)
+        addMemberExtensions(implicitReceiverTypes, receiverTypes, callType, kindFilter, nameFilter)
+        addScopeAndSyntheticExtensions(resolutionScope, receiverTypes, callType, kindFilter, nameFilter)
+    }
+
+    private fun MutableSet<DeclarationDescriptor>.addMemberExtensions(
+            dispatchReceiverTypes: Collection<JetType>,
+            extensionReceiverTypes: Collection<JetType>,
+            callType: CallType,
+            kindFilter: DescriptorKindFilter,
+            nameFilter: (Name) -> Boolean
+    ) {
+        val memberFilter = kindFilter exclude DescriptorKindExclude.NonExtensions
+        for (dispatchReceiverType in dispatchReceiverTypes) {
+            for (member in dispatchReceiverType.memberScope.getDescriptorsFiltered(memberFilter, nameFilter)) {
+                addAll((member as CallableDescriptor).substituteExtensionIfCallable(extensionReceiverTypes, callType))
+            }
+        }
+    }
+
+    private fun MutableSet<DeclarationDescriptor>.addNonExtensionMembers(
+            receiverTypes: Collection<JetType>,
+            callType: CallType,
+            kindFilter: DescriptorKindFilter,
+            nameFilter: (Name) -> Boolean
     ) {
         val memberFilter = kindFilter exclude DescriptorKindExclude.Extensions
-        val containingDeclaration = resolutionScope.getContainingDeclaration()
-
-        val receiverTypes = SmartCastUtils.getSmartCastVariantsWithLessSpecificExcluded(receiverValue, context, containingDeclaration, dataFlowInfo)
-
         for (receiverType in receiverTypes) {
-            val members = receiverType.getMemberScope().getDescriptorsFiltered(DescriptorKindFilter.ALL, nameFilter) // filter by kind later because of constructors
+            val members = receiverType.memberScope.getDescriptorsFiltered(DescriptorKindFilter.ALL, nameFilter) // filter by kind later because of constructors
             for (member in members) {
                 if (member is ClassDescriptor) {
-                    if (member.isInner()) {
-                        member.getConstructors().filterTo(this) { callType.canCall(it) && memberFilter.accepts(it) }
+                    if (member.isInner) {
+                        member.constructors.filterTo(this) { callType.canCall(it) && memberFilter.accepts(it) }
                     }
                 }
                 else if (callType.canCall(member) && memberFilter.accepts(member)) {
@@ -183,24 +199,42 @@ public class ReferenceVariantsHelper(
                 }
             }
         }
+    }
 
-        if (!kindFilter.excludes.contains(DescriptorKindExclude.Extensions)) {
-            fun processExtension(extension: CallableDescriptor) {
-                if (nameFilter(extension.name) && kindFilter.accepts(extension)) {
-                    addAll(extension.substituteExtensionIfCallable(receiverValue, callType, context, dataFlowInfo, containingDeclaration))
-                }
+    private fun MutableSet<DeclarationDescriptor>.addScopeAndSyntheticExtensions(
+            resolutionScope: JetScope,
+            receiverTypes: Collection<JetType>,
+            callType: CallType,
+            kindFilter: DescriptorKindFilter,
+            nameFilter: (Name) -> Boolean
+    ) {
+        if (kindFilter.excludes.contains(DescriptorKindExclude.Extensions)) return
+
+
+        fun process(extension: CallableDescriptor) {
+            if (nameFilter(extension.name) && kindFilter.accepts(extension)) {
+                addAll(extension.substituteExtensionIfCallable(receiverTypes, callType))
             }
+        }
 
-            if (kindFilter.acceptsKinds(DescriptorKindFilter.VARIABLES_MASK)) {
-                for (extension in resolutionScope.getSyntheticExtensionProperties(receiverTypes)) {
-                    processExtension(extension)
-                }
+        for (descriptor in resolutionScope.getDescriptors(kindFilter exclude DescriptorKindExclude.NonExtensions, nameFilter)) {
+            assert(descriptor !is CallableDescriptor || descriptor.dispatchReceiverParameter == null) {
+                "Resolution scope with member descriptor: $descriptor. Scope structure: ${JetScopeUtils.printStructure(resolutionScope)}"
             }
+            if (descriptor.isExtension) {
+                process(descriptor as CallableDescriptor)
+            }
+        }
 
-            if (kindFilter.acceptsKinds(DescriptorKindFilter.FUNCTIONS_MASK)) {
-                for (extension in resolutionScope.getSyntheticExtensionFunctions(receiverTypes)) {
-                    processExtension(extension)
-                }
+        if (kindFilter.acceptsKinds(DescriptorKindFilter.VARIABLES_MASK)) {
+            for (extension in resolutionScope.getSyntheticExtensionProperties(receiverTypes)) {
+                process(extension)
+            }
+        }
+
+        if (kindFilter.acceptsKinds(DescriptorKindFilter.FUNCTIONS_MASK)) {
+            for (extension in resolutionScope.getSyntheticExtensionFunctions(receiverTypes)) {
+                process(extension)
             }
         }
     }
@@ -227,35 +261,6 @@ public class ReferenceVariantsHelper(
         return type
     }
 
-    private fun MutableCollection<DeclarationDescriptor>.addCallableExtensions(
-            resolutionScope: JetScope,
-            receiver: ReceiverValue,
-            dataFlowInfo: DataFlowInfo,
-            callType: CallType,
-            kindFilter: DescriptorKindFilter,
-            nameFilter: (Name) -> Boolean
-    ) {
-        if (kindFilter.excludes.contains(DescriptorKindExclude.Extensions)) return
-        val extensionsFilter = kindFilter.exclude(DescriptorKindExclude.NonExtensions)
-
-        fun processExtension(descriptor: DeclarationDescriptor) {
-            addAll((descriptor as CallableDescriptor).substituteExtensionIfCallable(receiver, callType, context, dataFlowInfo, resolutionScope.getContainingDeclaration()))
-        }
-
-        // process member extensions from implicit receivers separately to filter out ones from implicit receivers with no instance
-        for (implicitReceiver in resolutionScope.getImplicitReceiversWithInstance()) {
-            for (extension in implicitReceiver.getType().getMemberScope().getDescriptorsFiltered(extensionsFilter, nameFilter)) {
-                processExtension(extension)
-            }
-        }
-
-        for (extension in resolutionScope.getDescriptorsFiltered(extensionsFilter, nameFilter)) {
-            if ((extension as CallableDescriptor).getDispatchReceiverParameter() == null) { // otherwise it should already be processed via implicit receivers
-                processExtension(extension)
-            }
-        }
-    }
-
     public fun getPackageReferenceVariants(
             expression: JetSimpleNameExpression,
             nameFilter: (Name) -> Boolean

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

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.resolve.scopes.JetScope
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.resolve.scopes.receivers.ThisReceiver
+import org.jetbrains.kotlin.types.JetType
 import org.jetbrains.kotlin.types.typeUtil.TypeNullability
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
 import org.jetbrains.kotlin.types.typeUtil.nullability
@@ -82,10 +83,18 @@ public fun CallableDescriptor.substituteExtensionIfCallable(
         containingDeclarationOrModule: DeclarationDescriptor
 ): Collection<CallableDescriptor> {
     if (!receiver.exists()) return listOf()
+
+    var types = SmartCastUtils.getSmartCastVariants(receiver, bindingContext, containingDeclarationOrModule, dataFlowInfo)
+    return substituteExtensionIfCallable(types, callType)
+}
+
+public fun CallableDescriptor.substituteExtensionIfCallable(
+        receiverTypes: Collection<JetType>,
+        callType: CallType
+): Collection<CallableDescriptor> {
     if (!callType.canCall(this)) return listOf()
 
-    var types = SmartCastUtils.getSmartCastVariants(receiver, bindingContext, containingDeclarationOrModule, dataFlowInfo).asSequence()
-
+    var types = receiverTypes.asSequence()
     if (callType == CallType.SAFE) {
         types = types.map { it.makeNotNullable() }
     }

idea/idea-completion/testData/basic/common/boldOrGrayed/NonPredictableSmartCast.kt

@@ -12,4 +12,6 @@ fun f(pair: Pair<out Any, out Any>) {
 // EXIST: { lookupString: "hashCode", attributes: "bold" }
 // EXIST: { lookupString: "forString", attributes: "grayed" }
 // EXIST: { lookupString: "forAny", attributes: "bold" }
-// EXIST: { lookupString: "forT", attributes: "" }
+
+/*TODO: { lookupString: "forT", attributes: "" }*/
+// EXIST: { lookupString: "forT", attributes: "grayed" }

idea/idea-completion/testData/basic/common/boldOrGrayed/NonPredictableSmartCastWithWrongReturnType.kt

@@ -0,0 +1,15 @@
+interface A {
+    fun foo(): Any
+    fun bar()
+}
+interface B: A {
+    override fun foo(): String
+}
+
+fun f(pair: Pair<out A, out Any>) {
+    if (pair.first !is B) return
+    pair.first.<caret>
+}
+/*TODO: { lookupString: "foo", "typeText":"Any", attributes: "bold" } */
+// EXIST: { lookupString: "foo", "typeText":"String", attributes: "bold" }
+// EXIST: { lookupString: "bar", attributes: "bold" }

idea/idea-completion/testData/basic/common/extensions/MemberExtensionWithoutExplicitReceiver.kt

@@ -0,0 +1,20 @@
+class B {
+    fun A.extFunInBForA() {}
+    fun B.extFunInBForB() {}
+
+}
+class A {
+    fun B.extFunInAForB() {}
+    fun A.extFunInAForA() {}
+}
+
+fun A.test() {
+    with(B()) {
+        <caret>
+    }
+}
+
+// EXIST: extFunInAForA
+// EXIST: extFunInBForA
+// EXIST: extFunInAForB
+// EXIST: extFunInBForB

idea/idea-completion/testData/basic/common/smartCast/MemberExtensionAfterThisSmartCast.kt

@@ -0,0 +1,12 @@
+class A {
+    fun String.foo() {}
+}
+
+fun Any.test() { // see KT-8865 Member extension after this smart cast not working
+    if (this is A) {
+        "".<caret>
+    }
+}
+
+
+// EXIST: foo

idea/idea-completion/testData/basic/common/smartCast/MembersAfterSmartCast.kt

@@ -0,0 +1,21 @@
+interface A {
+    fun memberA()
+}
+
+interface B {
+    fun memberB()
+}
+
+interface C {
+    fun memberC()
+}
+
+fun C.test(a: Any) {
+    if (a is A && a is B) {
+        a.<caret>
+    }
+}
+
+// EXIST: memberA
+// EXIST: memberB
+// ABSENT: memberC

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JSBasicCompletionTestGenerated.java

@@ -1110,6 +1110,12 @@ public class JSBasicCompletionTestGenerated extends AbstractJSBasicCompletionTes
                 doTest(fileName);
             }
 
+            @TestMetadata("NonPredictableSmartCastWithWrongReturnType.kt")
+            public void testNonPredictableSmartCastWithWrongReturnType() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/boldOrGrayed/NonPredictableSmartCastWithWrongReturnType.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("SafeCallAfterNullable.kt")
             public void testSafeCallAfterNullable() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/boldOrGrayed/SafeCallAfterNullable.kt");
@@ -1215,6 +1221,12 @@ public class JSBasicCompletionTestGenerated extends AbstractJSBasicCompletionTes
                 doTest(fileName);
             }
 
+            @TestMetadata("MemberExtensionWithoutExplicitReceiver.kt")
+            public void testMemberExtensionWithoutExplicitReceiver() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/extensions/MemberExtensionWithoutExplicitReceiver.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("MemberExtensionsFromMemberExtension.kt")
             public void testMemberExtensionsFromMemberExtension() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/extensions/MemberExtensionsFromMemberExtension.kt");
@@ -1800,6 +1812,18 @@ public class JSBasicCompletionTestGenerated extends AbstractJSBasicCompletionTes
                 JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/idea-completion/testData/basic/common/smartCast"), Pattern.compile("^(.+)\\.kt$"), true);
             }
 
+            @TestMetadata("MemberExtensionAfterThisSmartCast.kt")
+            public void testMemberExtensionAfterThisSmartCast() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/smartCast/MemberExtensionAfterThisSmartCast.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("MembersAfterSmartCast.kt")
+            public void testMembersAfterSmartCast() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/smartCast/MembersAfterSmartCast.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("SmartCastAfterIf.kt")
             public void testSmartCastAfterIf() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/smartCast/SmartCastAfterIf.kt");

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/JvmBasicCompletionTestGenerated.java

@@ -1110,6 +1110,12 @@ public class JvmBasicCompletionTestGenerated extends AbstractJvmBasicCompletionT
                 doTest(fileName);
             }
 
+            @TestMetadata("NonPredictableSmartCastWithWrongReturnType.kt")
+            public void testNonPredictableSmartCastWithWrongReturnType() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/boldOrGrayed/NonPredictableSmartCastWithWrongReturnType.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("SafeCallAfterNullable.kt")
             public void testSafeCallAfterNullable() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/boldOrGrayed/SafeCallAfterNullable.kt");
@@ -1215,6 +1221,12 @@ public class JvmBasicCompletionTestGenerated extends AbstractJvmBasicCompletionT
                 doTest(fileName);
             }
 
+            @TestMetadata("MemberExtensionWithoutExplicitReceiver.kt")
+            public void testMemberExtensionWithoutExplicitReceiver() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/extensions/MemberExtensionWithoutExplicitReceiver.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("MemberExtensionsFromMemberExtension.kt")
             public void testMemberExtensionsFromMemberExtension() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/extensions/MemberExtensionsFromMemberExtension.kt");
@@ -1800,6 +1812,18 @@ public class JvmBasicCompletionTestGenerated extends AbstractJvmBasicCompletionT
                 JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/idea-completion/testData/basic/common/smartCast"), Pattern.compile("^(.+)\\.kt$"), true);
             }
 
+            @TestMetadata("MemberExtensionAfterThisSmartCast.kt")
+            public void testMemberExtensionAfterThisSmartCast() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/smartCast/MemberExtensionAfterThisSmartCast.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("MembersAfterSmartCast.kt")
+            public void testMembersAfterSmartCast() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/smartCast/MembersAfterSmartCast.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("SmartCastAfterIf.kt")
             public void testSmartCastAfterIf() throws Exception {
                 String fileName = JetTestUtils.navigationMetadata("idea/idea-completion/testData/basic/common/smartCast/SmartCastAfterIf.kt");