Dsl completion: assign high priority for dsl elements that are members

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

@@ -238,6 +238,7 @@ class BasicLookupElementFactory(
             }
         }
 
+        var isMarkedAsDsl = false
         if (descriptor is CallableDescriptor) {
             appendContainerAndReceiverInformation(descriptor) { element = element.appendTailText(it, true) }
 
@@ -245,6 +246,7 @@ class BasicLookupElementFactory(
                 EditorColorsManager.getInstance().globalScheme.getAttributes(it)
             }
             if (dslTextAttributes != null) {
+                isMarkedAsDsl = true
                 element = element.withBoldness(dslTextAttributes.fontType == Font.BOLD)
                 dslTextAttributes.foregroundColor?.let { element = element.withItemTextForeground(it) }
             }
@@ -268,7 +270,9 @@ class BasicLookupElementFactory(
             element.putUserData(KotlinCompletionCharFilter.ACCEPT_OPENING_BRACE, Unit)
         }
 
-        return element.withIconFromLookupObject()
+        val result = element.withIconFromLookupObject()
+        result.isDslMember = isMarkedAsDsl
+        return result
     }
 
     fun appendContainerAndReceiverInformation(descriptor: CallableDescriptor, appendTailText: (String) -> Unit) {

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

@@ -6,8 +6,6 @@
 package org.jetbrains.kotlin.idea.completion
 
 import com.intellij.codeInsight.completion.PrefixMatcher
-import com.intellij.codeInsight.lookup.LookupElement
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.idea.core.KotlinIndicesHelper
 import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
 import org.jetbrains.kotlin.idea.util.ReceiverType
@@ -26,30 +24,6 @@ class DslMembersCompletion(
     private val nearestReceiver = receiverTypes?.firstOrNull()
     private val nearestReceiverMarkers = nearestReceiver?.takeIf { it.implicit }?.extractDslMarkers().orEmpty()
 
-    private val factory = object : AbstractLookupElementFactory {
-        override fun createLookupElement(
-            descriptor: DeclarationDescriptor,
-            useReceiverTypes: Boolean,
-            qualifyNestedClasses: Boolean,
-            includeClassTypeArguments: Boolean,
-            parametersAndTypeGrayed: Boolean
-        ): LookupElement? {
-            error("Should not be called")
-        }
-
-        override fun createStandardLookupElementsForDescriptor(
-            descriptor: DeclarationDescriptor,
-            useReceiverTypes: Boolean
-        ): Collection<LookupElement> {
-            return elementFactory.createStandardLookupElementsForDescriptor(descriptor, useReceiverTypes).also {
-                it.forEach { element ->
-                    element.isDslMember = true
-                }
-            }
-        }
-    }
-
-
     fun completeDslFunctions() {
         if (nearestReceiver == null || nearestReceiverMarkers.isEmpty()) return
 
@@ -64,7 +38,7 @@ class DslMembersCompletion(
             receiverTypes = listOf(nearestReceiver.type)
         )
         extensionDescriptors.forEach {
-            collector.addDescriptorElements(it, factory, notImported = true, withReceiverCast = false, prohibitDuplicates = true)
+            collector.addDescriptorElements(it, elementFactory, notImported = true, withReceiverCast = false, prohibitDuplicates = true)
         }
 
         collector.flushToResultSet()

idea/idea-completion/testData/weighers/basic/DslMemberCalls.Data1.kt

@@ -0,0 +1,38 @@
+package bar
+
+@DslMarker
+annotation class Dsl
+
+
+@Dsl
+class R {
+    fun foo1(i: Int) {
+
+    }
+
+    @Dsl
+    fun foo2(i : Int) {
+
+    }
+
+    val foo5: Int = 6
+    @Dsl
+    val foo6: Int = 7
+
+
+}
+
+fun r(body: R.() -> Unit) {
+
+}
+
+
+@Dsl
+fun R.foo4() {
+
+}
+
+@Dsl
+fun R.foo3() {
+
+}

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

@@ -0,0 +1,20 @@
+// RUNTIME
+package test
+
+import bar.r
+import bar.foo3
+
+fun main() {
+    val fooLocal = 3
+    r {
+        foo<caret>
+    }
+}
+
+// ORDER: foo2
+// ORDER: foo3
+// ORDER: foo6
+// ORDER: foo4
+// ORDER: foo1
+// ORDER: fooLocal
+// ORDER: foo5

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/weighers/BasicCompletionWeigherTestGenerated.java

@@ -79,6 +79,11 @@ public class BasicCompletionWeigherTestGenerated extends AbstractBasicCompletion
         runTest("idea/idea-completion/testData/weighers/basic/DslCallsWithMultipleReceivers.kt");
     }
 
+    @TestMetadata("DslMemberCalls.kt")
+    public void testDslMemberCalls() throws Exception {
+        runTest("idea/idea-completion/testData/weighers/basic/DslMemberCalls.kt");
+    }
+
     @TestMetadata("ExactMatchForKeyword.kt")
     public void testExactMatchForKeyword() throws Exception {
         runTest("idea/idea-completion/testData/weighers/basic/ExactMatchForKeyword.kt");