Skip to content

Commit

Permalink
defer resolution of annotations when possible.
Browse files Browse the repository at this point in the history
  • Loading branch information
@neetopia
neetopia committed Jul 9, 2024
1 parent dc38b13 commit c28c357
Show file tree
Hide file tree
Showing 10 changed files with 269 additions and 35 deletions.
8 changes: 5 additions & 3 deletions ...i/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/AbstractKSDeclarationImpl.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
Expand Down
42 changes: 23 additions & 19 deletions ...alysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSAnnotationImpl.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
}
}
Expand All @@ -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 ->
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSFileImpl.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
5 changes: 4 additions & 1 deletion ...-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyAccessorImpl.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.*
Expand All @@ -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)
}

Expand Down
21 changes: 17 additions & 4 deletions ...i/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSPropertyDeclarationImpl.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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) :
Expand All @@ -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
Expand All @@ -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()
)
}

Expand Down Expand Up @@ -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) {
Expand Down
7 changes: 5 additions & 2 deletions kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeImpl.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
15 changes: 14 additions & 1 deletion ...sis-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSTypeReferenceImpl.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion ...is-api/src/main/kotlin/com/google/devtools/ksp/impl/symbol/kotlin/KSValueParameterImpl.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>() {
Expand Down
Loading

0 comments on commit c28c357

Please sign in to comment.