diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt
index 808afc23849..5bdca67275e 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt
@@ -380,19 +380,20 @@ private class FirShorteningContext(val analysisSession: KtFirAnalysisSession) {
 }
 
 private sealed class ElementToShorten {
+    abstract val element: KtElement
     abstract val nameToImport: FqName?
     abstract val importAllInParent: Boolean
 }
 
 private class ShortenType(
-    val element: KtUserType,
+    override val element: KtUserType,
     val shortenedRef: String? = null,
     override val nameToImport: FqName? = null,
     override val importAllInParent: Boolean = false,
 ) : ElementToShorten()
 
 private class ShortenQualifier(
-    val element: KtDotQualifiedExpression,
+    override val element: KtDotQualifiedExpression,
     val shortenedRef: String? = null,
     override val nameToImport: FqName? = null,
     override val importAllInParent: Boolean = false
@@ -845,10 +846,16 @@ private class ElementsToShortenCollector(
 
     private fun createElementToShorten(
         element: KtElement,
-        nameToImport: FqName? = null,
+        referencedSymbol: FqName? = null,
         importAllInParent: Boolean = false,
         shortenedRef: String? = null,
     ): ElementToShorten {
+        var nameToImport = if (importAllInParent) {
+            referencedSymbol?.parentOrNull() ?: error("Provided FqName '$referencedSymbol' cannot be imported with a star")
+        } else {
+            referencedSymbol
+        }
+
         return when (element) {
             is KtUserType -> ShortenType(element, shortenedRef, nameToImport, importAllInParent)
             is KtDotQualifiedExpression -> ShortenQualifier(element, shortenedRef, nameToImport, importAllInParent)
@@ -1232,37 +1239,17 @@ private class ElementsToShortenCollector(
         }
     }
 
-    private fun addElementToShorten(element: KtElement, shortenedRef: String?, nameToImport: FqName?, isImportWithStar: Boolean) {
-        val qualifier = element.getQualifier() ?: return
+    private fun addElementToShorten(elementInfoToShorten: ElementToShorten) {
+        val qualifier = elementInfoToShorten.element.getQualifier() ?: return
         if (!qualifier.isAlreadyCollected()) {
             removeRedundantElements(qualifier)
-            when (element) {
-                is KtUserType -> typesToShorten.add(ShortenType(element, shortenedRef, nameToImport, isImportWithStar))
-                is KtDotQualifiedExpression -> qualifiersToShorten.add(
-                    ShortenQualifier(
-                        element, shortenedRef, nameToImport, isImportWithStar
-                    )
-                )
+            when (elementInfoToShorten) {
+                is ShortenType -> typesToShorten.add(elementInfoToShorten)
+                is ShortenQualifier -> qualifiersToShorten.add(elementInfoToShorten)
             }
         }
     }
 
-    private fun addElementToShorten(elementInfoToShorten: ElementToShorten) {
-        val (nameToImport, isImportWithStar) = if (elementInfoToShorten.importAllInParent && elementInfoToShorten.nameToImport?.parentOrNull()?.isRoot == false) {
-            elementInfoToShorten.nameToImport?.parent() to true
-        } else {
-            elementInfoToShorten.nameToImport to false
-        }
-        when (elementInfoToShorten) {
-            is ShortenType -> addElementToShorten(
-                elementInfoToShorten.element, elementInfoToShorten.shortenedRef, nameToImport, isImportWithStar
-            )
-            is ShortenQualifier -> addElementToShorten(
-                elementInfoToShorten.element, elementInfoToShorten.shortenedRef, nameToImport, isImportWithStar
-            )
-        }
-    }
-
     /**
      * Checks whether type reference of [this] type is considered to be in the [selection] text range.
      *
@@ -1421,5 +1408,5 @@ internal fun KtSimpleNameExpression.getDotQualifiedExpressionForSelector(): KtDo
 private fun KtElement.getQualifier(): KtElement? = when (this) {
     is KtUserType -> qualifier
     is KtDotQualifiedExpression -> receiverExpression
-    else -> null
+    else -> error("Unexpected ${this::class}")
 }