[AA FIR] avoid redundant resolve by annotations in IDE mode

^KT-56046

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

@@ -7,14 +7,11 @@ package org.jetbrains.kotlin.analysis.api.fir.annotations
 
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationApplication
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationsList
-import org.jetbrains.kotlin.analysis.api.fir.toKtAnnotationApplication
 import org.jetbrains.kotlin.analysis.api.impl.base.annotations.KtEmptyAnnotationsList
 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.hasAnnotation
-import org.jetbrains.kotlin.fir.declarations.toAnnotationClassId
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.name.ClassId
 
@@ -25,10 +22,7 @@ internal class KtFirAnnotationListForDeclaration private constructor(
 ) : KtAnnotationsList() {
     override val annotations: List<KtAnnotationApplication>
         get() = withValidityAssertion {
-            firSymbol.resolvedAnnotationsWithArguments
+            annotations(firSymbol, useSiteSession)
-                .map { annotation ->
-                    annotation.toKtAnnotationApplication(useSiteSession)
-                }
         }
 
     override fun hasAnnotation(
@@ -36,26 +30,21 @@ internal class KtFirAnnotationListForDeclaration private constructor(
         useSiteTarget: AnnotationUseSiteTarget?,
         acceptAnnotationsWithoutUseSite: Boolean,
     ): Boolean = withValidityAssertion {
-        firSymbol.resolvedAnnotationsWithClassIds.any {
+        hasAnnotation(firSymbol, classId, useSiteSession, useSiteTarget, acceptAnnotationsWithoutUseSite)
-            (it.useSiteTarget == useSiteTarget || acceptAnnotationsWithoutUseSite && it.useSiteTarget == null) &&
-                    it.toAnnotationClassId(useSiteSession) == classId
-        }
     }
 
     override fun hasAnnotation(classId: ClassId): Boolean = withValidityAssertion {
-        firSymbol.resolvedAnnotationsWithClassIds.hasAnnotation(classId, useSiteSession)
+        hasAnnotation(firSymbol, classId, useSiteSession)
     }
 
     override fun annotationsByClassId(classId: ClassId): List<KtAnnotationApplication> = withValidityAssertion {
-        firSymbol.resolvedAnnotationsWithArguments.mapNotNull { annotation ->
+        annotationsByClassId(firSymbol, classId, useSiteSession)
-            if (annotation.toAnnotationClassId(useSiteSession) != classId) return@mapNotNull null
-            annotation.toKtAnnotationApplication(useSiteSession)
-        }
     }
 
-
     override val annotationClassIds: Collection<ClassId>
-        get() = withValidityAssertion { firSymbol.resolvedAnnotationClassIds }
+        get() = withValidityAssertion {
+            annotationClassIds(firSymbol)
+        }
 
     companion object {
         fun create(

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

@@ -7,19 +7,13 @@ package org.jetbrains.kotlin.analysis.api.fir.annotations
 
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationApplication
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationsList
-import org.jetbrains.kotlin.analysis.api.fir.toKtAnnotationApplication
 import org.jetbrains.kotlin.analysis.api.impl.base.annotations.KtEmptyAnnotationsList
 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.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.hasAnnotation
-import org.jetbrains.kotlin.fir.declarations.toAnnotationClassId
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
-import org.jetbrains.kotlin.fir.symbols.resolvedAnnotationClassIds
-import org.jetbrains.kotlin.fir.symbols.resolvedAnnotationsWithArguments
-import org.jetbrains.kotlin.fir.symbols.resolvedAnnotationsWithClassIds
 import org.jetbrains.kotlin.name.ClassId
 
 internal class KtFirAnnotationListForReceiverParameter private constructor(
@@ -31,9 +25,7 @@ internal class KtFirAnnotationListForReceiverParameter private constructor(
 
     override val annotations: List<KtAnnotationApplication>
         get() = withValidityAssertion {
-            receiverParameter.resolvedAnnotationsWithArguments(firCallableSymbol).map { annotation ->
+            annotations(firCallableSymbol, useSiteSession, receiverParameter)
-                annotation.toKtAnnotationApplication(useSiteSession)
-            }
         }
 
     override fun hasAnnotation(
@@ -41,26 +33,20 @@ internal class KtFirAnnotationListForReceiverParameter private constructor(
         useSiteTarget: AnnotationUseSiteTarget?,
         acceptAnnotationsWithoutUseSite: Boolean,
     ): Boolean = withValidityAssertion {
-        receiverParameter.resolvedAnnotationsWithClassIds(firCallableSymbol).any {
+        hasAnnotation(firCallableSymbol, classId, useSiteSession, useSiteTarget, acceptAnnotationsWithoutUseSite, receiverParameter)
-            (it.useSiteTarget == useSiteTarget || acceptAnnotationsWithoutUseSite && it.useSiteTarget == null) &&
-                    it.toAnnotationClassId(useSiteSession) == classId
-        }
     }
 
     override fun hasAnnotation(classId: ClassId): Boolean = withValidityAssertion {
-        receiverParameter.resolvedAnnotationsWithClassIds(firCallableSymbol).hasAnnotation(classId, useSiteSession)
+        hasAnnotation(firCallableSymbol, classId, useSiteSession, receiverParameter)
     }
 
     override fun annotationsByClassId(classId: ClassId): List<KtAnnotationApplication> = withValidityAssertion {
-        receiverParameter.resolvedAnnotationsWithArguments(firCallableSymbol).mapNotNull { annotation ->
+        annotationsByClassId(firCallableSymbol, classId, useSiteSession, receiverParameter)
-            if (annotation.toAnnotationClassId(useSiteSession) != classId) return@mapNotNull null
-            annotation.toKtAnnotationApplication(useSiteSession)
-        }
     }
 
     override val annotationClassIds: Collection<ClassId>
         get() = withValidityAssertion {
-            receiverParameter.resolvedAnnotationClassIds(firCallableSymbol)
+            annotationClassIds(firCallableSymbol, receiverParameter)
         }
 
     companion object {

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

@@ -17,7 +17,6 @@ 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
@@ -82,7 +81,11 @@ private fun ConeKotlinType.customAnnotationsWithLazyResolve(phase: FirResolvePha
     val annotations = custom.annotations.ifEmpty { return emptyList() }
 
     for (containerSymbol in custom.containerSymbols) {
-        containerSymbol.lazyResolveToPhase(phase)
+        when (phase) {
+            FirResolvePhase.TYPES -> containerSymbol.resolvedAnnotationsWithClassIds
+            FirResolvePhase.ANNOTATIONS_ARGUMENTS_MAPPING -> containerSymbol.resolvedAnnotationsWithArguments
+            else -> {}
+        }
     }
 
     return annotations

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

@@ -5,12 +5,20 @@
 
 package org.jetbrains.kotlin.analysis.api.fir.annotations
 
+import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationApplication
+import org.jetbrains.kotlin.analysis.api.fir.toKtAnnotationApplication
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withFirEntry
 import org.jetbrains.kotlin.analysis.utils.errors.checkWithAttachmentBuilder
 import org.jetbrains.kotlin.analysis.utils.errors.withClassEntry
-import org.jetbrains.kotlin.fir.declarations.resolved
+import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.CompilerRequiredAnnotationsHelper
+import org.jetbrains.kotlin.fir.symbols.*
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
 internal fun mapAnnotationParameters(annotation: FirAnnotation): Map<Name, FirExpression> {
@@ -22,3 +30,69 @@ internal fun mapAnnotationParameters(annotation: FirAnnotation): Map<Name, FirEx
 
     return annotation.argumentMapping.mapping.mapKeys { (name, _) -> name }
 }
+
+internal fun annotationsByClassId(
+    firSymbol: FirBasedSymbol<*>,
+    classId: ClassId,
+    useSiteSession: FirSession,
+    annotationContainer: FirAnnotationContainer = firSymbol.fir,
+): List<KtAnnotationApplication> =
+    if (firSymbol.isFromCompilerRequiredAnnotationsPhase(classId)) {
+        annotationContainer.resolvedCompilerRequiredAnnotations(firSymbol).mapNotNull { annotation ->
+            if (annotation.toAnnotationClassIdSafe(useSiteSession) != classId) return@mapNotNull null
+            annotation.toKtAnnotationApplication(useSiteSession)
+        }
+    } else {
+        annotationContainer.resolvedAnnotationsWithArguments(firSymbol).mapNotNull { annotation ->
+            if (annotation.toAnnotationClassId(useSiteSession) != classId) return@mapNotNull null
+            annotation.toKtAnnotationApplication(useSiteSession)
+        }
+    }
+
+internal fun annotations(
+    firSymbol: FirBasedSymbol<*>,
+    useSiteSession: FirSession,
+    annotationContainer: FirAnnotationContainer = firSymbol.fir,
+): List<KtAnnotationApplication> = annotationContainer.resolvedAnnotationsWithArguments(firSymbol).map { annotation ->
+    annotation.toKtAnnotationApplication(useSiteSession)
+}
+
+internal fun annotationClassIds(
+    firSymbol: FirBasedSymbol<*>,
+    annotationContainer: FirAnnotationContainer = firSymbol.fir,
+): Collection<ClassId> = annotationContainer.resolvedAnnotationClassIds(firSymbol)
+
+internal fun hasAnnotation(
+    firSymbol: FirBasedSymbol<*>,
+    classId: ClassId,
+    useSiteSession: FirSession,
+    useSiteTarget: AnnotationUseSiteTarget?,
+    acceptAnnotationsWithoutUseSite: Boolean,
+    annotationContainer: FirAnnotationContainer = firSymbol.fir,
+): Boolean =
+    if (firSymbol.isFromCompilerRequiredAnnotationsPhase(classId)) {
+        annotationContainer.resolvedCompilerRequiredAnnotations(firSymbol).any {
+            (it.useSiteTarget == useSiteTarget || acceptAnnotationsWithoutUseSite && it.useSiteTarget == null) &&
+                    it.toAnnotationClassIdSafe(useSiteSession) == classId
+        }
+    } else {
+        annotationContainer.resolvedAnnotationsWithClassIds(firSymbol).any {
+            (it.useSiteTarget == useSiteTarget || acceptAnnotationsWithoutUseSite && it.useSiteTarget == null) &&
+                    it.toAnnotationClassId(useSiteSession) == classId
+        }
+    }
+
+internal fun hasAnnotation(
+    firSymbol: FirBasedSymbol<*>,
+    classId: ClassId,
+    useSiteSession: FirSession,
+    annotationContainer: FirAnnotationContainer = firSymbol.fir,
+): Boolean =
+    if (firSymbol.isFromCompilerRequiredAnnotationsPhase(classId)) {
+        annotationContainer.resolvedCompilerRequiredAnnotations(firSymbol).hasAnnotationSafe(classId, useSiteSession)
+    } else {
+        annotationContainer.resolvedAnnotationsWithClassIds(firSymbol).hasAnnotation(classId, useSiteSession)
+    }
+
+private fun FirBasedSymbol<*>.isFromCompilerRequiredAnnotationsPhase(classId: ClassId): Boolean =
+    fir.resolvePhase < FirResolvePhase.TYPES && classId in CompilerRequiredAnnotationsHelper.REQUIRED_ANNOTATIONS

compiler/fir/providers/src/org/jetbrains/kotlin/fir/declarations/FirAnnotationUtils.kt

@@ -11,9 +11,9 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.toSymbol
-import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
@@ -36,7 +36,7 @@ private fun FirAnnotation.toAnnotationLookupTagSafe(session: FirSession): ConeCl
 fun FirAnnotation.toAnnotationClassId(session: FirSession): ClassId? =
     toAnnotationLookupTag(session)?.classId
 
-private fun FirAnnotation.toAnnotationClassIdSafe(session: FirSession): ClassId? =
+fun FirAnnotation.toAnnotationClassIdSafe(session: FirSession): ClassId? =
     toAnnotationLookupTagSafe(session)?.classId
 
 fun FirAnnotation.toAnnotationClassLikeSymbol(session: FirSession): FirClassLikeSymbol<*>? =

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/AbstractFirSpecificAnnotationResolveTransformer.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.
  */
 
@@ -22,6 +22,8 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.DesignationState
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSpecificTypeResolverTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.ScopeClassDeclaration
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.CompilerRequiredAnnotationsHelper.REQUIRED_ANNOTATIONS
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.CompilerRequiredAnnotationsHelper.REQUIRED_ANNOTATIONS_WITH_ARGUMENTS
 import org.jetbrains.kotlin.fir.resolve.transformers.withClassDeclarationCleanup
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.createImportingScopes
@@ -43,6 +45,22 @@ import org.jetbrains.kotlin.name.StandardClassIds.Annotations.JvmRecord
 import org.jetbrains.kotlin.name.StandardClassIds.Annotations.WasExperimental
 import org.jetbrains.kotlin.utils.addToStdlib.shouldNotBeCalled
 
+object CompilerRequiredAnnotationsHelper {
+    internal val REQUIRED_ANNOTATIONS_WITH_ARGUMENTS: Set<ClassId> = setOf(
+        Deprecated,
+        Annotations.Retention,
+        Annotations.Target,
+    )
+
+    val REQUIRED_ANNOTATIONS: Set<ClassId> = REQUIRED_ANNOTATIONS_WITH_ARGUMENTS + setOf(
+        DeprecatedSinceKotlin,
+        WasExperimental,
+        JvmRecord,
+        Annotations.Repeatable,
+        Annotations.Java.Repeatable,
+    )
+}
+
 internal abstract class AbstractFirSpecificAnnotationResolveTransformer(
     protected val session: FirSession,
     protected val scopeSession: ScopeSession,
@@ -50,20 +68,6 @@ internal abstract class AbstractFirSpecificAnnotationResolveTransformer(
     containingDeclarations: List<FirDeclaration> = emptyList()
 ) : FirDefaultTransformer<Nothing?>() {
     companion object {
-        private val REQUIRED_ANNOTATIONS_WITH_ARGUMENTS: Set<ClassId> = setOf(
-            Deprecated,
-            Annotations.Retention,
-            Annotations.Target,
-        )
-
-        private val REQUIRED_ANNOTATIONS: Set<ClassId> = REQUIRED_ANNOTATIONS_WITH_ARGUMENTS + setOf(
-            DeprecatedSinceKotlin,
-            WasExperimental,
-            JvmRecord,
-            Annotations.Repeatable,
-            Annotations.Java.Repeatable,
-        )
-
         private val REQUIRED_ANNOTATION_NAMES: Set<Name> = REQUIRED_ANNOTATIONS.mapTo(mutableSetOf()) { it.shortClassName }
     }