FIR IDE: Enable types shortening for nested classes and nested types

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

@@ -45,5 +45,20 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
         public void testParameterType() throws Exception {
             runTest("idea/testData/shortenRefsFir/types/ParameterType.kt");
         }
+
+        @TestMetadata("ParameterTypeFunctionalType.kt")
+        public void testParameterTypeFunctionalType() throws Exception {
+            runTest("idea/testData/shortenRefsFir/types/ParameterTypeFunctionalType.kt");
+        }
+
+        @TestMetadata("ParameterTypeGenericTypes.kt")
+        public void testParameterTypeGenericTypes() throws Exception {
+            runTest("idea/testData/shortenRefsFir/types/ParameterTypeGenericTypes.kt");
+        }
+
+        @TestMetadata("ParameterTypeNestedType.kt")
+        public void testParameterTypeNestedType() throws Exception {
+            runTest("idea/testData/shortenRefsFir/types/ParameterTypeNestedType.kt");
+        }
     }
 }

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.components.KtReferenceShortener
 import org.jetbrains.kotlin.idea.frontend.api.components.ShortenCommand
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtTypeReference
@@ -45,28 +46,27 @@ internal class KtFirReferenceShortener(
             }
 
             override fun visitResolvedTypeRef(resolvedTypeRef: FirResolvedTypeRef) {
-                val targetTypeReference = resolvedTypeRef.psi as? KtTypeReference ?: return
-                val targetType = targetTypeReference.typeElement as? KtUserType ?: return
+                resolvedTypeRef.acceptChildren(this)
 
-                if (targetType.qualifier == null) return
+                val wholeTypeReference = resolvedTypeRef.psi as? KtTypeReference ?: return
 
-                val targetClassId = resolvedTypeRef.type.classId
-                val targetClassName = targetClassId?.shortClassName ?: return
+                val wholeTypeElement = wholeTypeReference.typeElement as? KtUserType ?: return
+                if (wholeTypeElement.qualifier == null) return
 
-                val positionScopes = findScopesAtPosition(targetTypeReference) ?: return
+                val wholeClassifierId = resolvedTypeRef.type.classId ?: return
 
-                val firstFoundClass = positionScopes.asSequence()
-                    .mapNotNull { scope -> scope.findFirstClassifierByName(targetClassName) }
-                    .mapNotNull { classifierSymbol -> classifierSymbol.toLookupTag() as? ConeClassLikeLookupTag }
-                    .map { it.classId }
-                    .firstOrNull()
+                val allClassIds = generateSequence(wholeClassifierId) { it.outerClassId }
+                val allTypeElements = generateSequence(wholeTypeElement) { it.qualifier }
 
-                if (firstFoundClass == null) {
-                    // this class should be imported
-                }
+                val positionScopes = findScopesAtPosition(wholeTypeReference) ?: return
 
-                if (firstFoundClass == targetClassId) {
-                    typesToShorten.add(targetType)
+                for ((classId, typeElement) in allClassIds.zip(allTypeElements)) {
+                    val firstFoundClass = findFirstClassifierInScopesByName(positionScopes, classId.shortClassName)
+
+                    if (firstFoundClass == classId) {
+                        typesToShorten.add(typeElement)
+                        break
+                    }
                 }
             }
         })
@@ -74,6 +74,18 @@ internal class KtFirReferenceShortener(
         return ShortenCommand(file, emptyList(), typesToShorten.map { it.createSmartPointer() })
     }
 
+    private fun findFirstClassifierInScopesByName(positionScopes: List<FirScope>, targetClassName: Name): ClassId? {
+        for (scope in positionScopes) {
+            val classifierSymbol = scope.findFirstClassifierByName(targetClassName) ?: continue
+            val classifierLookupTag = classifierSymbol.toLookupTag() as? ConeClassLikeLookupTag ?: continue
+
+            return classifierLookupTag.classId
+        }
+
+        return null
+    }
+
+
     private fun resolveFileToBodyResolve(file: KtFile) {
         for (declaration in file.declarations) {
             declaration.getOrBuildFir(firResolveState) // temporary hack, resolves declaration to BODY_RESOLVE stage
@@ -81,8 +93,17 @@ internal class KtFirReferenceShortener(
     }
 
     @OptIn(ExperimentalStdlibApi::class)
-    private fun FirScope.findFirstClassifierByName(name: Name): FirClassifierSymbol<*>? =
-        buildList { processClassifiersByName(name, this::add) }.firstOrNull()
+    private fun FirScope.findFirstClassifierByName(name: Name): FirClassifierSymbol<*>? {
+        var element: FirClassifierSymbol<*>? = null
+
+        processClassifiersByName(name) {
+            if (element == null) {
+                element = it
+            }
+        }
+
+        return element
+    }
 
     private fun findScopesAtPosition(targetTypeReference: KtTypeReference): List<FirScope>? {
         val towerDataContext = firResolveState.getTowerDataContextForElement(targetTypeReference) ?: return null

idea/testData/shortenRefsFir/types/ParameterTypeFunctionalType.kt

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+class T {
+    class TT
+}
+
+<selection>fun foo(t: (test.T) -> test.T.TT) {}</selection>

idea/testData/shortenRefsFir/types/ParameterTypeFunctionalType.kt.after

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+class T {
+    class TT
+}
+
+fun foo(t: (T) -> T.TT) {}

idea/testData/shortenRefsFir/types/ParameterTypeGenericTypes.kt

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+class T {
+    class TT
+}
+
+class Generic<T, TT>
+
+<selection>fun foo(t: test.Generic<test.T, test.T.TT>) {}</selection>

idea/testData/shortenRefsFir/types/ParameterTypeGenericTypes.kt.after

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+class T {
+    class TT
+}
+
+class Generic<T, TT>
+
+fun foo(t: Generic<T, T.TT>) {}

idea/testData/shortenRefsFir/types/ParameterTypeNestedType.kt

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+class T {
+    class TT
+}
+
+<selection>fun foo(t: test.T.TT) {}</selection>

idea/testData/shortenRefsFir/types/ParameterTypeNestedType.kt.after

@@ -0,0 +1,8 @@
+// FIR_COMPARISON
+package test
+
+class T {
+    class TT
+}
+
+fun foo(t: T.TT) {}