From 26673d2b08f01dec1a9007b9b75436a50fa497e9 Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Tue, 12 Oct 2021 14:21:42 -0700
Subject: [PATCH] FIR IDE: find source of deserialized enum entry

---
 .../FirIdeDeserializedDeclarationSourceProvider.kt   | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirIdeDeserializedDeclarationSourceProvider.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirIdeDeserializedDeclarationSourceProvider.kt
index f895a34bfa6..4f60c76c6ca 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirIdeDeserializedDeclarationSourceProvider.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/FirIdeDeserializedDeclarationSourceProvider.kt
@@ -32,6 +32,7 @@ object FirIdeDeserializedDeclarationSourceProvider {
             is FirClass -> provideSourceForClass(fir, project)
             is FirTypeAlias -> provideSourceForTypeAlias(fir, project)
             is FirConstructor -> provideSourceForConstructor(fir, project)
+            is FirEnumEntry -> provideSourceForEnumEntry(fir, project)
             else -> null
         }
     }
@@ -97,6 +98,17 @@ object FirIdeDeserializedDeclarationSourceProvider {
         return null
     }
 
+    private fun provideSourceForEnumEntry(
+        enumEntry: FirEnumEntry,
+        project: Project
+    ): PsiElement? {
+        val candidates = enumEntry.containingKtClass(project)?.body?.enumEntries
+            ?.filter { it.name == enumEntry.name.asString() }
+            .orEmpty()
+
+        return candidates.firstOrNull(KtElement::isCompiled)
+    }
+
     private fun FirDeclaration.scope(project: Project): GlobalSearchScope {
         return GlobalSearchScope.allScope(project)
         /* TODO: