FIR IDE: Create fake scopes to avoid import duplicates

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

@@ -66,6 +66,11 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
             runTest("idea/testData/shortenRefsFir/types/ParameterTypeGenericTypes.kt");
         }
 
+        @TestMetadata("ParameterTypeImportedNestedClass.kt")
+        public void testParameterTypeImportedNestedClass() throws Exception {
+            runTest("idea/testData/shortenRefsFir/types/ParameterTypeImportedNestedClass.kt");
+        }
+
         @TestMetadata("ParameterTypeImportedTypeWins.kt")
         public void testParameterTypeImportedTypeWins() throws Exception {
             runTest("idea/testData/shortenRefsFir/types/ParameterTypeImportedTypeWins.kt");
@@ -81,6 +86,16 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
             runTest("idea/testData/shortenRefsFir/types/ParameterTypeNonImportedClass.kt");
         }
 
+        @TestMetadata("ParameterTypeNonImportedClassTwice.kt")
+        public void testParameterTypeNonImportedClassTwice() throws Exception {
+            runTest("idea/testData/shortenRefsFir/types/ParameterTypeNonImportedClassTwice.kt");
+        }
+
+        @TestMetadata("ParameterTypeNotImportedNestedClass.kt")
+        public void testParameterTypeNotImportedNestedClass() throws Exception {
+            runTest("idea/testData/shortenRefsFir/types/ParameterTypeNotImportedNestedClass.kt");
+        }
+
         @TestMetadata("ParameterTypeStarImportedTypeLoses.kt")
         public void testParameterTypeStarImportedTypeLoses() throws Exception {
             runTest("idea/testData/shortenRefsFir/types/ParameterTypeStarImportedTypeLoses.kt");

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

@@ -8,10 +8,14 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.components
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.SmartPsiElementPointer
+import com.intellij.util.containers.addIfNotNull
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
 import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.scopes.FirScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirExplicitSimpleImportingScope
 import org.jetbrains.kotlin.fir.scopes.processClassifiersByName
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassifierSymbol
@@ -21,9 +25,10 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
-import org.jetbrains.kotlin.idea.frontend.api.components.ShortenCommand
 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.idea.frontend.api.fir.utils.addImportToFile
 import org.jetbrains.kotlin.name.ClassId
@@ -31,6 +36,8 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.createSmartPointer
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
+import org.jetbrains.kotlin.resolve.ImportPath
 
 internal class KtFirReferenceShortener(
     override val analysisSession: KtFirAnalysisSession,
@@ -78,11 +85,34 @@ internal class KtFirReferenceShortener(
         return element
     }
 
-    private fun findScopesAtPosition(targetTypeReference: KtElement): List<FirScope>? {
+    @OptIn(ExperimentalStdlibApi::class)
+    private fun findScopesAtPosition(targetTypeReference: KtElement, newImports: List<FqName>): List<FirScope>? {
         val towerDataContext = firResolveState.getTowerDataContextForElement(targetTypeReference) ?: return null
-        val availableScopes = towerDataContext.towerDataElements.mapNotNull { it.scope }
 
-        return availableScopes.asReversed()
+        val result = buildList<FirScope> {
+            addAll(towerDataContext.nonLocalTowerDataElements.mapNotNull { it.scope })
+            addIfNotNull(createFakeImportingScope(targetTypeReference.project, newImports))
+            addAll(towerDataContext.localScopes)
+        }
+
+        return result.asReversed()
+    }
+
+    private fun createFakeImportingScope(
+        project: Project,
+        newImports: List<FqName>
+    ): FirScope? {
+        if (newImports.isEmpty()) return null
+
+        val psiFactory = KtPsiFactory(project)
+
+        val resolvedNewImports = newImports
+            .map { psiFactory.createImportDirective(ImportPath(it, isAllUnder = false)) }
+            .mapNotNull { it.getOrBuildFirSafe<FirResolvedImport>(firResolveState) }
+
+        if (resolvedNewImports.isEmpty()) return null
+
+        return FirExplicitSimpleImportingScope(resolvedNewImports, firResolveState.rootModuleSession, ScopeSession())
     }
 
     private inner class TypesCollectingVisitor(
@@ -115,7 +145,7 @@ internal class KtFirReferenceShortener(
             val allClassIds = generateSequence(wholeClassifierId) { it.outerClassId }
             val allTypeElements = generateSequence(wholeTypeElement) { it.qualifier }
 
-            val positionScopes = findScopesAtPosition(wholeTypeElement) ?: return
+            val positionScopes = findScopesAtPosition(wholeTypeElement, typesToImport) ?: return
 
             for ((classId, typeElement) in allClassIds.zip(allTypeElements)) {
                 val firstFoundClass = findFirstClassifierInScopesByName(positionScopes, classId.shortClassName)

idea/testData/shortenRefsFir/types/ParameterTypeImportedNestedClass.dependency.kt

@@ -0,0 +1,5 @@
+package dependency
+
+class T {
+    class TT
+}

idea/testData/shortenRefsFir/types/ParameterTypeImportedNestedClass.kt

@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+import dependency.T.TT
+
+<selection>fun foo(p: dependency.T.TT) {}</selection>

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

@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+import dependency.T.TT
+
+fun foo(p: TT) {}

idea/testData/shortenRefsFir/types/ParameterTypeNonImportedClassTwice.dependency.kt

@@ -0,0 +1,3 @@
+package dependency
+
+class T

idea/testData/shortenRefsFir/types/ParameterTypeNonImportedClassTwice.kt

@@ -0,0 +1,4 @@
+// FIR_COMPARISON
+package test
+
+<selection>fun foo(p1: dependency.T, p2: dependency.T) {}</selection>

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

@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+import dependency.T
+
+fun foo(p1: T, p2: T) {}

idea/testData/shortenRefsFir/types/ParameterTypeNotImportedNestedClass.dependency.kt

@@ -0,0 +1,5 @@
+package dependency
+
+class T {
+    class TT
+}

idea/testData/shortenRefsFir/types/ParameterTypeNotImportedNestedClass.kt

@@ -0,0 +1,4 @@
+// FIR_COMPARISON
+package test
+
+<selection>fun foo(p: dependency.T.TT) {}</selection>

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

@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+import dependency.T
+
+fun foo(p: T.TT) {}