[SLC] LazyAnnotationsBox: rewrite from SymbolLightAbstractAnnotationWithClassId to PsiAnnotation

and update AdditionalAnnotationsProvider interface

^KT-56046

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/AdditionalAnnotationsProvider.kt

@@ -5,7 +5,10 @@
 
 package org.jetbrains.kotlin.light.classes.symbol.annotations
 
+import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiModifierList
+
 internal interface AdditionalAnnotationsProvider {
-    fun addAllAnnotations(currentRawAnnotations: MutableList<out SymbolLightAbstractAnnotationWithClassId>)
-    fun findAdditionalAnnotation(box: LazyAnnotationsBox, qualifiedName: String): SymbolLightAbstractAnnotationWithClassId?
+    fun addAllAnnotations(currentRawAnnotations: MutableList<in PsiAnnotation>, foundQualifiers: MutableSet<String>, owner: PsiModifierList)
+    fun findAdditionalAnnotation(annotationsBox: LazyAnnotationsBox, qualifiedName: String, owner: PsiModifierList): PsiAnnotation?
 }

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/LazyAnnotationsBox.kt

@@ -25,45 +25,43 @@ internal class LazyAnnotationsBox(
     private val annotatedSymbolPointer: KtSymbolPointer<KtAnnotatedSymbol>,
     private val ktModule: KtModule,
     private val owner: PsiModifierList,
-    private val additionalAnnotationsProviders: Collection<AdditionalAnnotationsProvider>,
+    private val additionalAnnotationsProvider: AdditionalAnnotationsProvider,
 ) : PsiAnnotationOwner {
     private inline fun <T> withAnnotatedSymbol(crossinline action: context(KtAnalysisSession) (KtAnnotatedSymbol) -> T): T =
         annotatedSymbolPointer.withSymbol(ktModule, action)
 
-    private val annotationsArray: AtomicReference<Array<SymbolLightAbstractAnnotationWithClassId>?> = AtomicReference()
-    private var specialAnnotations: SmartList<SymbolLightAbstractAnnotationWithClassId>? = null
+    private val annotationsArray: AtomicReference<Array<PsiAnnotation>?> = AtomicReference()
+    private var specialAnnotations: SmartList<PsiAnnotation>? = null
     private val monitor = Any()
 
-    override fun getAnnotations(): Array<SymbolLightAbstractAnnotationWithClassId> {
+    override fun getAnnotations(): Array<PsiAnnotation> {
         annotationsArray.get()?.let { return it }
 
         val classIds = withAnnotatedSymbol { ktAnnotatedSymbol ->
             ktAnnotatedSymbol.annotationClassIds
         }
 
-        val annotations = classIds.withIndex().mapTo(SmartList<SymbolLightAbstractAnnotationWithClassId>()) { (index, classId) ->
+        val annotations = classIds.withIndex().mapTo(SmartList<PsiAnnotation>()) { (index, classId) ->
             SymbolLightLazyAnnotation(classId, annotatedSymbolPointer, ktModule, index, owner)
         }
 
         val valueToReturn = if (annotations.isEmpty()) {
-            setAnnotationsArray(emptyArray)
+            setAnnotationsArray(PsiAnnotation.EMPTY_ARRAY)
         } else {
             synchronized(monitor) {
-                specialAnnotations?.forEach { lazyAnnotation ->
-                    val index = annotations.indexOfFirst { it.classId == lazyAnnotation.classId }
+                specialAnnotations?.forEach { specialAnnotation ->
+                    val index = annotations.indexOfFirst { it.qualifiedName == specialAnnotation.qualifiedName }
                     if (index != -1) {
-                        annotations[index] = lazyAnnotation
+                        annotations[index] = specialAnnotation
                     } else {
-                        annotations += lazyAnnotation
+                        annotations += specialAnnotation
                     }
                 }
 
-                for (annotationsProvider in additionalAnnotationsProviders) {
-                    annotationsProvider.addAllAnnotations(annotations)
-                }
+                val foundQualifiers = hashSetOf<String>()
+                additionalAnnotationsProvider.addAllAnnotations(annotations, foundQualifiers, owner)
 
                 specialAnnotations = null
-
                 setAnnotationsArray(annotations.toTypedArray())
             }
         }
@@ -71,16 +69,18 @@ internal class LazyAnnotationsBox(
         return valueToReturn
     }
 
-    private fun setAnnotationsArray(array: Array<SymbolLightAbstractAnnotationWithClassId>): Array<SymbolLightAbstractAnnotationWithClassId> =
+    private fun setAnnotationsArray(array: Array<PsiAnnotation>): Array<PsiAnnotation> =
         if (annotationsArray.compareAndSet(null, array)) {
             array
         } else {
             annotationsArray.get() ?: error("Unexpected state")
         }
 
-    override fun getApplicableAnnotations(): Array<SymbolLightAbstractAnnotationWithClassId> = annotations
+    override fun getApplicableAnnotations(): Array<PsiAnnotation> = annotations
 
-    override fun findAnnotation(qualifiedName: String): SymbolLightAbstractAnnotationWithClassId? {
+    override fun findAnnotation(qualifiedName: String): PsiAnnotation? = findAnnotation(qualifiedName, withAdditionalAnnotations = true)
+
+    fun findAnnotation(qualifiedName: String, withAdditionalAnnotations: Boolean): PsiAnnotation? {
         annotationsArray.get()?.let { array ->
             return array.find { it.qualifiedName == qualifiedName }
         }
@@ -98,8 +98,10 @@ internal class LazyAnnotationsBox(
                 annotationApplication,
                 owner,
             )
+        } else if (withAdditionalAnnotations) {
+            additionalAnnotationsProvider.findAdditionalAnnotation(this, qualifiedName, owner)
         } else {
-            additionalAnnotationsProviders.firstNotNullOfOrNull { it.findAdditionalAnnotation(this, qualifiedName) }
+            null
         }
 
         if (specialAnnotation == null) {
@@ -113,7 +115,7 @@ internal class LazyAnnotationsBox(
 
             if (specialAnnotations != null) {
                 val specialAnnotations = specialAnnotations!!
-                val oldAnnotation = specialAnnotations.find { it.classId == specialAnnotation.classId }
+                val oldAnnotation = specialAnnotations.find { it.qualifiedName == specialAnnotation.qualifiedName }
                 if (oldAnnotation != null) {
                     oldAnnotation
                 } else {
@@ -149,14 +151,15 @@ internal class LazyAnnotationsBox(
         private val specialAnnotationsList: Map<String, ClassId> = listOf(
             StandardClassIds.Annotations.Deprecated,
 //        Annotations.Retention,
+//        Annotations.Java.Retention,
 //        Annotations.Target,
+//        Annotations.Java.Target,
+//        Annotations.Java.Override,
             StandardClassIds.Annotations.DeprecatedSinceKotlin,
             StandardClassIds.Annotations.WasExperimental,
             StandardClassIds.Annotations.JvmRecord,
 //        Annotations.Repeatable,
 //        Annotations.Java.Repeatable,
         ).associateBy { it.asFqNameString() }
-
-        private val emptyArray: Array<SymbolLightAbstractAnnotationWithClassId> = emptyArray<SymbolLightAbstractAnnotationWithClassId>()
     }
 }

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/SymbolLightAbstractAnnotationWithClassId.kt

@@ -1,16 +0,0 @@
-/*
- * 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.light.classes.symbol.annotations
-
-import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.name.ClassId
-
-internal abstract class SymbolLightAbstractAnnotationWithClassId(
-    val classId: ClassId,
-    owner: PsiElement,
-) : SymbolLightAbstractAnnotation(owner) {
-    override fun getQualifiedName(): String = classId.asFqNameString()
-}

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/annotations/SymbolLightLazyAnnotation.kt

@@ -9,30 +9,35 @@ import com.intellij.psi.PsiAnnotationParameterList
 import com.intellij.psi.PsiModifierList
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationApplication
-import org.jetbrains.kotlin.analysis.api.annotations.annotations
+import org.jetbrains.kotlin.analysis.api.annotations.annotationsByClassId
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtAnnotatedSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.asJava.classes.lazyPub
 import org.jetbrains.kotlin.light.classes.symbol.withSymbol
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtCallElement
 
 internal class SymbolLightLazyAnnotation private constructor(
-    classId: ClassId,
+    private val classId: ClassId,
     private val annotatedSymbolPointer: KtSymbolPointer<KtAnnotatedSymbol>,
     private val ktModule: KtModule,
     private val index: Int?,
     specialAnnotationApplication: KtAnnotationApplication?,
     owner: PsiModifierList,
-) : SymbolLightAbstractAnnotationWithClassId(classId, owner) {
+) : SymbolLightAbstractAnnotation(owner) {
     init {
         require(index != null || specialAnnotationApplication != null)
     }
 
-    private val annotationApplication: Lazy<KtAnnotationApplication> = specialAnnotationApplication?.let(::lazyOf) ?: lazyPub {
+    private val fqName: FqName = classId.asSingleFqName()
+
+    val annotationApplication: Lazy<KtAnnotationApplication> = specialAnnotationApplication?.let(::lazyOf) ?: lazyPub {
         withAnnotatedSymbol { ktAnnotatedSymbol ->
-            ktAnnotatedSymbol.annotations[index!!]
+            val applications = ktAnnotatedSymbol.annotationsByClassId(classId)
+            applications.find { it.index == index }
+                ?: error("expected index: $index, actual indices: ${applications.map { it.index }}")
         }
     }
 
@@ -66,12 +71,12 @@ internal class SymbolLightLazyAnnotation private constructor(
         owner = owner,
     )
 
-    private inline fun <T> withAnnotatedSymbol(crossinline action: context(KtAnalysisSession) (KtAnnotatedSymbol) -> T): T =
+    inline fun <T> withAnnotatedSymbol(crossinline action: context(KtAnalysisSession) (KtAnnotatedSymbol) -> T): T =
         annotatedSymbolPointer.withSymbol(ktModule, action)
 
     override val kotlinOrigin: KtCallElement? get() = annotationApplication.value.psi
 
-    override fun getQualifiedName(): String = classId.asFqNameString()
+    override fun getQualifiedName(): String = fqName.asString()
 
     private val _parameterList: PsiAnnotationParameterList by lazyPub {
         SymbolLightLazyAnnotationParameterList(this, lazyPub { annotationApplication.value.arguments })
@@ -81,10 +86,10 @@ internal class SymbolLightLazyAnnotation private constructor(
 
     override fun equals(other: Any?): Boolean = this === other ||
             other is SymbolLightLazyAnnotation &&
-            other.classId == classId &&
+            other.fqName == fqName &&
             other.index == index &&
             other.ktModule == ktModule &&
             other.parent == parent
 
-    override fun hashCode(): Int = classId.hashCode()
+    override fun hashCode(): Int = fqName.hashCode()
 }