FIR IDE: fix shortenings for enums

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");

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) {

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>
+}

idea/testData/shortenRefsFir/quailfiers/Enum.kt.after

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+enum class A {
+    B
+}
+
+fun usage() {
+    A.B
+}