LightAnnotations: wrapping nested annotations with KtLightAnnotationForSourceEntry

compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/lightAnnotations.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.classes.cannotModify
 import org.jetbrains.kotlin.asJava.classes.lazyPub
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.ClassConstructorDescriptor
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
@@ -36,6 +37,7 @@ import org.jetbrains.kotlin.resolve.calls.model.DefaultValueArgument
 import org.jetbrains.kotlin.resolve.calls.model.ExpressionValueArgument
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.VarargValueArgument
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.types.TypeUtils
 
@@ -202,6 +204,9 @@ class KtLightAnnotationForSourceEntry(
                 wrapAnnotationValue(memberValue, this, {
                     originalExpression.let { ktOrigin ->
                         when {
+                            ktOrigin is KtCallExpression
+                                    && memberValue is PsiAnnotation
+                                    && isAnnotationConstructorCall(ktOrigin, memberValue) -> ktOrigin
                             ktOrigin is KtValueArgumentList -> ktOrigin.arguments.getOrNull(i)?.getArgumentExpression()
                             ktOrigin is KtCallElement -> ktOrigin.valueArguments.getOrNull(i)?.getArgumentExpression()
                             ktOrigin is KtCollectionLiteralExpression -> ktOrigin.getInnerExpressions().getOrNull(i)
@@ -255,6 +260,25 @@ class KtLightAnnotationForSourceEntry(
 
     override fun findDeclaredAttributeValue(name: String?) = clsDelegate.findDeclaredAttributeValue(name)?.let { wrapAnnotationValue(it) }
 
+    override fun getParameterList(): PsiAnnotationParameterList = KtLightAnnotationParameterList(super.getParameterList())
+
+    inner class KtLightAnnotationParameterList(private val list: PsiAnnotationParameterList) : KtLightElementBase(this),
+        PsiAnnotationParameterList {
+        override val kotlinOrigin get() = null
+        override fun getAttributes(): Array<PsiNameValuePair> = list.attributes.map { KtLightPsiNameValuePair(it) }.toTypedArray()
+
+        inner class KtLightPsiNameValuePair(private val psiNameValuePair: PsiNameValuePair) : KtLightElementBase(this),
+            PsiNameValuePair {
+            override fun setValue(newValue: PsiAnnotationMemberValue): PsiAnnotationMemberValue = psiNameValuePair.setValue(newValue)
+            override fun getNameIdentifier(): PsiIdentifier? = psiNameValuePair.nameIdentifier
+            override fun getLiteralValue(): String? = (value as? PsiLiteralValue)?.value?.toString()
+            override val kotlinOrigin: KtElement? = null
+
+            override fun getValue(): PsiAnnotationMemberValue? = psiNameValuePair.value?.let { wrapAnnotationValue(it) }
+        }
+    }
+
+
     override fun delete() = kotlinOrigin.delete()
 
     override fun toString() = "@$qualifiedName"
@@ -337,6 +361,9 @@ private fun KtElement.getResolvedCall(): ResolvedCall<out CallableDescriptor>? {
     return this.getResolvedCall(context)
 }
 
+private fun isAnnotationConstructorCall(callExpression: KtCallExpression, psiAnnotation: PsiAnnotation) =
+    (callExpression.getResolvedCall()?.resultingDescriptor as? ClassConstructorDescriptor)?.constructedClass?.fqNameUnsafe?.toString() == psiAnnotation.qualifiedName
+
 private fun PsiElement.asKtCall(): KtCallElement? = (this as? KtElement)?.getResolvedCall()?.call?.callElement as? KtCallElement
 
 private fun psiReport(psiElement: PsiElement?): String {

compiler/light-classes/src/org/jetbrains/kotlin/asJava/lightClassUtils.kt

@@ -164,15 +164,31 @@ private val DEFAULT_IMPLS_CLASS_NAME = Name.identifier(JvmAbi.DEFAULT_IMPLS_CLAS
 fun FqName.defaultImplsChild() = child(DEFAULT_IMPLS_CLASS_NAME)
 
 @Suppress("unused")
-fun KtAnnotationEntry.toLightAnnotation(): PsiAnnotation? {
+fun KtElement.toLightAnnotation(): PsiAnnotation? {
     val ktDeclaration = getStrictParentOfType<KtModifierList>()?.parent as? KtDeclaration ?: return null
     for (lightElement in ktDeclaration.toLightElements()) {
         if (lightElement !is PsiModifierListOwner) continue
-        lightElement.modifierList?.annotations?.firstOrNull { it is KtLightAnnotationForSourceEntry && it.kotlinOrigin == this }?.let { return it }
+        for (rootAnnotation in lightElement.modifierList?.annotations ?: continue) {
+            for (annotation in rootAnnotation.withNestedAnnotations()) {
+                if (annotation is KtLightAnnotationForSourceEntry && annotation.kotlinOrigin == this)
+                    return annotation
+            }
+        }
     }
     return null
 }
 
+private fun PsiAnnotation.withNestedAnnotations(): Sequence<PsiAnnotation> {
+    fun handleValue(memberValue: PsiAnnotationMemberValue?): Sequence<PsiAnnotation> =
+        when (memberValue) {
+            is PsiArrayInitializerMemberValue ->
+                memberValue.initializers.asSequence().flatMap { handleValue(it) }
+            is PsiAnnotation -> memberValue.withNestedAnnotations()
+            else -> emptySequence()
+        }
+    return sequenceOf(this) + parameterList.attributes.asSequence().flatMap { handleValue(it.value) }
+}
+
 fun propertyNameByAccessor(name: String, accessor: KtLightMethod): String? {
     val toRename = accessor.kotlinOrigin ?: return null
     if (toRename !is KtProperty && toRename !is KtParameter) return null