Take into account refined applicability of dsl-marker in IDE

After the previous change, when being applied to a function type
with receiver it's assumed to work just as it's applied to receiver type

Thus, it's necessary to fix relevant IDE features relied on DSL markers

 #KT-23255 Fixed

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/DslMarkerUtils.kt

@@ -25,7 +25,9 @@ object DslMarkerUtils {
     data class DslMarkersFromReceiver(
         val common: Set<FqName>,
         val fromContainingFunctionType: Set<FqName>
-    )
+    ) {
+        fun all() = common + fromContainingFunctionType
+    }
 
     fun extractDslMarkerFqNames(receiver: ReceiverValue): DslMarkersFromReceiver {
         val errorLevel = extractDslMarkerFqNames(receiver.type)

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

@@ -213,7 +213,17 @@ sealed class CallTypeAndReceiver<TReceiver : KtElement?, out TCallType : CallTyp
     }
 }
 
-data class ReceiverType(val type: KotlinType, val receiverIndex: Int, val implicit: Boolean = false)
+data class ReceiverType(
+    val type: KotlinType,
+    val receiverIndex: Int,
+    val implicitValue: ReceiverValue? = null
+) {
+    val implicit: Boolean get() = implicitValue != null
+
+    fun extractDslMarkers() =
+        implicitValue?.let(DslMarkerUtils::extractDslMarkerFqNames)?.all()
+            ?: DslMarkerUtils.extractDslMarkerFqNames(type)
+}
 
 fun CallTypeAndReceiver<*, *>.receiverTypes(
         bindingContext: BindingContext,
@@ -305,9 +315,13 @@ fun CallTypeAndReceiver<*, *>.receiverTypesWithIndex(
     var receiverIndex = 0
 
     fun addReceiverType(receiverValue: ReceiverValue, implicit: Boolean) {
-        val types = receiverValueTypes(receiverValue, dataFlowInfo, bindingContext, moduleDescriptor, stableSmartCastsOnly, languageVersionSettings,
-                                       resolutionFacade.frontendService<DataFlowValueFactory>())
-        types.mapTo(result) { ReceiverType(it, receiverIndex, implicit) }
+        val types = receiverValueTypes(
+            receiverValue, dataFlowInfo, bindingContext, moduleDescriptor, stableSmartCastsOnly, languageVersionSettings,
+            resolutionFacade.frontendService<DataFlowValueFactory>()
+        )
+
+        types.mapTo(result) { type -> ReceiverType(type, receiverIndex, receiverValue.takeIf { implicit }) }
+
         receiverIndex++
     }
     if (withImplicitReceiversWhenExplicitPresent || expressionReceiver == null) {
@@ -350,7 +364,7 @@ fun Collection<ReceiverType>.shadowedByDslMarkers(): Set<ReceiverType> {
 
     this
             .mapNotNull { receiver ->
-                val dslMarkers = DslMarkerUtils.extractDslMarkerFqNames(receiver.type)
+                val dslMarkers = receiver.extractDslMarkers()
                 (receiver to dslMarkers).takeIf { dslMarkers.isNotEmpty() }
             }
             .forEach { (v, dslMarkers) -> dslMarkers.forEach { typesByDslScopes.getOrPut(it, { mutableListOf() }) += v } }
@@ -359,4 +373,4 @@ fun Collection<ReceiverType>.shadowedByDslMarkers(): Set<ReceiverType> {
     typesByDslScopes.flatMapTo(shadowedDslReceivers) { (_, v) -> v.asSequence().drop(1).asIterable() }
 
     return shadowedDslReceivers
-}
+}

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

@@ -24,8 +24,7 @@ class DslMembersCompletion(
     private val callTypeAndReceiver: CallTypeAndReceiver<*, *>
 ) {
     private val nearestReceiver = receiverTypes?.firstOrNull()
-    private val nearestReceiverMarkers = nearestReceiver?.takeIf { it.implicit }
-        ?.let { DslMarkerUtils.extractDslMarkerFqNames(it.type) }.orEmpty()
+    private val nearestReceiverMarkers = nearestReceiver?.takeIf { it.implicit }?.extractDslMarkers().orEmpty()
 
     private val factory = object : AbstractLookupElementFactory {
         override fun createLookupElement(

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

@@ -0,0 +1,41 @@
+package bar
+
+@DslMarker
+annotation class Dsl
+
+
+class R
+
+fun r(body: @Dsl R.() -> Unit) {
+
+}
+
+fun foo1(i: Int) {
+
+}
+
+fun foo3() {
+
+}
+
+@Dsl
+fun R.foo2() {}
+
+@Dsl
+fun R.foo4() {
+
+}
+
+@Dsl
+fun R.fooloooooong() {
+
+}
+
+val R.fooval
+    get() = Unit
+
+@Dsl
+fun R.SomethingSomethingFooSomething() {
+
+}
+

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

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

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

@@ -69,6 +69,11 @@ public class BasicCompletionWeigherTestGenerated extends AbstractBasicCompletion
         runTest("idea/idea-completion/testData/weighers/basic/DslCalls.kt");
     }
 
+    @TestMetadata("DslCallsAnnotatedFunctionType.kt")
+    public void testDslCallsAnnotatedFunctionType() throws Exception {
+        runTest("idea/idea-completion/testData/weighers/basic/DslCallsAnnotatedFunctionType.kt");
+    }
+
     @TestMetadata("DslCallsWithMultipleReceivers.kt")
     public void testDslCallsWithMultipleReceivers() throws Exception {
         runTest("idea/idea-completion/testData/weighers/basic/DslCallsWithMultipleReceivers.kt");

idea/testData/quickfix/autoImports/dslMarkersOnReceiver.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/dslMarkersOnReceiver.before.Dependency.kt

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

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

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

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

@@ -459,6 +459,11 @@ public class QuickFixMultiFileTestGenerated extends AbstractQuickFixMultiFileTes
             runTest("idea/testData/quickfix/autoImports/dslMarkers.before.Main.kt");
         }
 
+        @TestMetadata("dslMarkersOnReceiver.before.Main.kt")
+        public void testDslMarkersOnReceiver() throws Exception {
+            runTest("idea/testData/quickfix/autoImports/dslMarkersOnReceiver.before.Main.kt");
+        }
+
         @TestMetadata("extensionFunctionImport.before.Main.kt")
         public void testExtensionFunctionImport() throws Exception {
             runTest("idea/testData/quickfix/autoImports/extensionFunctionImport.before.Main.kt");