KT-60940 [Analysis API] Properly handle vararg types in KtFirReferenceShortener

^KT-60940 Fixed
^KTIJ-26518 Fixed

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

@@ -9,6 +9,7 @@ import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.SmartPsiElementPointer
+import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.analysis.api.components.KtReferenceShortener
 import org.jetbrains.kotlin.analysis.api.components.ShortenCommand
 import org.jetbrains.kotlin.analysis.api.components.ShortenOption
@@ -416,12 +417,7 @@ private class ElementsToShortenCollector(
     }
 
     private fun processTypeRef(resolvedTypeRef: FirResolvedTypeRef) {
-        val typeElement = when (val realPsi = resolvedTypeRef.realPsi) {
+        val typeElement = resolvedTypeRef.correspondingTypePsi ?: return
-            is KtTypeReference -> realPsi.typeElement
-            is KtNameReferenceExpression -> realPsi.parent as? KtTypeElement
-            else -> null
-        }?.unwrapNullability() as? KtUserType ?: return
-
         if (typeElement.referenceExpression?.textRange?.intersects(selection) != true) return
         if (typeElement.qualifier == null) return
 
@@ -430,6 +426,36 @@ private class ElementsToShortenCollector(
         findTypeToShorten(classifierId, typeElement)?.let(::addElementToShorten)
     }
 
+    /**
+     * Retrieves the corresponding [KtUserType] PSI the given [FirResolvedTypeRef].
+     *
+     * This code handles some quirks of FIR sources and PSI:
+     * - in `vararg args: String` declaration, `String` type reference has fake source, but `Array<String>` has real source
+     * (see [KtFakeSourceElementKind.ArrayTypeFromVarargParameter]).
+     * - if FIR reference points to the type with generic parameters (like `Foo<Bar>`), its source is not [KtTypeReference], but
+     * [KtNameReferenceExpression].
+     */
+    private val FirResolvedTypeRef.correspondingTypePsi: KtUserType?
+        get() {
+            val sourcePsi = when {
+                // array type for vararg parameters is not present in the code, so no need to handle it
+                delegatedTypeRef?.source?.kind == KtFakeSourceElementKind.ArrayTypeFromVarargParameter -> null
+
+                // but the array's underlying type is present with a fake source, and needs to be handled
+                source?.kind == KtFakeSourceElementKind.ArrayTypeFromVarargParameter -> psi
+
+                else -> realPsi
+            }
+
+            val outerTypeElement = when (sourcePsi) {
+                is KtTypeReference -> sourcePsi.typeElement
+                is KtNameReferenceExpression -> sourcePsi.parent as? KtTypeElement
+                else -> null
+            }
+
+            return outerTypeElement?.unwrapNullability() as? KtUserType
+        }
+
     val ConeKotlinType.candidateClassId: ClassId?
         get() {
             return when (this) {

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

@@ -418,6 +418,12 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/typeParams2.kt");
     }
 
+    @Test
+    @TestMetadata("vararg.kt")
+    public void testVararg() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/vararg.kt");
+    }
+
     @Test
     @TestMetadata("variable.kt")
     public void testVariable() throws Exception {

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

@@ -418,6 +418,12 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenera
         runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/typeParams2.kt");
     }
 
+    @Test
+    @TestMetadata("vararg.kt")
+    public void testVararg() throws Exception {
+        runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/vararg.kt");
+    }
+
     @Test
     @TestMetadata("variable.kt")
     public void testVariable() throws Exception {

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/vararg.kt

@@ -0,0 +1,7 @@
+// FILE: main.kt
+fun test(vararg args: <expr>dependency.Foo</expr>) {}
+
+// FILE: dependency.kt
+package dependency
+
+class Foo

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/vararg.txt

@@ -0,0 +1,7 @@
+Before shortening: dependency.Foo
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] dependency.Foo
+with SHORTEN_AND_STAR_IMPORT:
+[type] dependency.Foo