KTIJ-28080 [AA] Move selection checking to ElementsToShortenCollector.findClassifierQualifierToShorten

This way, even if the whole qualifier is not selected, but some of its
parts are, the reference shortener will correctly find and shorten
those parts

^KTIJ-28080 Fixed

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

@@ -524,8 +524,6 @@ private class ElementsToShortenCollector(
         val typeElement = resolvedTypeRef.correspondingTypePsi ?: return
         if (typeElement.qualifier == null) return
 
-        if (!typeElement.inSelection) return
-
         val classifierId = resolvedTypeRef.type.lowerBoundIfFlexible().candidateClassId ?: return
 
         findTypeQualifierToShorten(classifierId, typeElement)?.let(::addElementToShorten)
@@ -597,8 +595,6 @@ private class ElementsToShortenCollector(
             else -> return
         }
 
-        if (!wholeQualifierElement.inSelection) return
-
         findTypeQualifierToShorten(wholeClassQualifier, wholeQualifierElement)?.let(::addElementToShorten)
     }
 
@@ -829,6 +825,10 @@ private class ElementsToShortenCollector(
     /**
      * Finds the longest qualifier in [wholeQualifierElement] which can be safely shortened in the [positionScopes].
      * [wholeQualifierClassId] is supposed to reflect the class which is referenced by the [wholeQualifierElement].
+     *
+     * N.B. Even if the [wholeQualifierElement] is not strictly in the [selection],
+     * some outer part of it might be, and we want to shorten that.
+     * So we have to check all the outer qualifiers.
      */
     private fun findClassifierQualifierToShorten(
         positionScopes: List<FirScope>,
@@ -839,6 +839,8 @@ private class ElementsToShortenCollector(
         val allQualifiedElements = wholeQualifierElement.qualifiedElementsWithSelf
 
         for ((classId, element) in allClassIds.zip(allQualifiedElements)) {
+            if (!element.inSelection) continue
+
             val classSymbol = shorteningContext.toClassSymbol(classId) ?: return null
             val option = classShortenStrategy(classSymbol)
             if (option == ShortenStrategy.DO_NOT_SHORTEN) continue
@@ -885,7 +887,11 @@ private class ElementsToShortenCollector(
                 }
             }
         }
-        return findFakePackageToShorten(allQualifiedElements.last())
+
+        val lastQualifier = allQualifiedElements.last()
+        if (!lastQualifier.inSelection) return null
+
+        return findFakePackageToShorten(lastQualifier)
     }
 
     private fun createElementToShorten(
@@ -1361,6 +1367,15 @@ private class ElementsToShortenCollector(
         }
     }
 
+    private val KtElement.inSelection: Boolean
+        get() = when (this) {
+            is KtUserType -> inSelection
+            is KtDotQualifiedExpression -> inSelection
+            is KtThisExpression -> inSelection
+
+            else -> error("Unexpected ${this::class}")
+        }
+
     /**
      * Checks whether type reference of [this] type is considered to be in the [selection] text range.
      *

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/references/FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated.java

@@ -418,6 +418,18 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
         runTest("analysis/analysis-api/testData/components/referenceShortener/shortenRange/partiallySelectedType2.kt");
     }
 
+    @Test
+    @TestMetadata("partiallySelectedType3.kt")
+    public void testPartiallySelectedType3() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/shortenRange/partiallySelectedType3.kt");
+    }
+
+    @Test
+    @TestMetadata("partiallySelectedType4.kt")
+    public void testPartiallySelectedType4() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/shortenRange/partiallySelectedType4.kt");
+    }
+
     @Test
     @TestMetadata("partiallySelectedTypeQualifier1.kt")
     public void testPartiallySelectedTypeQualifier1() throws Exception {
@@ -430,6 +442,18 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
         runTest("analysis/analysis-api/testData/components/referenceShortener/shortenRange/partiallySelectedTypeQualifier2.kt");
     }
 
+    @Test
+    @TestMetadata("partiallySelectedTypeQualifier3.kt")
+    public void testPartiallySelectedTypeQualifier3() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/shortenRange/partiallySelectedTypeQualifier3.kt");
+    }
+
+    @Test
+    @TestMetadata("partiallySelectedTypeQualifier4.kt")
+    public void testPartiallySelectedTypeQualifier4() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/shortenRange/partiallySelectedTypeQualifier4.kt");
+    }
+
     @Test
     @TestMetadata("qualifierOfUnresolvedReference.kt")
     public void testQualifierOfUnresolvedReference() throws Exception {
@@ -628,6 +652,12 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
         runTest("analysis/analysis-api/testData/components/referenceShortener/shortenRange/variable_invokeOperator.kt");
     }
 
+    @Test
+    @TestMetadata("_rootIdePackage_IsNotRemovedIfNotSelected.kt")
+    public void test_rootIdePackage_IsNotRemovedIfNotSelected() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/shortenRange/_rootIdePackage_IsNotRemovedIfNotSelected.kt");
+    }
+
     @Nested
     @TestMetadata("analysis/analysis-api/testData/components/referenceShortener/shortenRange/nestedClasses")
     @TestDataPath("$PROJECT_ROOT")