From 944bc31f696582652819e51cfa2ae9f2b7f429cb Mon Sep 17 00:00:00 2001
From: Nicolay Mitropolsky <nicolay.mitropolsky@jetbrains.com>
Date: Tue, 26 Mar 2019 13:26:22 +0300
Subject: [PATCH] Making Kotlin light elements `isEquivalentTo` to it's origins
 (KT-30583)

---
 .../KtLightClassForSourceDeclaration.kt       |  3 ++-
 .../asJava/elements/KtLightElementBase.kt     |  2 ++
 .../asJava/elements/KtLightMemberImpl.kt      |  2 ++
 .../asJava/elements/KtLightParameter.kt       |  2 ++
 .../lightClasses/LightElementsEqualsTest.kt   | 19 +++++++++++++++++++
 5 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/KtLightClassForSourceDeclaration.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/KtLightClassForSourceDeclaration.kt
index 5689f7c9e16..f2d6dc712dc 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/KtLightClassForSourceDeclaration.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/KtLightClassForSourceDeclaration.kt
@@ -139,7 +139,8 @@ abstract class KtLightClassForSourceDeclaration(protected val classOrObject: KtC
     override fun getNavigationElement(): PsiElement = classOrObject
 
     override fun isEquivalentTo(another: PsiElement?): Boolean {
-        return another is KtLightClassForSourceDeclaration && Comparing.equal(another.qualifiedName, qualifiedName)
+        return kotlinOrigin.isEquivalentTo(another) ||
+                another is KtLightClassForSourceDeclaration && Comparing.equal(another.qualifiedName, qualifiedName)
     }
 
     override fun getElementIcon(flags: Int): Icon? {
diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightElementBase.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightElementBase.kt
index 7d209b506e1..58c4a083ca6 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightElementBase.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightElementBase.kt
@@ -40,4 +40,6 @@ abstract class KtLightElementBase(private val parent: PsiElement): LightElement(
     override fun getPresentation() = (kotlinOrigin ?: this).let { ItemPresentationProviders.getItemPresentation(it) }
     override fun isValid() = parent.isValid
     override fun findElementAt(offset: Int) = kotlinOrigin?.findElementAt(offset)
+    override fun isEquivalentTo(another: PsiElement?): Boolean =
+        super.isEquivalentTo(another) || kotlinOrigin?.isEquivalentTo(another) == true
 }
diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightMemberImpl.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightMemberImpl.kt
index 21e85ca0dfe..3925b845202 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightMemberImpl.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightMemberImpl.kt
@@ -66,6 +66,8 @@ abstract class KtLightMemberImpl<out D : PsiMember>(
     override fun isDeprecated() = (clsDelegate as PsiDocCommentOwner).isDeprecated
 
     override fun isEquivalentTo(another: PsiElement?): Boolean {
+        if (lightMemberOrigin?.originalElement?.isEquivalentTo(another) == true) return true
+
         val isEquivalentByOrigin =
                 another is KtLightMember<*> &&
                 lightMemberOrigin?.isEquivalentTo(another.lightMemberOrigin) == true
diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightParameter.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightParameter.kt
index 40a9eca4b82..af877347574 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightParameter.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightParameter.kt
@@ -102,6 +102,8 @@ class KtLightParameter(
     override fun isEquivalentTo(another: PsiElement?): Boolean {
         val result = ApplicationManager.getApplication().runReadAction(Computable {
             val kotlinOrigin = kotlinOrigin
+            if (kotlinOrigin?.isEquivalentTo(another) == true) return@Computable true
+
             if (another is KtLightParameter && kotlinOrigin != null) {
                 kotlinOrigin == another.kotlinOrigin && clsDelegate == another.clsDelegate
             }
diff --git a/idea/tests/org/jetbrains/kotlin/idea/lightClasses/LightElementsEqualsTest.kt b/idea/tests/org/jetbrains/kotlin/idea/lightClasses/LightElementsEqualsTest.kt
index aba29523b59..343b0c19259 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/lightClasses/LightElementsEqualsTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/lightClasses/LightElementsEqualsTest.kt
@@ -53,6 +53,25 @@ class LightElementsEqualsTest : KotlinLightCodeInsightFixtureTestCase() {
 
     }
 
+
+    fun `test light elements are "isEquivalentTo" to it's origins`() {
+        val psiFile = myFixture.configureByText("a.kt", SAMPLE_SOURCE)
+        psiFile.accept(object : PsiElementVisitor() {
+            override fun visitElement(element: PsiElement) {
+                if (element is KtElement) {
+                    for (lightElement in element.toLightElements()) {
+                        TestCase.assertTrue(
+                            "light element '$lightElement'[${lightElement.javaClass}] should be \"isEquivalentTo\" to it's origin '$element'",
+                            lightElement.isEquivalentTo(element)
+                        )
+                    }
+                }
+                element.acceptChildren(this)
+            }
+        })
+
+    }
+
     fun testToLightMethodsConvertedEquality() {
         myFixture.configureByText("a.kt", SAMPLE_SOURCE)