FIR IDE: Create fake scopes to avoid import duplicates

2020-12-22 16:30:37 +03:00
parent e34370554d
commit 08e271411f
11 changed files with 95 additions and 5 deletions
@@ -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");
@@ -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)
@@ -0,0 +1,5 @@
+package dependency
+
+class T {
+    class TT
+}
@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+import dependency.T.TT
+
+<selection>fun foo(p: dependency.T.TT) {}</selection>
@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+import dependency.T.TT
+
+fun foo(p: TT) {}
@@ -0,0 +1,3 @@
+package dependency
+
+class T
@@ -0,0 +1,4 @@
+// FIR_COMPARISON
+package test
+
+<selection>fun foo(p1: dependency.T, p2: dependency.T) {}</selection>
@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+import dependency.T
+
+fun foo(p1: T, p2: T) {}
@@ -0,0 +1,5 @@
+package dependency
+
+class T {
+    class TT
+}
@@ -0,0 +1,4 @@
+// FIR_COMPARISON
+package test
+
+<selection>fun foo(p: dependency.T.TT) {}</selection>
@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+import dependency.T
+
+fun foo(p: T.TT) {}