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

idea/tests/org/jetbrains/kotlin/asJava/KtLightAnnotationTest.kt

@@ -21,6 +21,7 @@ import com.intellij.openapi.externalSystem.service.project.IdeModifiableModelsPr
 import com.intellij.psi.*
 import com.intellij.testFramework.LightProjectDescriptor
 import junit.framework.TestCase
+import org.jetbrains.kotlin.asJava.elements.KtLightAnnotationForSourceEntry
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.idea.facet.configureFacet
 import org.jetbrains.kotlin.idea.facet.getOrCreateFacet
@@ -187,6 +188,7 @@ class KtLightAnnotationTest : KotlinLightCodeInsightFixtureTestCase() {
         annotationAttributeVal as PsiAnnotation
         val innerAnnotationAttributeVal = annotationAttributeVal.findAttributeValue("anno3") as PsiElement
         assertTextAndRange("Anno3()", innerAnnotationAttributeVal)
+        assertIsKtLightAnnotation("Anno3()", innerAnnotationAttributeVal)
     }
 
     fun testKotlinAnnotationWithStringArray() {
@@ -235,16 +237,25 @@ class KtLightAnnotationTest : KotlinLightCodeInsightFixtureTestCase() {
         """.trimIndent())
         myFixture.testHighlighting("AnnotatedClass.kt")
 
-        val annotations = myFixture.findClass("AnnotatedClass").expectAnnotations(1)
-        val annotationAttributeVal = annotations.first().findAttributeValue("anno2") as PsiElement
-        assertTextAndRange("arrayOf(Anno2(1), Anno2(2))", annotationAttributeVal)
+        val annotation = myFixture.findClass("AnnotatedClass").expectAnnotations(1).single()
 
-        annotationAttributeVal as PsiArrayInitializerMemberValue
-        val innerAnnotationAttributeVal = annotationAttributeVal.initializers[0]
-        assertTextAndRange("Anno2(1)", innerAnnotationAttributeVal)
-        innerAnnotationAttributeVal as PsiAnnotation
-        val value = innerAnnotationAttributeVal.findAttributeValue("v")!!
-        assertTextAndRange("1", value)
+        val annotationAttributeVal = annotation.findAttributeValue("anno2") as PsiArrayInitializerMemberValue
+        assertTextAndRange("arrayOf(Anno2(1), Anno2(2))", annotationAttributeVal)
+        annotationAttributeVal.initializers[0].let { innerAnnotationAttributeVal ->
+            assertTextAndRange("Anno2(1)", innerAnnotationAttributeVal)
+            assertIsKtLightAnnotation("Anno2(1)", innerAnnotationAttributeVal)
+            innerAnnotationAttributeVal as PsiAnnotation
+            val value = innerAnnotationAttributeVal.findAttributeValue("v")!!
+            assertTextAndRange("1", value)
+        }
+
+
+        val attributeValueFromParameterList = annotation.parameterList.attributes.single().value as PsiArrayInitializerMemberValue
+        assertTextAndRange("arrayOf(Anno2(1), Anno2(2))", attributeValueFromParameterList)
+        attributeValueFromParameterList.initializers[0].let { innerAnnotationAttributeVal ->
+            assertTextAndRange("Anno2(1)", innerAnnotationAttributeVal)
+            assertIsKtLightAnnotation("Anno2(1)", innerAnnotationAttributeVal)
+        }
 
     }
 
@@ -270,6 +281,7 @@ class KtLightAnnotationTest : KotlinLightCodeInsightFixtureTestCase() {
             annotationAttributeVal as PsiArrayInitializerMemberValue
             annotationAttributeVal.initializers[0].let { innerAnnotationAttributeVal ->
                 assertTextAndRange("Inner()", innerAnnotationAttributeVal)
+                assertIsKtLightAnnotation("Inner()", innerAnnotationAttributeVal)
             }
         }
 
@@ -491,6 +503,10 @@ class KtLightAnnotationTest : KotlinLightCodeInsightFixtureTestCase() {
         TestCase.assertEquals(psiElement, PsiAnchor.create(psiElement).retrieve())
     }
 
+    private fun assertIsKtLightAnnotation(expected: String, psiElement: PsiElement) {
+        TestCase.assertEquals(expected, (psiElement as KtLightAnnotationForSourceEntry).kotlinOrigin.text)
+    }
+
     private fun assertTextRangeAndValue(expected: String, value: Any?, psiElement: PsiElement) {
         assertTextAndRange(expected, psiElement)
         val result = JavaPsiFacade.getInstance(project).constantEvaluationHelper.computeConstantExpression(psiElement)