defer resolution of annotations when possible.

neetopia · neetopia · commit c28c3573a179 · 2024-07-09T13:06:36.000-07:00
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/AbstractKSDeclarationImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/AbstractKSDeclarationImpl.kt
@@ -35,7 +35,7 @@ import com.intellij.psi.PsiModifierListOwner
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtNamedSymbol
 import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
-import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtModifierListOwner
 
 abstract class AbstractKSDeclarationImpl(val ktDeclarationSymbol: KtDeclarationSymbol) : KSDeclaration, Deferrable {
@@ -119,8 +119,10 @@ abstract class AbstractKSDeclarationImpl(val ktDeclarationSymbol: KtDeclarationS
     override val docString: String?
         get() = ktDeclarationSymbol.toDocString()
 
-    internal val originalAnnotations: Sequence<KSAnnotation> by lazy {
-        if (ktDeclarationSymbol.psi is KtElement || ktDeclarationSymbol.psi == null) {
+    internal open val originalAnnotations: Sequence<KSAnnotation> by lazy {
+        if (ktDeclarationSymbol.psi is KtAnnotated) {
+            (ktDeclarationSymbol.psi as KtAnnotated).annotations(ktDeclarationSymbol, this)
+        } else if (ktDeclarationSymbol.psi == null) {
             ktDeclarationSymbol.annotations(this)
         } else {
             (ktDeclarationSymbol.psi as PsiJvmModifiersOwner)
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSAnnotationImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSAnnotationImpl.kt
@@ -23,43 +23,47 @@ import com.google.devtools.ksp.common.impl.KSNameImpl
 import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.impl.symbol.java.KSValueArgumentLiteImpl
 import com.google.devtools.ksp.impl.symbol.java.calcValue
-import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolvedImpl
 import com.google.devtools.ksp.symbol.*
 import com.intellij.psi.PsiAnnotationMethod
 import com.intellij.psi.PsiArrayInitializerMemberValue
 import com.intellij.psi.PsiClass
 import com.intellij.psi.impl.compiled.ClsClassImpl
-import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
+import org.jetbrains.kotlin.analysis.api.KaAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationApplicationWithArgumentsInfo
 import org.jetbrains.kotlin.analysis.api.annotations.KtNamedAnnotationValue
 import org.jetbrains.kotlin.analysis.api.annotations.KtUnsupportedAnnotationValue
-import org.jetbrains.kotlin.analysis.api.components.buildClassType
 import org.jetbrains.kotlin.analysis.api.platform.lifetime.KotlinAlwaysAccessibleLifetimeToken
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget.*
 import org.jetbrains.kotlin.psi.KtAnnotationEntry
 
 // TODO: implement a psi based version of annotation application.
 class KSAnnotationImpl private constructor(
-    private val annotationApplication: KtAnnotationApplicationWithArgumentsInfo,
-    override val parent: KSNode?
+    private val ktAnnotationEntry: KtAnnotationEntry,
+    override val parent: KSNode?,
+    private val resolveToAnnotationApplication: () -> KtAnnotationApplicationWithArgumentsInfo
 ) : KSAnnotation {
-    companion object : KSObjectCache<IdKeyPair<KtAnnotationApplicationWithArgumentsInfo, KSNode?>, KSAnnotationImpl>() {
-        fun getCached(annotationApplication: KtAnnotationApplicationWithArgumentsInfo, parent: KSNode? = null) =
-            cache.getOrPut(IdKeyPair(annotationApplication, parent)) { KSAnnotationImpl(annotationApplication, parent) }
+    companion object : KSObjectCache<IdKeyPair<KtAnnotationEntry, KSNode?>, KSAnnotationImpl>() {
+        fun getCached(
+            ktAnnotationEntry: KtAnnotationEntry,
+            parent: KSNode? = null,
+            resolveToAnnotationApplication: () -> KtAnnotationApplicationWithArgumentsInfo
+        ) =
+            cache.getOrPut(IdKeyPair(ktAnnotationEntry, parent)) {
+                KSAnnotationImpl(ktAnnotationEntry, parent, resolveToAnnotationApplication)
+            }
+    }
+
+    private val annotationApplication by lazy {
+        resolveToAnnotationApplication()
     }
 
     override val annotationType: KSTypeReference by lazy {
         analyze {
-            if (annotationApplication.psi is KtAnnotationEntry) {
+            ktAnnotationEntry.typeReference!!.let {
                 KSTypeReferenceImpl.getCached(
-                    (annotationApplication.psi as KtAnnotationEntry).typeReference!!,
-                    parent = this@KSAnnotationImpl
-                )
-            } else {
-                KSTypeReferenceResolvedImpl.getCached(
-                    buildClassType(annotationApplication.classId!!),
-                    parent = this@KSAnnotationImpl
+                    it,
+                    this@KSAnnotationImpl
                 )
             }
         }
@@ -74,7 +78,7 @@ class KSAnnotationImpl private constructor(
         presentArgs + absentArgs
     }
 
-    @OptIn(KtAnalysisApiInternals::class)
+    @OptIn(KaAnalysisApiInternals::class)
     override val defaultArguments: List<KSValueArgument> by lazy {
         analyze {
             annotationApplication.classId?.toKtClassSymbol()?.let { symbol ->
@@ -120,11 +124,11 @@ class KSAnnotationImpl private constructor(
     }
 
     override val shortName: KSName by lazy {
-        KSNameImpl.getCached(annotationApplication.classId!!.shortClassName.asString())
+        KSNameImpl.getCached(ktAnnotationEntry.shortName!!.asString())
     }
 
     override val useSiteTarget: AnnotationUseSiteTarget? by lazy {
-        when (annotationApplication.useSiteTarget) {
+        when (ktAnnotationEntry.useSiteTarget?.getAnnotationUseSiteTarget()) {
             null -> null
             FILE -> AnnotationUseSiteTarget.FILE
             PROPERTY -> AnnotationUseSiteTarget.PROPERTY
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFileImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFileImpl.kt
@@ -85,7 +85,7 @@ class KSFileImpl private constructor(internal val ktFileSymbol: KtFileSymbol) :
     }
 
     override val annotations: Sequence<KSAnnotation> by lazy {
-        ktFileSymbol.annotations(this)
+        (ktFileSymbol.psi as? KtFile)?.annotations(ktFileSymbol) ?: ktFileSymbol.annotations(this)
     }
 
     override fun toString(): String {
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt
@@ -18,6 +18,7 @@
 package com.google.devtools.ksp.impl.symbol.kotlin
 
 import com.google.devtools.ksp.common.KSObjectCache
+import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSAnnotationResolvedImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolvedImpl
 import com.google.devtools.ksp.impl.symbol.util.toKSModifiers
 import com.google.devtools.ksp.symbol.*
@@ -37,13 +38,15 @@ abstract class KSPropertyAccessorImpl(
 ) : KSPropertyAccessor, Deferrable {
 
     override val annotations: Sequence<KSAnnotation> by lazy {
+        // (ktPropertyAccessorSymbol.psi as? KtPropertyAccessor)?.annotations(ktPropertyAccessorSymbol, this) ?:
         ktPropertyAccessorSymbol.annotations.asSequence()
             .filter { it.useSiteTarget != AnnotationUseSiteTarget.SETTER_PARAMETER }
-            .map { KSAnnotationImpl.getCached(it, this) }
+            .map { KSAnnotationResolvedImpl.getCached(it, this) }
             .plus(findAnnotationFromUseSiteTarget())
     }
 
     internal val originalAnnotations: Sequence<KSAnnotation> by lazy {
+        // (ktPropertyAccessorSymbol.psi as? KtPropertyAccessor)?.annotations(ktPropertyAccessorSymbol, this) ?:
         ktPropertyAccessorSymbol.annotations(this)
     }
 
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt
@@ -25,13 +25,13 @@ import com.google.devtools.ksp.common.impl.KSNameImpl
 import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.impl.recordLookupForPropertyOrMethod
 import com.google.devtools.ksp.impl.recordLookupWithSupertypes
+import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSAnnotationResolvedImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeReferenceResolvedImpl
 import com.google.devtools.ksp.impl.symbol.util.BinaryClassInfoCache
 import com.google.devtools.ksp.symbol.*
 import com.intellij.psi.PsiClass
 import org.jetbrains.kotlin.analysis.api.KaConstantInitializerValue
 import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationApplication
-import org.jetbrains.kotlin.analysis.api.annotations.annotations
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KaFirKotlinPropertySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtKotlinPropertySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtPropertySymbol
@@ -42,6 +42,7 @@ import org.jetbrains.kotlin.descriptors.java.JavaVisibilities
 import org.jetbrains.kotlin.load.java.structure.impl.JavaClassImpl
 import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement
+import org.jetbrains.kotlin.psi.KtAnnotationEntry
 import org.jetbrains.kotlin.psi.KtProperty
 
 class KSPropertyDeclarationImpl private constructor(internal val ktPropertySymbol: KtPropertySymbol) :
@@ -53,10 +54,13 @@ class KSPropertyDeclarationImpl private constructor(internal val ktPropertySymbo
             cache.getOrPut(ktPropertySymbol) { KSPropertyDeclarationImpl(ktPropertySymbol) }
     }
 
+    override val originalAnnotations: Sequence<KSAnnotation>
+        get() = annotations
+
     override val annotations: Sequence<KSAnnotation> by lazy {
         ktPropertySymbol.annotations.asSequence()
             .filter { !it.isUseSiteTargetAnnotation() }
-            .map { KSAnnotationImpl.getCached(it, this) }
+            .map { KSAnnotationResolvedImpl.getCached(it, this) }
             .plus(
                 if (ktPropertySymbol.isFromPrimaryConstructor) {
                     (parentDeclaration as? KSClassDeclaration)?.primaryConstructor?.parameters
@@ -74,8 +78,9 @@ class KSPropertyDeclarationImpl private constructor(internal val ktPropertySymbo
                     } ?: false
                 }
             }.plus(
+                // TODO: optimize for psi
                 ktPropertySymbol.backingFieldSymbol?.annotations
-                    ?.map { KSAnnotationImpl.getCached(it) } ?: emptyList()
+                    ?.map { KSAnnotationResolvedImpl.getCached(it, this@KSPropertyDeclarationImpl) } ?: emptyList()
             )
     }
 
@@ -181,7 +186,15 @@ internal fun KtAnnotationApplication.isUseSiteTargetAnnotation(): Boolean {
             it == AnnotationUseSiteTarget.CONSTRUCTOR_PARAMETER
     } ?: false
 }
-
+internal fun KtAnnotationEntry.isUseSiteTargetAnnotation(): Boolean {
+    return this.useSiteTarget?.getAnnotationUseSiteTarget()?.let {
+        it == AnnotationUseSiteTarget.PROPERTY_GETTER ||
+            it == AnnotationUseSiteTarget.PROPERTY_SETTER ||
+            it == AnnotationUseSiteTarget.SETTER_PARAMETER ||
+            it == AnnotationUseSiteTarget.CONSTRUCTOR_PARAMETER ||
+            it == AnnotationUseSiteTarget.FIELD
+    } ?: false
+}
 internal fun KtPropertySymbol.toModifiers(): Set<Modifier> {
     val result = mutableSetOf<Modifier>()
     if (visibility != JavaVisibilities.PackageVisibility) {
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeImpl.kt
@@ -22,6 +22,7 @@ import com.google.devtools.ksp.common.KSObjectCache
 import com.google.devtools.ksp.common.errorTypeOnInconsistentArguments
 import com.google.devtools.ksp.impl.ResolverAAImpl
 import com.google.devtools.ksp.impl.recordLookupWithSupertypes
+import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSAnnotationResolvedImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.resolved.KSTypeArgumentResolvedImpl
 import com.google.devtools.ksp.impl.symbol.kotlin.synthetic.getExtensionFunctionTypeAnnotation
 import com.google.devtools.ksp.symbol.KSAnnotation
@@ -88,9 +89,11 @@ class KSTypeImpl private constructor(internal val type: KtType) : KSType {
     override val annotations: Sequence<KSAnnotation>
         get() = type.annotations() +
             if (type is KtFunctionalType && type.receiverType != null) {
-                sequenceOf(KSAnnotationImpl.getCached(getExtensionFunctionTypeAnnotation(type.annotations.size)))
+                sequenceOf(
+                    KSAnnotationResolvedImpl.getCached(getExtensionFunctionTypeAnnotation(type.annotations.size))
+                )
             } else {
-                emptySequence<KSAnnotation>()
+                emptySequence()
             }
 
     override fun isAssignableFrom(that: KSType): Boolean {
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt
@@ -30,6 +30,7 @@ import com.google.devtools.ksp.symbol.Location
 import com.google.devtools.ksp.symbol.Modifier
 import com.google.devtools.ksp.symbol.Origin
 import org.jetbrains.kotlin.analysis.api.types.KtType
+import org.jetbrains.kotlin.psi.KtAnnotationEntry
 import org.jetbrains.kotlin.psi.KtDynamicType
 import org.jetbrains.kotlin.psi.KtNullableType
 import org.jetbrains.kotlin.psi.KtTypeReference
@@ -70,7 +71,19 @@ class KSTypeReferenceImpl(
     }
 
     override val annotations: Sequence<KSAnnotation> by lazy {
-        ktType.annotations.map { KSAnnotationImpl.getCached(it) }.asSequence()
+        val innerAnnotations = mutableListOf<Sequence<KtAnnotationEntry>>()
+        visitNullableType {
+            innerAnnotations.add(it.annotationEntries.asSequence())
+        }
+
+        (ktTypeReference.annotationEntries.asSequence() + innerAnnotations.asSequence().flatten())
+            .map { annotationEntry ->
+                KSAnnotationImpl.getCached(annotationEntry, this@KSTypeReferenceImpl) {
+                    ktType.annotations.single {
+                        it.psi == annotationEntry
+                    }
+                }
+            }
     }
 
     override val origin: Origin = parent?.origin ?: Origin.SYNTHETIC
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt
@@ -31,7 +31,7 @@ import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.psi.KtParameter
 
 class KSValueParameterImpl private constructor(
-    private val ktValueParameterSymbol: KtValueParameterSymbol,
+    internal val ktValueParameterSymbol: KtValueParameterSymbol,
     override val parent: KSAnnotated
 ) : KSValueParameter, Deferrable {
     companion object : KSObjectCache<KtValueParameterSymbol, KSValueParameterImpl>() {
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSAnnotationResolvedImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/resolved/KSAnnotationResolvedImpl.kt
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/util.kt

Original file line number	Diff line number	Diff line change
`@@ -85,7 +85,7 @@ class KSFileImpl private constructor(internal val ktFileSymbol: KtFileSymbol) :`
`85`	`85`	`}`
`86`	`86`
`87`	`87`	`override val annotations: Sequence<KSAnnotation> by lazy {`
`88`		`- ktFileSymbol.annotations(this)`
	`88`	`+ (ktFileSymbol.psi as? KtFile)?.annotations(ktFileSymbol) ?: ktFileSymbol.annotations(this)`
`89`	`89`	`}`
`90`	`90`
`91`	`91`	`override fun toString(): String {`