Skip to content

Commit

Permalink
AA: support java annotation argument evaluation
Browse files Browse the repository at this point in the history
  • Loading branch information
neetopia committed Jul 26, 2023
1 parent b360931 commit c1feea0
Show file tree
Hide file tree
Showing 8 changed files with 346 additions and 39 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,180 @@
package com.google.devtools.ksp.impl.symbol.java

import com.google.devtools.ksp.KSObjectCache
import com.google.devtools.ksp.getClassDeclarationByName
import com.google.devtools.ksp.impl.ResolverAAImpl
import com.google.devtools.ksp.impl.symbol.kotlin.KSTypeImpl
import com.google.devtools.ksp.impl.symbol.kotlin.KSTypeReferenceImpl
import com.google.devtools.ksp.impl.symbol.kotlin.KSValueArgumentImpl
import com.google.devtools.ksp.impl.symbol.kotlin.analyze
import com.google.devtools.ksp.impl.symbol.kotlin.classifierSymbol
import com.google.devtools.ksp.impl.symbol.kotlin.getDefaultValue
import com.google.devtools.ksp.impl.symbol.kotlin.toKSDeclaration
import com.google.devtools.ksp.impl.symbol.kotlin.toKtClassSymbol
import com.google.devtools.ksp.impl.symbol.kotlin.toLocation
import com.google.devtools.ksp.processing.impl.KSNameImpl
import com.google.devtools.ksp.symbol.AnnotationUseSiteTarget
import com.google.devtools.ksp.symbol.ClassKind
import com.google.devtools.ksp.symbol.KSAnnotation
import com.google.devtools.ksp.symbol.KSClassDeclaration
import com.google.devtools.ksp.symbol.KSName
import com.google.devtools.ksp.symbol.KSNode
import com.google.devtools.ksp.symbol.KSTypeReference
import com.google.devtools.ksp.symbol.KSValueArgument
import com.google.devtools.ksp.symbol.KSVisitor
import com.google.devtools.ksp.symbol.Location
import com.google.devtools.ksp.symbol.Origin
import com.intellij.lang.jvm.JvmClassKind
import com.intellij.psi.JavaPsiFacade
import com.intellij.psi.PsiAnnotation
import com.intellij.psi.PsiAnnotationMemberValue
import com.intellij.psi.PsiArrayInitializerMemberValue
import com.intellij.psi.PsiClass
import com.intellij.psi.PsiField
import com.intellij.psi.PsiLiteralValue
import com.intellij.psi.PsiReference
import com.intellij.psi.PsiType
import com.intellij.psi.impl.source.PsiAnnotationMethodImpl
import org.jetbrains.kotlin.analysis.api.annotations.KtNamedAnnotationValue
import org.jetbrains.kotlin.analysis.api.symbols.KtClassOrObjectSymbol
import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
import org.jetbrains.kotlin.analysis.api.types.KtType
import org.jetbrains.kotlin.name.ClassId

class KSAnnotationJavaImpl private constructor(private val psi: PsiAnnotation) : KSAnnotation {
companion object : KSObjectCache<PsiAnnotation, KSAnnotationJavaImpl>() {
fun getCached(psi: PsiAnnotation) =
KSAnnotationJavaImpl.cache.getOrPut(psi) { KSAnnotationJavaImpl(psi) }
}

private val type: KtType by lazy {
(ResolverAAImpl.instance.getClassDeclarationByName(psi.qualifiedName!!)!!.asStarProjectedType() as KSTypeImpl)
.type
}

override val annotationType: KSTypeReference by lazy {
KSTypeReferenceImpl.getCached(type, this)
}

override val arguments: List<KSValueArgument> by lazy {
val annotationConstructor = analyze {
(type.classifierSymbol() as KtClassOrObjectSymbol).getMemberScope().getConstructors().singleOrNull()
}
val presentArgs = psi.parameterList.attributes.mapIndexed { index, it ->
val name = it.name ?: annotationConstructor?.valueParameters?.getOrNull(index)?.name?.asString()
val value = it.value
val calculatedValue: Any? = if (value is PsiArrayInitializerMemberValue) {
value.initializers.map {
calcValue(it)
}
} else {
calcValue(it.value)
}
KSValueArgumentLiteImpl.getCached(
name?.let { KSNameImpl.getCached(it) },
calculatedValue,
Origin.JAVA
)
}
val presentValueArgumentNames = presentArgs.map { it.name?.asString() ?: "" }
presentArgs + defaultArguments.filter { it.name?.asString() !in presentValueArgumentNames }
}

override val defaultArguments: List<KSValueArgument> by lazy {
analyze {
(type.classifierSymbol() as KtClassOrObjectSymbol).getMemberScope().getConstructors().singleOrNull()
?.let { symbol ->
if (symbol.origin == KtSymbolOrigin.JAVA && symbol.psi != null) {
(symbol.psi as PsiClass).allMethods.filterIsInstance<PsiAnnotationMethodImpl>()
.mapNotNull { annoMethod ->
annoMethod.defaultValue?.let {
KSValueArgumentLiteImpl.getCached(
KSNameImpl.getCached(annoMethod.name),
calcValue(it),
Origin.SYNTHETIC
)
}
}
} else {
symbol.valueParameters.mapNotNull { valueParameterSymbol ->
valueParameterSymbol.getDefaultValue()?.let { constantValue ->
KSValueArgumentImpl.getCached(
KtNamedAnnotationValue(
valueParameterSymbol.name, constantValue,
),
Origin.SYNTHETIC
)
}
}
}
}
} ?: emptyList()
}

override val shortName: KSName by lazy {
KSNameImpl.getCached(psi.qualifiedName!!.split(".").last())
}

override val useSiteTarget: AnnotationUseSiteTarget? = null

override val origin: Origin = Origin.JAVA

override val location: Location
get() = psi.toLocation()

override val parent: KSNode?
get() = TODO("Not yet implemented")

override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
return visitor.visitAnnotation(this, data)
}

override fun toString(): String {
return "@${shortName.asString()}"
}
}

fun calcValue(value: PsiAnnotationMemberValue?): Any? {
if (value is PsiAnnotation) {
return KSAnnotationJavaImpl.getCached(value)
}
val result = when (value) {
is PsiReference -> value.resolve()?.let { resolved ->
JavaPsiFacade.getInstance(value.project).constantEvaluationHelper.computeConstantExpression(value)
?: resolved
}
else -> value?.let {
JavaPsiFacade.getInstance(value.project).constantEvaluationHelper.computeConstantExpression(value)
}
}
return when (result) {
is PsiType -> {
analyze {
(ClassId.fromString(result.canonicalText).toKtClassSymbol()?.toKSDeclaration() as? KSClassDeclaration)
?.asStarProjectedType()
}
}
is PsiLiteralValue -> {
result.value
}
is PsiField -> {
// manually handle enums as constant expression evaluator does not seem to be resolving them.
val containingClass = result.containingClass
if (containingClass?.classKind == JvmClassKind.ENUM) {
// this is an enum entry
containingClass.qualifiedName?.let {
ResolverAAImpl.instance!!.getClassDeclarationByName(it)
}?.declarations?.find {
it is KSClassDeclaration && it.classKind == ClassKind.ENUM_ENTRY &&
it.simpleName.asString() == result.name
}?.let { (it as KSClassDeclaration).asStarProjectedType() }
?.let {
return it
}
} else {
null
}
}
else -> result
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
package com.google.devtools.ksp.impl.symbol.java

import com.google.devtools.ksp.IdKeyPair
import com.google.devtools.ksp.KSObjectCache
import com.google.devtools.ksp.symbol.KSAnnotation
import com.google.devtools.ksp.symbol.KSName
import com.google.devtools.ksp.symbol.KSNode
import com.google.devtools.ksp.symbol.KSValueArgument
import com.google.devtools.ksp.symbol.KSVisitor
import com.google.devtools.ksp.symbol.Location
import com.google.devtools.ksp.symbol.Origin

class KSValueArgumentLiteImpl private constructor(
override val name: KSName?,
override val value: Any?,
override val origin: Origin
) : KSValueArgument {
companion object : KSObjectCache<IdKeyPair<KSName?, Any?>, KSValueArgumentLiteImpl>() {
fun getCached(name: KSName?, value: Any?, origin: Origin) =
KSValueArgumentLiteImpl.cache
.getOrPut(IdKeyPair(name, value)) { KSValueArgumentLiteImpl(name, value, origin) }
}
override val isSpread: Boolean = false

override val annotations: Sequence<KSAnnotation> = emptySequence()

override val location: Location
get() = TODO("Not yet implemented")

override val parent: KSNode?
get() = TODO("Not yet implemented")

override fun <D, R> accept(visitor: KSVisitor<D, R>, data: D): R {
return visitor.visitValueArgument(this, data)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
*/
package com.google.devtools.ksp.impl.symbol.kotlin

import com.google.devtools.ksp.impl.symbol.java.KSAnnotationJavaImpl
import com.google.devtools.ksp.processing.impl.KSNameImpl
import com.google.devtools.ksp.symbol.KSAnnotation
import com.google.devtools.ksp.symbol.KSDeclaration
Expand All @@ -27,12 +28,14 @@ import com.google.devtools.ksp.symbol.Location
import com.google.devtools.ksp.symbol.Modifier
import com.google.devtools.ksp.symbol.Origin
import com.google.devtools.ksp.toKSModifiers
import com.intellij.psi.PsiJvmModifiersOwner
import com.intellij.psi.PsiModifierListOwner
import org.jetbrains.kotlin.analysis.api.symbols.KtClassOrObjectSymbol
import org.jetbrains.kotlin.analysis.api.symbols.KtDeclarationSymbol
import org.jetbrains.kotlin.analysis.api.symbols.KtFunctionLikeSymbol
import org.jetbrains.kotlin.analysis.api.symbols.KtPropertySymbol
import org.jetbrains.kotlin.analysis.api.symbols.markers.KtNamedSymbol
import org.jetbrains.kotlin.psi.KtElement
import org.jetbrains.kotlin.psi.KtModifierListOwner

abstract class AbstractKSDeclarationImpl(val ktDeclarationSymbol: KtDeclarationSymbol) : KSDeclaration {
Expand Down Expand Up @@ -98,5 +101,10 @@ abstract class AbstractKSDeclarationImpl(val ktDeclarationSymbol: KtDeclarationS
override val docString: String?
get() = ktDeclarationSymbol.toDocString()

internal val originalAnnotations = ktDeclarationSymbol.annotations(this)
internal val originalAnnotations = if (ktDeclarationSymbol.psi is KtElement || ktDeclarationSymbol.psi == null) {
ktDeclarationSymbol.annotations(this)
} else {
(ktDeclarationSymbol.psi as PsiJvmModifiersOwner)
.annotations.map { KSAnnotationJavaImpl.getCached(it) }.asSequence()
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,17 +19,17 @@ package com.google.devtools.ksp.impl.symbol.kotlin

import com.google.devtools.ksp.IdKeyPair
import com.google.devtools.ksp.KSObjectCache
import com.google.devtools.ksp.impl.symbol.java.KSValueArgumentLiteImpl
import com.google.devtools.ksp.impl.symbol.java.calcValue
import com.google.devtools.ksp.processing.impl.KSNameImpl
import com.google.devtools.ksp.symbol.*
import com.intellij.psi.PsiClass
import com.intellij.psi.impl.source.PsiAnnotationMethodImpl
import org.jetbrains.kotlin.analysis.api.annotations.KtAnnotationApplicationWithArgumentsInfo
import org.jetbrains.kotlin.analysis.api.annotations.KtConstantAnnotationValue
import org.jetbrains.kotlin.analysis.api.annotations.KtNamedAnnotationValue
import org.jetbrains.kotlin.analysis.api.base.KtConstantValue
import org.jetbrains.kotlin.analysis.api.components.KtConstantEvaluationMode
import org.jetbrains.kotlin.analysis.api.components.buildClassType
import org.jetbrains.kotlin.analysis.api.symbols.KtValueParameterSymbol
import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget.*
import org.jetbrains.kotlin.psi.KtParameter

class KSAnnotationImpl private constructor(
private val annotationApplication: KtAnnotationApplicationWithArgumentsInfo,
Expand All @@ -47,31 +47,46 @@ class KSAnnotationImpl private constructor(
}

override val arguments: List<KSValueArgument> by lazy {
val presentArgs = annotationApplication.arguments.map { KSValueArgumentImpl.getCached(it) }
val presentArgs = annotationApplication.arguments.map { KSValueArgumentImpl.getCached(it, Origin.KOTLIN) }
val presentNames = presentArgs.mapNotNull { it.name?.asString() }
val absentArgs = analyze {
val absentArgs = defaultArguments.filter {
it.name?.asString() !in presentNames
}
presentArgs + absentArgs
}

override val defaultArguments: List<KSValueArgument> by lazy {
analyze {
annotationApplication.classId?.toKtClassSymbol()?.let { symbol ->
symbol.getMemberScope().getConstructors().singleOrNull()?.let {
it.valueParameters.filter { valueParameter ->
valueParameter.name.asString() !in presentNames
}.mapNotNull { valueParameterSymbol ->
valueParameterSymbol.getDefaultValue()?.let { constantValue ->
KSValueArgumentImpl.getCached(
KtNamedAnnotationValue(
valueParameterSymbol.name, KtConstantAnnotationValue(constantValue)
if (symbol.origin == KtSymbolOrigin.JAVA && symbol.psi != null) {
(symbol.psi as PsiClass).allMethods.filterIsInstance<PsiAnnotationMethodImpl>()
.mapNotNull { annoMethod ->
annoMethod.defaultValue?.let {
KSValueArgumentLiteImpl.getCached(
KSNameImpl.getCached(annoMethod.name),
calcValue(it),
Origin.SYNTHETIC
)
}
}
} else {
symbol.getMemberScope().getConstructors().singleOrNull()?.let {
it.valueParameters.mapNotNull { valueParameterSymbol ->
valueParameterSymbol.getDefaultValue()?.let { constantValue ->
KSValueArgumentImpl.getCached(
KtNamedAnnotationValue(
valueParameterSymbol.name, constantValue,
),
Origin.SYNTHETIC
)
)
}
}
}
}
} ?: emptyList<KSValueArgument>()
} ?: emptyList()
}
presentArgs + absentArgs
}

override val defaultArguments: List<KSValueArgument>
get() = TODO("Not yet implemented")

override val shortName: KSName by lazy {
KSNameImpl.getCached(annotationApplication.classId!!.shortClassName.asString())
}
Expand Down Expand Up @@ -105,14 +120,3 @@ class KSAnnotationImpl private constructor(
return "@${shortName.asString()}"
}
}

internal fun KtValueParameterSymbol.getDefaultValue(): KtConstantValue? {
return this.psi?.let {
when (it) {
is KtParameter -> analyze {
it.defaultValue?.evaluate(KtConstantEvaluationMode.CONSTANT_EXPRESSION_EVALUATION)
}
else -> throw IllegalStateException("Unhandled default value type ${it.javaClass}")
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,14 @@ class KSFunctionDeclarationImpl private constructor(internal val ktFunctionSymbo
(origin == Origin.JAVA && ktFunctionSymbol.psi == null || ktFunctionSymbol.psi is PsiClass)
}

override val annotations: Sequence<KSAnnotation> by lazy {
if (isSyntheticConstructor()) {
emptySequence()
} else {
super.annotations
}
}

override fun toString(): String {
// TODO: fix origin for implicit Java constructor in AA
// TODO: should we change the toString() behavior for synthetic constructors?
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,12 @@ import org.jetbrains.kotlin.analysis.api.annotations.KtNamedAnnotationValue
import org.jetbrains.kotlin.analysis.api.annotations.KtUnsupportedAnnotationValue

class KSValueArgumentImpl private constructor(
private val namedAnnotationValue: KtNamedAnnotationValue
private val namedAnnotationValue: KtNamedAnnotationValue,
override val origin: Origin
) : KSValueArgument {
companion object : KSObjectCache<KtNamedAnnotationValue, KSValueArgumentImpl>() {
fun getCached(namedAnnotationValue: KtNamedAnnotationValue) =
cache.getOrPut(namedAnnotationValue) { KSValueArgumentImpl(namedAnnotationValue) }
fun getCached(namedAnnotationValue: KtNamedAnnotationValue, origin: Origin) =
cache.getOrPut(namedAnnotationValue) { KSValueArgumentImpl(namedAnnotationValue, origin) }
}

override val name: KSName? by lazy {
Expand All @@ -47,8 +48,6 @@ class KSValueArgumentImpl private constructor(

override val annotations: Sequence<KSAnnotation> = emptySequence()

override val origin: Origin = Origin.KOTLIN

override val location: Location by lazy {
namedAnnotationValue.expression.sourcePsi?.toLocation() ?: NonExistLocation
}
Expand Down
Loading

0 comments on commit c1feea0

Please sign in to comment.