From e53db4c771f7a284165cdbda1878f82b0fdb5cc6 Mon Sep 17 00:00:00 2001
From: Nicolay Mitropolsky <nicolay.mitropolsky@jetbrains.com>
Date: Thu, 30 Jul 2020 21:18:49 +0300
Subject: [PATCH] Uast: fixing `WrappedUAnnotation` identifiers

---
 .../KtLightPsiJavaCodeReferenceElement.kt        |  2 +-
 .../uast/kotlin/declarations/KotlinUVariable.kt  | 12 ++++++++++--
 .../kotlin/declarations/UastLightIdentifier.kt   | 16 ++++++++++------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightPsiJavaCodeReferenceElement.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightPsiJavaCodeReferenceElement.kt
index f1578f67080..ec86fb95482 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightPsiJavaCodeReferenceElement.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightPsiJavaCodeReferenceElement.kt
@@ -25,7 +25,7 @@ class KtLightPsiJavaCodeReferenceElement(
     override fun advancedResolve(incompleteCode: Boolean): JavaResolveResult =
         delegate?.advancedResolve(incompleteCode) ?: JavaResolveResult.EMPTY
 
-    override fun getReferenceNameElement(): PsiElement? = ktElement
+    override fun getReferenceNameElement(): PsiElement? = delegate?.referenceNameElement
 
     override fun getTypeParameters(): Array<PsiType> = delegate?.typeParameters ?: emptyArray()
 
diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt
index 1f4898b8c37..194fd05e0f9 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt
@@ -22,6 +22,7 @@ import com.intellij.psi.util.PsiTypesUtil
 import org.jetbrains.annotations.NotNull
 import org.jetbrains.annotations.Nullable
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
+import org.jetbrains.kotlin.asJava.elements.KtLightAbstractAnnotation
 import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.psi.*
@@ -31,6 +32,7 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.TypeNullability
 import org.jetbrains.kotlin.types.typeUtil.nullability
 import org.jetbrains.kotlin.utils.SmartList
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 import org.jetbrains.uast.*
 import org.jetbrains.uast.internal.acceptList
 import org.jetbrains.uast.kotlin.declarations.KotlinUIdentifier
@@ -123,13 +125,19 @@ abstract class AbstractKotlinUVariable(givenParent: UElement?) : KotlinAbstractU
 
         override val javaPsi: PsiAnnotation = psiAnnotation
         override val psi: PsiAnnotation = javaPsi
-        override val sourcePsi: PsiElement? = null
+        override val sourcePsi: PsiElement? = psiAnnotation.safeAs<KtLightAbstractAnnotation>()?.kotlinOrigin
 
         override val attributeValues: List<UNamedExpression> by lz {
             psi.parameterList.attributes.map { WrappedUNamedExpression(it, this) }
         }
 
-        override val uastAnchor by lazy { KotlinUIdentifier(javaPsi.nameReferenceElement?.referenceNameElement, null, this) }
+        override val uastAnchor by lazy {
+            KotlinUIdentifier(
+                { javaPsi.nameReferenceElement?.referenceNameElement },
+                sourcePsi.safeAs<KtAnnotationEntry>()?.typeReference?.nameElement,
+                this
+            )
+        }
 
         class WrappedUNamedExpression(pair: PsiNameValuePair, override val uastParent: UElement?) : UNamedExpression {
             override val name: String? = pair.name
diff --git a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/UastLightIdentifier.kt b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/UastLightIdentifier.kt
index d7e74ebe48e..d5445528964 100644
--- a/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/UastLightIdentifier.kt
+++ b/plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/UastLightIdentifier.kt
@@ -34,12 +34,16 @@ class UastLightIdentifier(lightOwner: PsiNameIdentifierOwner, ktDeclaration: KtD
     override fun getContainingFile(): PsiFile = unwrapFakeFileForLightClass(super.getContainingFile())
 }
 
-class KotlinUIdentifier private constructor(
-    override val javaPsi: PsiElement?,
+class KotlinUIdentifier constructor(
+    javaPsiSupplier: () -> PsiElement?,
     override val sourcePsi: PsiElement?,
-    override val psi: PsiElement?,
     givenParent: UElement?
-) : UIdentifier(psi, givenParent) {
+) : UIdentifier(sourcePsi, givenParent) {
+
+    override val javaPsi: PsiElement? by lazy(javaPsiSupplier) // don't know any real need to call it in production
+
+    override val psi: PsiElement?
+        get() = javaPsi ?: sourcePsi
 
     init {
         if (ApplicationManager.getApplication().isUnitTestMode && !acceptableSourcePsi(sourcePsi))
@@ -72,6 +76,6 @@ class KotlinUIdentifier private constructor(
         return null
     }
 
-    constructor(javaPsi: PsiElement?, sourcePsi: PsiElement?, uastParent: UElement?) : this(javaPsi, sourcePsi, javaPsi, uastParent)
-    constructor(sourcePsi: PsiElement?, uastParent: UElement?) : this(null, sourcePsi, sourcePsi, uastParent)
+    constructor(javaPsi: PsiElement?, sourcePsi: PsiElement?, uastParent: UElement?) : this({ javaPsi }, sourcePsi, uastParent)
+    constructor(sourcePsi: PsiElement?, uastParent: UElement?) : this({ null }, sourcePsi, uastParent)
 }
\ No newline at end of file