[FIR] Fix collection of DSL markers from function types

When we're checking a context parameter, we only should collect the
DSL markers from the respective context parameter type of the function
type, and especially we shouldn't collect the DSL markers from the
receiver type.

#KT-77301 Fixed

Author: cypressious

compiler/fir/analysis-tests/testData/resolve/contextParameters/declarationAndUsages/dslMarkerKt77301.kt

@@ -1,4 +1,4 @@
-// RUN_PIPELINE_TILL: FRONTEND
+// RUN_PIPELINE_TILL: BACKEND
 // ISSUE: KT-77301
 // LANGUAGE: +ContextParameters
 @DslMarker
@@ -17,6 +17,6 @@ val pageContext: PageContext
 
 fun test() {
     postprocess {
-        <!DSL_SCOPE_VIOLATION!>pageContext<!>
+        pageContext
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/stages/ResolutionStages.kt

@@ -497,22 +497,33 @@ object CheckDslScopeViolation : ResolutionStage() {
         context: ResolutionContext,
     ): Set<ClassId> {
         return buildSet {
-            (boundSymbol.containingDeclarationIfParameter() as? FirAnonymousFunctionSymbol)?.fir?.matchingParameterFunctionType?.let {
-                // collect annotations in the function type at declaration site. For example, the `@A` and `@B` in the following code.
+            (boundSymbol.containingDeclarationIfParameter() as? FirAnonymousFunctionSymbol)?.let { anonymousFunctionSymbol ->
+                val matchingParameterFunctionType = anonymousFunctionSymbol.fir.matchingParameterFunctionType ?: return@let
+
+                // Collect annotations in the function type at declaration site. For example, the `@A`, `@B` and `@C in the following code.
                 // ```
-                // fun <T> body(block: @A ((@B T).() -> Unit)) { ... }
+                // fun <T, R> body(block: @A (context(@B T) (@C R).() -> Unit)) { ... }
                 // ```
+                // @A should be collected unconditionally.
+                // @B should only be collected if `boundSymbol` resolves to the respective context parameter of the anonymous function.
+                // @C should only be collected if `boundSymbol` resolves to the receiver parameter of the anonymous function.
 
                 // Collect the annotation on the function type, or `@A` in the example above.
-                collectDslMarkerAnnotations(context, it.customAnnotations)
+                collectDslMarkerAnnotations(context, matchingParameterFunctionType.customAnnotations)
 
-                // Collect the annotation on the extension receiver, or `@B` in the example above.
-                it.receiverType(context.session)?.let { receiverType ->
-                    collectDslMarkerAnnotations(context, receiverType)
+                // Collect the annotation on the context parameter, or `@B` in the example above.
+                if (boundSymbol is FirValueParameterSymbol) {
+                    val index = anonymousFunctionSymbol.contextParameterSymbols.indexOf(boundSymbol)
+                    matchingParameterFunctionType.contextParameterTypes(context.session).elementAtOrNull(index)?.let { contextType ->
+                        collectDslMarkerAnnotations(context, contextType)
+                    }
                 }
 
-                it.contextParameterTypes(context.session).forEach { contextType ->
-                    collectDslMarkerAnnotations(context, contextType)
+                // Collect the annotation on the extension receiver, or `@C` in the example above.
+                if (boundSymbol is FirReceiverParameterSymbol) {
+                    matchingParameterFunctionType.receiverType(context.session)?.let { receiverType ->
+                        collectDslMarkerAnnotations(context, receiverType)
+                    }
                 }
             }