Optimize completion for overridden descriptors

core/descriptors/src/org/jetbrains/kotlin/resolve/DescriptorUtils.kt

@@ -233,3 +233,20 @@ fun CallableMemberDescriptor.overriddenTreeAsSequence(useOriginal: Boolean): Seq
     with(if (useOriginal) original else this) {
         sequenceOf(this) + overriddenDescriptors.asSequence().flatMap { it.overriddenTreeAsSequence(useOriginal) }
     }
+
+fun CallableDescriptor.overriddenTreeUniqueAsSequence(useOriginal: Boolean): Sequence<CallableDescriptor> {
+    val set = hashSetOf<CallableDescriptor>()
+
+    fun CallableDescriptor.doBuildOverriddenTreeAsSequence(): Sequence<CallableDescriptor> {
+        return with(if (useOriginal) original else this) {
+            if (original in set)
+                emptySequence()
+            else {
+                set += original
+                sequenceOf(this) + overriddenDescriptors.asSequence().flatMap { it.doBuildOverriddenTreeAsSequence() }
+            }
+        }
+    }
+
+    return doBuildOverriddenTreeAsSequence()
+}

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

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.renderer.DescriptorRenderer
 import org.jetbrains.kotlin.renderer.render
 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.synthetic.SamAdapterExtensionFunctionDescriptor
@@ -269,11 +270,6 @@ class LookupElementFactory(
         if (receiverTypes == null) return null
         if (descriptor !is CallableDescriptor) return null
 
-        val overridden = descriptor.overriddenDescriptors
-        if (overridden.isNotEmpty()) {
-            return overridden.map { callableWeight(it)!! }.min()!!
-        }
-
         // don't treat synthetic extensions as real extensions
         if (descriptor is SyntheticJavaPropertyDescriptor) {
             return callableWeight(descriptor.getMethod)
@@ -282,29 +278,52 @@ class LookupElementFactory(
             return callableWeight(descriptor.sourceFunction)
         }
 
-        val receiverParameter = descriptor.extensionReceiverParameter ?: descriptor.dispatchReceiverParameter
-        if (receiverParameter != null) {
-            return if (receiverTypes.any { TypeUtils.equalTypes(it, receiverParameter.type) }) {
-                when {
-                    descriptor.isExtensionForTypeParameter() -> CallableWeight.typeParameterExtension
-                    descriptor.isExtension -> CallableWeight.thisTypeExtension
-                    else -> CallableWeight.thisClassMember
-                }
-            }
-            else if (receiverTypes.any { it.isSubtypeOf(receiverParameter.type) }) {
-                if (descriptor.isExtension) CallableWeight.baseTypeExtension else CallableWeight.baseClassMember
-            }
-            else {
-                CallableWeight.receiverCastRequired
-            }
+        if (descriptor.overriddenDescriptors.isNotEmpty()) {
+            // Optimization: when one of direct overridden fits, then nothing can fit better
+            descriptor.overriddenDescriptors.mapNotNull {
+                it.callableWeightBasedOnReceiver(receiverTypes, onReceiverTypeMismatch = null)
+            }.min()?.let { return it }
+
+            val overridden = descriptor.overriddenTreeUniqueAsSequence(useOriginal = false)
+            return overridden.map { callableWeightBasic(it, receiverTypes)!! }.min()!!
         }
 
+        return callableWeightBasic(descriptor, receiverTypes)
+    }
+
+    private fun callableWeightBasic(descriptor: CallableDescriptor, receiverTypes: Collection<KotlinType>): CallableWeight? {
+        descriptor.callableWeightBasedOnReceiver(receiverTypes, CallableWeight.receiverCastRequired)?.let { return it }
+
         return when (descriptor.containingDeclaration) {
             is PackageFragmentDescriptor, is ClassifierDescriptor -> CallableWeight.globalOrStatic
             else -> CallableWeight.local
         }
     }
 
+    private fun CallableDescriptor.callableWeightBasedOnReceiver(
+            receiverTypes: Collection<KotlinType>,
+            onReceiverTypeMismatch: CallableWeight?
+    ): CallableWeight? {
+        val receiverParameter = extensionReceiverParameter ?: dispatchReceiverParameter
+        if (receiverParameter != null) {
+            return if (receiverTypes.any { TypeUtils.equalTypes(it, receiverParameter.type) }) {
+                when {
+                    isExtensionForTypeParameter() -> CallableWeight.typeParameterExtension
+                    isExtension -> CallableWeight.thisTypeExtension
+                    else -> CallableWeight.thisClassMember
+                }
+            }
+            else if (receiverTypes.any { it.isSubtypeOf(receiverParameter.type) }) {
+                if (isExtension) CallableWeight.baseTypeExtension else CallableWeight.baseClassMember
+            }
+            else {
+                onReceiverTypeMismatch
+            }
+        }
+
+        return null
+    }
+
     private fun CallableDescriptor.isExtensionForTypeParameter(): Boolean {
         val receiverParameter = original.extensionReceiverParameter ?: return false
         val typeParameter = receiverParameter.type.constructor.declarationDescriptor as? TypeParameterDescriptor ?: return false