More correct completion list sorting by receivers

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

@@ -45,7 +45,6 @@ import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindExclude
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
-import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker
 import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
@@ -135,7 +134,7 @@ abstract class CompletionSession(
     protected val referenceVariantsHelper = ReferenceVariantsHelper(bindingContext, resolutionFacade, moduleDescriptor, isVisibleFilter)
 
     protected val callTypeAndReceiver: CallTypeAndReceiver<*, *>
-    protected val receiverTypes: Collection<KotlinType>?
+    protected val receiverTypes: Collection<ReceiverType>?
 
     init {
         val (callTypeAndReceiver, receiverTypes) = detectCallTypeAndReceiverTypes()
@@ -370,7 +369,7 @@ abstract class CompletionSession(
         val filteredNotImportedExtensions = filterVariantsForRuntimeReceiverType(notImportedExtensions, referenceVariants.notImportedExtensions)
 
         val runtimeVariants = ReferenceVariants(filteredVariants, filteredNotImportedExtensions)
-        return Pair(runtimeVariants, lookupElementFactory.copy(receiverTypes = listOf(runtimeType)))
+        return Pair(runtimeVariants, lookupElementFactory.copy(receiverTypes = listOf(ReceiverType(runtimeType, 0))))
     }
 
     private fun <TDescriptor : DeclarationDescriptor> filterVariantsForRuntimeReceiverType(
@@ -430,19 +429,19 @@ abstract class CompletionSession(
                                     callTypeAndReceiver.callType, inDescriptor, contextVariablesProvider)
     }
 
-    private fun detectCallTypeAndReceiverTypes(): Pair<CallTypeAndReceiver<*, *>, Collection<KotlinType>?> {
+    private fun detectCallTypeAndReceiverTypes(): Pair<CallTypeAndReceiver<*, *>, Collection<ReceiverType>?> {
         if (nameExpression == null) {
             return CallTypeAndReceiver.UNKNOWN to null
         }
 
         val callTypeAndReceiver = CallTypeAndReceiver.detect(nameExpression)
 
-        var receiverTypes = callTypeAndReceiver.receiverTypes(
+        var receiverTypes = callTypeAndReceiver.receiverTypesWithIndex(
                 bindingContext, nameExpression, moduleDescriptor, resolutionFacade,
                 stableSmartCastsOnly = true /* we don't include smart cast receiver types for "unstable" receiver value to mark members grayed */)
 
         if (callTypeAndReceiver is CallTypeAndReceiver.SAFE || isDebuggerContext) {
-            receiverTypes = receiverTypes?.map { it.makeNotNullable() }
+            receiverTypes = receiverTypes?.map { ReceiverType(it.type.makeNotNullable(), it.receiverIndex) }
         }
 
         return callTypeAndReceiver to receiverTypes

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

@@ -24,15 +24,15 @@ import org.jetbrains.kotlin.builtins.isExtensionFunctionType
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.util.CallType
+import org.jetbrains.kotlin.idea.util.ReceiverType
 import org.jetbrains.kotlin.idea.util.substituteExtensionIfCallable
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.resolve.calls.tasks.createSynthesizedInvokes
-import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import java.util.*
 
 class ExtensionFunctionTypeValueCompletion(
-        private val receiverTypes: Collection<KotlinType>,
+        private val receiverTypes: Collection<ReceiverType>,
         private val callType: CallType<*>,
         private val lookupElementFactory: LookupElementFactory
 ) {
@@ -49,7 +49,7 @@ class ExtensionFunctionTypeValueCompletion(
 
             val invokes = variableType.memberScope.getContributedFunctions(OperatorNameConventions.INVOKE, NoLookupLocation.FROM_IDE)
             for (invoke in createSynthesizedInvokes(invokes)) {
-                for (substituted in invoke.substituteExtensionIfCallable(receiverTypes, callType)) {
+                for (substituted in invoke.substituteExtensionIfCallable(receiverTypes.map { it.type }, callType)) {
                     val factory = object : AbstractLookupElementFactory {
                         override fun createStandardLookupElementsForDescriptor(descriptor: DeclarationDescriptor, useReceiverTypes: Boolean): Collection<LookupElement> {
                             if (!useReceiverTypes) return emptyList()

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

@@ -30,15 +30,16 @@ import org.jetbrains.kotlin.idea.completion.handlers.KotlinFunctionInsertHandler
 import org.jetbrains.kotlin.idea.completion.handlers.lambdaPresentation
 import org.jetbrains.kotlin.idea.core.moveCaret
 import org.jetbrains.kotlin.idea.util.CallType
+import org.jetbrains.kotlin.idea.util.ReceiverType
 import org.jetbrains.kotlin.idea.util.toFuzzyType
 import org.jetbrains.kotlin.renderer.DescriptorRenderer
 import org.jetbrains.kotlin.renderer.render
+import org.jetbrains.kotlin.resolve.calls.util.getValueParametersCountFromFunctionType
 import org.jetbrains.kotlin.resolve.descriptorUtil.hasDefaultValue
 import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.descriptorUtil.overriddenTreeUniqueAsSequence
 import org.jetbrains.kotlin.resolve.descriptorUtil.parentsWithSelf
 import org.jetbrains.kotlin.resolve.findOriginalTopMostOverriddenDescriptors
-import org.jetbrains.kotlin.resolve.calls.util.getValueParametersCountFromFunctionType
 import org.jetbrains.kotlin.synthetic.SamAdapterExtensionFunctionDescriptor
 import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
 import org.jetbrains.kotlin.types.KotlinType
@@ -62,7 +63,7 @@ interface AbstractLookupElementFactory {
 data /* we need copy() */
 class LookupElementFactory(
         val basicFactory: BasicLookupElementFactory,
-        private val receiverTypes: Collection<KotlinType>?,
+        private val receiverTypes: Collection<ReceiverType>?,
         private val callType: CallType<*>,
         private val inDescriptor: DeclarationDescriptor,
         private val contextVariablesProvider: ContextVariablesProvider,
@@ -324,7 +325,7 @@ class LookupElementFactory(
         return callableWeightBasic(descriptor, receiverTypes)
     }
 
-    private fun callableWeightBasic(descriptor: CallableDescriptor, receiverTypes: Collection<KotlinType>): CallableWeight? {
+    private fun callableWeightBasic(descriptor: CallableDescriptor, receiverTypes: Collection<ReceiverType>): CallableWeight? {
         descriptor.callableWeightBasedOnReceiver(receiverTypes, CallableWeight.receiverCastRequired)?.let { return it }
 
         return when (descriptor.containingDeclaration) {
@@ -334,21 +335,33 @@ class LookupElementFactory(
     }
 
     private fun CallableDescriptor.callableWeightBasedOnReceiver(
-            receiverTypes: Collection<KotlinType>,
+            receiverTypes: Collection<ReceiverType>,
             onReceiverTypeMismatch: CallableWeight?
     ): CallableWeight? {
         val receiverParameter = extensionReceiverParameter
                                 ?: dispatchReceiverParameter
                                 ?: return null
 
-        for ((receiverIndex, receiverType) in receiverTypes.withIndex()) {
-            val weight = callableWeightForReceiverType(receiverType, receiverParameter.type)
+        for (receiverType in receiverTypes) {
+            val weight = callableWeightForReceiverType(receiverType.type, receiverParameter.type)
             if (weight != null) {
-                // if descriptor is matching all receivers then use null as receiverIndex - otherwise e.g. all members of Any would have too high priority
-                val receiverIndexToUse = if (receiverIndex == 0 && receiverTypes.drop(1).all { callableWeightForReceiverType(it, receiverParameter.type) != null })
-                    null
-                else
-                    receiverIndex
+                val receiverIndex = receiverType.receiverIndex
+
+                var receiverIndexToUse: Int? = receiverIndex
+                if (receiverIndex == 0) {
+                    val maxReceiverIndex = receiverTypes.map { it.receiverIndex }.max()!!
+                    if (maxReceiverIndex > 0) {
+                        val matchesAllReceivers = receiverTypes
+                                .filter { it.receiverIndex != 0 && callableWeightForReceiverType(it.type, receiverParameter.type) != null }
+                                .map { it.receiverIndex }
+                                .toSet()
+                                .containsAll((1..maxReceiverIndex).toList())
+                        if (matchesAllReceivers) { // if descriptor is matching all receivers then use null as receiverIndex - otherwise e.g. all members of Any would have too high priority
+                            receiverIndexToUse = null
+                        }
+                    }
+                }
+
                 return CallableWeight(weight, receiverIndexToUse)
             }
         }

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/smart/SmartCompletionSession.kt

@@ -86,7 +86,7 @@ class SmartCompletionSession(
         }
 
         val filter = smartCompletion!!.descriptorFilter
-        var contextVariableTypesForReferenceVariants = filter?.let {
+        val contextVariableTypesForReferenceVariants = filter?.let {
             withCollectRequiredContextVariableTypes { lookupElementFactory ->
                 val (imported, notImported) = referenceVariantsWithNonInitializedVarExcluded ?: return@withCollectRequiredContextVariableTypes
                 imported.forEach { collector.addElements(filter(it, lookupElementFactory)) }

idea/idea-completion/testData/weighers/basic/Callables.kt

@@ -17,7 +17,9 @@ class Derived : Base() {
         fun aaLocalFun(){}
 
         with (y) {
-            aa<caret>
+            if (this is Z1 && this is Z2) {
+                aa<caret>
+            }
         }
     }
 }
@@ -25,10 +27,21 @@ class Derived : Base() {
 interface X {
     fun aaX()
 }
+
 interface Y : X {
     fun aaY()
 }
 
+interface Z1 {
+    fun aaaZ1()
+    fun aabZ1()
+}
+
+interface Z2 {
+    fun aaaZ2()
+    fun aabZ2()
+}
+
 fun Any.aaAnyExtensionFun(){}
 fun Derived.aaExtensionFun(){}
 
@@ -43,6 +56,10 @@ fun Y.aaYExt(){}
 // ORDER: aaLocalVal
 // ORDER: aaLocalFun
 // ORDER: aaY
+// ORDER: aaaZ1
+// ORDER: aaaZ2
+// ORDER: aabZ1
+// ORDER: aabZ2
 // ORDER: aaX
 // ORDER: aaYExt
 // ORDER: aaXExt