[FIR] provide containerDeclaration to CustomAnnotationTypeAttribute for lazy resolve

tested by FirIdeNormalAnalysisSourceModuleSymbolByPsiTestGenerated.testTypeAnnotations

^KTIJ-23547 Fixed
^KTIJ-24141 Fixed

compiler/fir/providers/src/org/jetbrains/kotlin/fir/CopyUtils.kt

@@ -1,27 +1,19 @@
 /*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
 package org.jetbrains.kotlin.fir
 
 import org.jetbrains.kotlin.builtins.StandardNames
-import org.jetbrains.kotlin.descriptors.EffectiveVisibility
-import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.descriptors.Visibility
-import org.jetbrains.kotlin.fir.declarations.FirDeclarationStatus
-import org.jetbrains.kotlin.fir.declarations.FirResolvedDeclarationStatus
-import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
-import org.jetbrains.kotlin.fir.declarations.builder.buildTypeParameter
-import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
-import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.utils.expandedConeType
 import org.jetbrains.kotlin.fir.diagnostics.ConeDiagnostic
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.builder.*
+import org.jetbrains.kotlin.fir.expressions.builder.FirImplicitInvokeCallBuilder
+import org.jetbrains.kotlin.fir.expressions.builder.buildImplicitInvokeCall
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.typeAttributeExtensions
-import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.resolve.directExpansionType
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.types.*
@@ -80,7 +72,11 @@ fun FirTypeRef.errorTypeFromPrototype(
     }
 }
 
-fun List<FirAnnotation>.computeTypeAttributes(session: FirSession, predefined: List<ConeAttribute<*>> = emptyList()): ConeAttributes {
+fun List<FirAnnotation>.computeTypeAttributes(
+    session: FirSession,
+    predefined: List<ConeAttribute<*>> = emptyList(),
+    containerDeclaration: FirDeclaration? = null,
+): ConeAttributes {
     if (this.isEmpty()) {
         if (predefined.isEmpty()) return ConeAttributes.Empty
         return ConeAttributes.create(predefined)
@@ -113,7 +109,7 @@ fun List<FirAnnotation>.computeTypeAttributes(session: FirSession, predefined: L
         }
     }
     if (customAnnotations.isNotEmpty()) {
-        attributes += CustomAnnotationTypeAttribute(customAnnotations)
+        attributes += CustomAnnotationTypeAttribute(customAnnotations, containerDeclaration)
     }
     return ConeAttributes.create(attributes)
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirTypeResolverImpl.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -362,7 +362,7 @@ class FirTypeResolverImpl(private val session: FirSession) : FirTypeResolver() {
         return symbol.constructType(
             resultingArguments,
             typeRef.isMarkedNullable,
-            typeRef.annotations.computeTypeAttributes(session)
+            typeRef.annotations.computeTypeAttributes(session, containerDeclaration = topContainer)
         ).also {
             val lookupTag = it.lookupTag
             if (lookupTag is ConeClassLikeLookupTagImpl && symbol is FirClassLikeSymbol<*>) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirTypeResolveTransformer.kt

@@ -60,6 +60,17 @@ open class FirTypeResolveTransformer(
 ) : FirAbstractTreeTransformer<Any?>(FirResolvePhase.TYPES) {
     private val scopes = mutableListOf<FirScope>()
     private val towerScope = scopes.asReversed()
+    private var currentDeclaration: FirDeclaration? = null
+
+    private inline fun <T> withDeclaration(declaration: FirDeclaration, crossinline action: () -> T): T {
+        val oldDeclaration = currentDeclaration
+        return try {
+            currentDeclaration = declaration
+            action()
+        } finally {
+            currentDeclaration = oldDeclaration
+        }
+    }
 
     init {
         scopes.addAll(initialScopes.asReversed())
@@ -208,11 +219,15 @@ open class FirTypeResolveTransformer(
         return implicitTypeRef
     }
 
+    override fun transformDeclaration(declaration: FirDeclaration, data: Any?): FirDeclaration = withDeclaration(declaration) {
+        super.transformDeclaration(declaration, data)
+    }
+
     override fun transformTypeRef(typeRef: FirTypeRef, data: Any?): FirResolvedTypeRef {
         return typeResolverTransformer.withFile(currentFile) {
             typeRef.transform(
                 typeResolverTransformer,
-                ScopeClassDeclaration(towerScope, classDeclarationsStack)
+                ScopeClassDeclaration(towerScope, classDeclarationsStack, currentDeclaration)
             )
         }
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/CustomAnnotationTypeAttribute.kt

@@ -5,18 +5,22 @@
 
 package org.jetbrains.kotlin.fir.types
 
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.render
 import kotlin.reflect.KClass
 
-class CustomAnnotationTypeAttribute(val annotations: List<FirAnnotation>) : ConeAttribute<CustomAnnotationTypeAttribute>() {
+class CustomAnnotationTypeAttribute(
+    val annotations: List<FirAnnotation>,
+    val containerDeclaration: FirDeclaration? = null,
+) : ConeAttribute<CustomAnnotationTypeAttribute>() {
     override fun union(other: CustomAnnotationTypeAttribute?): CustomAnnotationTypeAttribute? = null
 
     override fun intersect(other: CustomAnnotationTypeAttribute?): CustomAnnotationTypeAttribute? = null
 
     override fun add(other: CustomAnnotationTypeAttribute?): CustomAnnotationTypeAttribute {
         if (other == null || other === this) return this
-        return CustomAnnotationTypeAttribute(annotations + other.annotations)
+        return CustomAnnotationTypeAttribute(annotations + other.annotations, containerDeclaration)
     }
 
     override fun isSubtypeOf(other: CustomAnnotationTypeAttribute?): Boolean = true