[FIR] provide containerDeclaration to CustomAnnotationTypeAttribute for lazy resolve

tested by FirIdeNormalAnalysisSourceModuleSymbolByPsiTestGenerated.testTypeAnnotations

^KTIJ-23547 Fixed
^KTIJ-24141 Fixed

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/annotations/KtFirAnnotationListForType.kt

@@ -13,9 +13,13 @@ import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.hasAnnotation
 import org.jetbrains.kotlin.fir.declarations.toAnnotationClassId
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.custom
 import org.jetbrains.kotlin.fir.types.customAnnotations
 import org.jetbrains.kotlin.name.ClassId
 
@@ -25,32 +29,38 @@ internal class KtFirAnnotationListForType private constructor(
     override val token: KtLifetimeToken,
 ) : KtAnnotationsList() {
     override val annotations: List<KtAnnotationApplication>
-        get() = withValidityAssertion { coneType.customAnnotations.map { it.toKtAnnotationApplication(useSiteSession) } }
+        get() = withValidityAssertion {
+            coneType.customAnnotationsWithLazyResolve(FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING).map {
+                it.toKtAnnotationApplication(useSiteSession)
+            }
+        }
 
     override fun hasAnnotation(
         classId: ClassId,
         useSiteTarget: AnnotationUseSiteTarget?,
         acceptAnnotationsWithoutUseSite: Boolean,
     ): Boolean = withValidityAssertion {
-        coneType.customAnnotations.any {
+        coneType.customAnnotationsWithLazyResolve(FirResolvePhase.TYPES).any {
             (it.useSiteTarget == useSiteTarget || acceptAnnotationsWithoutUseSite && it.useSiteTarget == null) &&
                     it.toAnnotationClassId(useSiteSession) == classId
         }
     }
 
     override fun hasAnnotation(classId: ClassId): Boolean = withValidityAssertion {
-        coneType.customAnnotations.hasAnnotation(classId, useSiteSession)
+        coneType.customAnnotationsWithLazyResolve(FirResolvePhase.TYPES).hasAnnotation(classId, useSiteSession)
     }
 
     override fun annotationsByClassId(classId: ClassId): List<KtAnnotationApplication> = withValidityAssertion {
-        coneType.customAnnotations.mapNotNull { annotation ->
+        coneType.customAnnotationsWithLazyResolve(FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING).mapNotNull { annotation ->
             if (annotation.toAnnotationClassId(useSiteSession) != classId) return@mapNotNull null
             annotation.toKtAnnotationApplication(useSiteSession)
         }
     }
 
     override val annotationClassIds: Collection<ClassId>
-        get() = withValidityAssertion { coneType.customAnnotations.mapNotNull { it.toAnnotationClassId(useSiteSession) } }
+        get() = withValidityAssertion {
+            coneType.customAnnotationsWithLazyResolve(FirResolvePhase.TYPES).mapNotNull { it.toAnnotationClassId(useSiteSession) }
+        }
 
     companion object {
         fun create(
@@ -67,3 +77,10 @@ internal class KtFirAnnotationListForType private constructor(
     }
 }
 
+private fun ConeKotlinType.customAnnotationsWithLazyResolve(phase: FirResolvePhase): List<FirAnnotation> {
+    val custom = attributes.custom ?: return emptyList()
+    val annotations = custom.annotations.ifEmpty { return emptyList() }
+
+    custom.containerDeclaration?.lazyResolveToPhase(phase)
+    return annotations
+}

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.fir.declarations.FirDeclaration
-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.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