FIR IDE: Remove fake root prefix even when the element is not shortened

idea/idea-fir/tests/org/jetbrains/kotlin/shortenRefs/FirShortenRefsTestGenerated.java

@@ -126,11 +126,6 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
             runTest("idea/testData/shortenRefsFir/calls/rootPackage.kt");
         }
 
-        @TestMetadata("rootPackageShortenFakeRootPackage.kt")
-        public void testRootPackageShortenFakeRootPackage() throws Exception {
-            runTest("idea/testData/shortenRefsFir/calls/rootPackageShortenFakeRootPackage.kt");
-        }
-
         @TestMetadata("selfReferencingFunction.kt")
         public void testSelfReferencingFunction() throws Exception {
             runTest("idea/testData/shortenRefsFir/calls/selfReferencingFunction.kt");
@@ -142,6 +137,24 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
         }
     }
 
+    @TestMetadata("idea/testData/shortenRefsFir/fakeRootPackage")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class FakeRootPackage extends AbstractFirShortenRefsTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTestWithMuting, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInFakeRootPackage() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/shortenRefsFir/fakeRootPackage"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("rootPackageShortenFakeRootPackage.kt")
+        public void testRootPackageShortenFakeRootPackage() throws Exception {
+            runTest("idea/testData/shortenRefsFir/fakeRootPackage/rootPackageShortenFakeRootPackage.kt");
+        }
+    }
+
     @TestMetadata("idea/testData/shortenRefsFir/quailfiers")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt

@@ -11,6 +11,7 @@ import com.intellij.psi.SmartPsiElementPointer
 import com.intellij.util.containers.addIfNotNull
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE
 import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClass
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
@@ -209,6 +210,17 @@ internal class KtFirReferenceShortener(
 
             if (availableClassifier == null || availableClassifier.isFromStarOrPackageImport) {
                 addTypeToImportAndShorten(mostTopLevelClassId.asSingleFqName(), mostTopLevelTypeElement)
+            } else {
+                addFakePackagePrefixToShortenIfPresent(mostTopLevelTypeElement)
+            }
+        }
+
+        private fun addFakePackagePrefixToShortenIfPresent(typeElement: KtUserType) {
+            val deepestTypeWithQualifier = typeElement.qualifiersWithSelf.last().parent as? KtUserType
+                ?: error("Type element should have at least one qualifier, instead it was ${typeElement.text}")
+
+            if (deepestTypeWithQualifier.hasFakeRootPrefix()) {
+                addTypeToShorten(deepestTypeWithQualifier)
             }
         }
 
@@ -260,11 +272,17 @@ internal class KtFirReferenceShortener(
             val scopes = findScopesAtPosition(callExpression, namesToImport) ?: return
             val availableCallables = findFunctionsInScopes(scopes, callableId.callableName)
 
-            if (availableCallables.isEmpty()) {
-                val additionalImport = callableId.asImportableFqName() ?: return
-                addElementToImportAndShorten(additionalImport, qualifiedCallExpression)
-            } else if (availableCallables.all { it.callableId == callableId }) {
-                addElementToShorten(qualifiedCallExpression)
+            when {
+                availableCallables.isEmpty() -> {
+                    val additionalImport = callableId.asImportableFqName() ?: return
+                    addElementToImportAndShorten(additionalImport, qualifiedCallExpression)
+                }
+                availableCallables.all { it.callableId == callableId } -> {
+                    addElementToShorten(qualifiedCallExpression)
+                }
+                else -> {
+                    addFakePackagePrefixToShortenIfPresent(qualifiedCallExpression)
+                }
             }
         }
 
@@ -346,6 +364,15 @@ internal class KtFirReferenceShortener(
 
             if (availableClassifier == null || availableClassifier.isFromStarOrPackageImport) {
                 addElementToImportAndShorten(mostTopLevelClassId.asSingleFqName(), mostTopLevelQualifier)
+            } else {
+                addFakePackagePrefixToShortenIfPresent(mostTopLevelQualifier)
+            }
+        }
+
+        private fun addFakePackagePrefixToShortenIfPresent(wholeQualifiedExpression: KtDotQualifiedExpression) {
+            val deepestQualifier = wholeQualifiedExpression.qualifiersWithSelf.last()
+            if (deepestQualifier.hasFakeRootPrefix()) {
+                addElementToShorten(deepestQualifier)
             }
         }
 
@@ -395,6 +422,12 @@ private class ShortenCommandImpl(
     }
 }
 
+private fun KtUserType.hasFakeRootPrefix(): Boolean =
+    qualifier?.referencedName == ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE
+
+private fun KtDotQualifiedExpression.hasFakeRootPrefix(): Boolean =
+    (receiverExpression as? KtNameReferenceExpression)?.getReferencedName() == ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE
+
 private fun CallableId.asImportableFqName(): FqName? = if (classId == null) packageName.child(callableName) else null
 
 private fun KtElement.getDotQualifiedExpressionForSelector(): KtDotQualifiedExpression? =

idea/testData/shortenRefsFir/calls/rootPackageShortenFakeRootPackage.kt

@@ -1,10 +0,0 @@
-// FIR_IGNORE
-package test
-
-fun test() {}
-
-fun usage() {
-    fun test() {}
-
-    <selection>_root_ide_package_.test.test()</selection>
-}

idea/testData/shortenRefsFir/fakeRootPackage/rootPackageShortenFakeRootPackage.kt

@@ -0,0 +1,18 @@
+// FIR_COMPARISON
+package test
+
+class Test
+
+fun test() {}
+
+fun usage() {
+    class Test
+
+    fun test() {}
+
+    <selection>
+    _root_ide_package_.test.test()
+    _root_ide_package_.test.Test
+    val t: _root_ide_package_.test.Test
+    </selection>
+}

idea/testData/shortenRefsFir/fakeRootPackage/rootPackageShortenFakeRootPackage.kt.after

@@ -1,10 +1,18 @@
 // FIR_COMPARISON
 package test
 
+class Test
+
 fun test() {}
 
 fun usage() {
+    class Test
+
     fun test() {}
 
+
     test.test()
+    test.Test
+    val t: test.Test
+
 }