From edbfff71e9d1f8581637da9d0ebe7b7348695c60 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Mon, 31 May 2021 15:25:34 +0200
Subject: [PATCH] FIR IDE: fix shortenings for enums

---
 .../FirShortenRefsTestGenerated.java           |  5 +++++
 .../fir/components/KtFirReferenceShortener.kt  | 18 +++++++++++++++---
 .../testData/shortenRefsFir/quailfiers/Enum.kt | 10 ++++++++++
 .../shortenRefsFir/quailfiers/Enum.kt.after    | 10 ++++++++++
 4 files changed, 40 insertions(+), 3 deletions(-)
 create mode 100644 idea/testData/shortenRefsFir/quailfiers/Enum.kt
 create mode 100644 idea/testData/shortenRefsFir/quailfiers/Enum.kt.after

diff --git a/idea/idea-fir/tests/org/jetbrains/kotlin/idea/fir/shortenRefs/FirShortenRefsTestGenerated.java b/idea/idea-fir/tests/org/jetbrains/kotlin/idea/fir/shortenRefs/FirShortenRefsTestGenerated.java
index bceb55ddaeb..8ece95740e7 100644
--- a/idea/idea-fir/tests/org/jetbrains/kotlin/idea/fir/shortenRefs/FirShortenRefsTestGenerated.java
+++ b/idea/idea-fir/tests/org/jetbrains/kotlin/idea/fir/shortenRefs/FirShortenRefsTestGenerated.java
@@ -177,6 +177,11 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
             runTest("idea/testData/shortenRefsFir/quailfiers/AlreadyImportedNestedType.kt");
         }
 
+        @TestMetadata("Enum.kt")
+        public void testEnum() throws Exception {
+            runTest("idea/testData/shortenRefsFir/quailfiers/Enum.kt");
+        }
+
         @TestMetadata("NestedTypeInSameFile.kt")
         public void testNestedTypeInSameFile() throws Exception {
             runTest("idea/testData/shortenRefsFir/quailfiers/NestedTypeInSameFile.kt");
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt
index 1ac5c9a7acf..bba0eb88d92 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt
@@ -188,7 +188,10 @@ private class FirShorteningContext(val firResolveState: FirModuleResolveState) {
 
 private sealed class ElementToShorten
 private class ShortenType(val element: KtUserType, val nameToImport: FqName? = null) : ElementToShorten()
-private class ShortenQualifier(val element: KtDotQualifiedExpression, val nameToImport: FqName? = null) : ElementToShorten()
+
+private class ShortenQualifier(val element: KtDotQualifiedExpression, val nameToImport: FqName? = null) : ElementToShorten() {
+    fun withElement(newElement: KtDotQualifiedExpression) = ShortenQualifier(newElement, nameToImport)
+}
 
 private class ElementsToShortenCollector(
     private val shorteningContext: FirShorteningContext,
@@ -325,7 +328,8 @@ private class ElementsToShortenCollector(
         val referenceExpression = resolvedNamedReference.psi as? KtNameReferenceExpression
         val qualifiedProperty = referenceExpression?.getDotQualifiedExpressionForSelector() ?: return
 
-        val propertyId = (resolvedNamedReference.resolvedSymbol as? FirCallableSymbol<*>)?.callableId ?: return
+        val firCallableSymbol = resolvedNamedReference.resolvedSymbol as? FirCallableSymbol<*> ?: return
+        val propertyId = firCallableSymbol.callableId
 
         val scopes = shorteningContext.findScopesAtPosition(qualifiedProperty, namesToImport, towerContextProvider) ?: return
         val singleAvailableProperty = shorteningContext.findPropertiesInScopes(scopes, propertyId.callableName)
@@ -336,7 +340,15 @@ private class ElementsToShortenCollector(
             else -> findFakePackageToShorten(qualifiedProperty)
         }
 
-        propertyToShorten?.let(::addElementToShorten)
+        propertyToShorten?.withQualifierForCallablesToShorten(firCallableSymbol)?.let(::addElementToShorten)
+    }
+
+    private fun ShortenQualifier.withQualifierForCallablesToShorten(callableSymbol: FirCallableSymbol<*>): ShortenQualifier? {
+        if (callableSymbol.fir is FirEnumEntry) {
+            val newQualifier = element.receiverExpression as? KtDotQualifiedExpression ?: return null
+            return withElement(newQualifier)
+        }
+        return this
     }
 
     private fun processFunctionCall(functionCall: FirFunctionCall) {
diff --git a/idea/testData/shortenRefsFir/quailfiers/Enum.kt b/idea/testData/shortenRefsFir/quailfiers/Enum.kt
new file mode 100644
index 00000000000..02f110edda8
--- /dev/null
+++ b/idea/testData/shortenRefsFir/quailfiers/Enum.kt
@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+enum class A {
+    B
+}
+
+fun usage() {
+    <selection>test.A.B</selection>
+}
\ No newline at end of file
diff --git a/idea/testData/shortenRefsFir/quailfiers/Enum.kt.after b/idea/testData/shortenRefsFir/quailfiers/Enum.kt.after
new file mode 100644
index 00000000000..36640817a35
--- /dev/null
+++ b/idea/testData/shortenRefsFir/quailfiers/Enum.kt.after
@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+enum class A {
+    B
+}
+
+fun usage() {
+    A.B
+}