FIR: Use outer scope for annotation arguments resolution

^KT-41980 Fixed

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirRegularTowerDataContexts.kt

@@ -11,6 +11,7 @@ import java.util.*
 
 enum class FirTowerDataMode {
     REGULAR,
+    CLASS_HEADER_ANNOTATIONS,
     NESTED_CLASS,
     COMPANION_OBJECT,
     CONSTRUCTOR_HEADER,
@@ -26,6 +27,7 @@ class FirRegularTowerDataContexts private constructor(
 ) {
     constructor(
         regular: FirTowerDataContext,
+        forClassHeaderAnnotations: FirTowerDataContext? = null,
         forNestedClasses: FirTowerDataContext? = null,
         forCompanionObject: FirTowerDataContext? = null,
         forConstructorHeaders: FirTowerDataContext? = null,
@@ -33,7 +35,7 @@ class FirRegularTowerDataContexts private constructor(
         primaryConstructorPureParametersScope: FirLocalScope? = null,
         primaryConstructorAllParametersScope: FirLocalScope? = null,
     ) : this(
-        enumMap(regular, forNestedClasses, forCompanionObject, forConstructorHeaders, forEnumEntries),
+        enumMap(regular, forClassHeaderAnnotations, forNestedClasses, forCompanionObject, forConstructorHeaders, forEnumEntries),
         primaryConstructorPureParametersScope,
         primaryConstructorAllParametersScope,
         FirTowerDataMode.REGULAR
@@ -68,6 +70,7 @@ class FirRegularTowerDataContexts private constructor(
     companion object {
         private fun enumMap(
             regular: FirTowerDataContext,
+            forClassHeaderAnnotations: FirTowerDataContext?,
             forNestedClasses: FirTowerDataContext?,
             forCompanionObject: FirTowerDataContext?,
             forConstructorHeaders: FirTowerDataContext?,
@@ -75,6 +78,7 @@ class FirRegularTowerDataContexts private constructor(
         ): EnumMap<FirTowerDataMode, FirTowerDataContext> {
             val modeMap = EnumMap<FirTowerDataMode, FirTowerDataContext>(FirTowerDataMode::class.java)
             modeMap[FirTowerDataMode.REGULAR] = regular
+            modeMap[FirTowerDataMode.CLASS_HEADER_ANNOTATIONS] = forClassHeaderAnnotations
             modeMap[FirTowerDataMode.NESTED_CLASS] = forNestedClasses
             modeMap[FirTowerDataMode.COMPANION_OBJECT] = forCompanionObject
             modeMap[FirTowerDataMode.CONSTRUCTOR_HEADER] = forConstructorHeaders

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -478,6 +478,7 @@ class BodyResolveContext(
 
         val newContexts = FirRegularTowerDataContexts(
             forMembersResolution,
+            forClassHeaderAnnotations = base,
             newTowerDataContextForStaticNestedClasses,
             statics,
             scopeForConstructorHeader,
@@ -655,6 +656,16 @@ class BodyResolveContext(
         f: () -> T
     ): T = withTowerDataMode(FirTowerDataMode.ENUM_ENTRY, f)
 
+    @OptIn(PrivateForInline::class)
+    inline fun <T> forAnnotation(
+        f: () -> T
+    ): T {
+        return when (containerIfAny) {
+            is FirRegularClass -> withTowerDataMode(FirTowerDataMode.CLASS_HEADER_ANNOTATIONS, f)
+            else -> f()
+        }
+    }
+
     @OptIn(PrivateForInline::class)
     inline fun <T> withAnonymousInitializer(
         anonymousInitializer: FirAnonymousInitializer,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -685,12 +685,12 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
         dataFlowAnalyzer.enterFunction(constructor)
 
-        constructor.transformTypeParameters(transformer, data)
-            .transformAnnotations(transformer, data)
-            .transformReceiverTypeRef(transformer, data)
-            .transformReturnTypeRef(transformer, data)
-
         context.withConstructor(constructor) {
+            constructor.transformTypeParameters(transformer, data)
+                .transformAnnotations(transformer, data)
+                .transformReceiverTypeRef(transformer, data)
+                .transformReturnTypeRef(transformer, data)
+
             context.forConstructorParameters(constructor, owningClass, components) {
                 constructor.transformValueParameters(transformer, data)
             }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -1021,15 +1021,17 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
     override fun transformAnnotationCall(annotationCall: FirAnnotationCall, data: ResolutionMode): FirStatement {
         if (annotationCall.resolved) return annotationCall
         annotationCall.transformAnnotationTypeRef(transformer, ResolutionMode.ContextIndependent)
-        return withFirArrayOfCallTransformer {
-            dataFlowAnalyzer.enterAnnotation(annotationCall)
-            val result = callResolver.resolveAnnotationCall(annotationCall)
-            dataFlowAnalyzer.exitAnnotation(result ?: annotationCall)
-            if (result == null) return annotationCall
-            callCompleter.completeCall(result, noExpectedType)
-            // TODO: FirBlackBoxCodegenTestGenerated.Annotations.testDelegatedPropertySetter, it fails with hard cast
-            (result.argumentList as? FirResolvedArgumentList)?.let { annotationCall.replaceArgumentMapping((it).toAnnotationArgumentMapping()) }
-            annotationCall
+        return context.forAnnotation {
+            withFirArrayOfCallTransformer {
+                dataFlowAnalyzer.enterAnnotation(annotationCall)
+                val result = callResolver.resolveAnnotationCall(annotationCall)
+                dataFlowAnalyzer.exitAnnotation(result ?: annotationCall)
+                if (result == null) return annotationCall
+                callCompleter.completeCall(result, noExpectedType)
+                // TODO: FirBlackBoxCodegenTestGenerated.Annotations.testDelegatedPropertySetter, it fails with hard cast
+                (result.argumentList as? FirResolvedArgumentList)?.let { annotationCall.replaceArgumentMapping((it).toAnnotationArgumentMapping()) }
+                annotationCall
+            }
         }
     }
 

compiler/testData/diagnostics/tests/scopes/classHeader/annotationOnClass.fir.kt

@@ -15,15 +15,15 @@ annotation class Ann(
 )
 
 @Ann(
-        Nested::class,
-        Inner::class,
-        Interface::class,
-        CONST,
-        Companion.CONST,
-        Nested.CONST,
-        Interface.CONST,
-        <!ANNOTATION_ARGUMENT_MUST_BE_CONST!>a<!>,
-        b()
+        <!ANNOTATION_ARGUMENT_MUST_BE_CONST!><!UNRESOLVED_REFERENCE!>Nested<!>::class<!>,
+        <!ANNOTATION_ARGUMENT_MUST_BE_CONST!><!UNRESOLVED_REFERENCE!>Inner<!>::class<!>,
+        <!ANNOTATION_ARGUMENT_MUST_BE_CONST!><!UNRESOLVED_REFERENCE!>Interface<!>::class<!>,
+        <!UNRESOLVED_REFERENCE!>CONST<!>,
+        <!UNRESOLVED_REFERENCE!>Companion<!>.CONST,
+        <!UNRESOLVED_REFERENCE!>Nested<!>.CONST,
+        <!UNRESOLVED_REFERENCE!>Interface<!>.CONST,
+        <!UNRESOLVED_REFERENCE!>a<!>,
+        <!UNRESOLVED_REFERENCE!>b<!>()
 )
 class A {