FIR IDE: Decouple ElementsToShortenCollector from KtFirReferenceShortener

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

@@ -11,6 +11,7 @@ import com.intellij.psi.SmartPsiElementPointer
 import com.intellij.util.containers.addIfNotNull
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE
 import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClass
 import org.jetbrains.kotlin.fir.declarations.*
@@ -37,6 +38,7 @@ import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.lowerBoundIfFlexible
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
@@ -60,11 +62,13 @@ internal class KtFirReferenceShortener(
     override val token: ValidityToken,
     override val firResolveState: FirModuleResolveState,
 ) : KtReferenceShortener(), KtFirAnalysisSessionComponent {
+    private val context = FirShorteningContext(firResolveState)
+
     override fun collectShortenings(file: KtFile, selection: TextRange): ShortenCommand {
         resolveFileToBodyResolve(file)
         val firFile = file.getOrBuildFirOfType<FirFile>(firResolveState)
 
-        val collector = ElementsToShortenCollector()
+        val collector = ElementsToShortenCollector(context)
         firFile.acceptChildren(collector)
 
         return ShortenCommandImpl(
@@ -75,9 +79,21 @@ internal class KtFirReferenceShortener(
         )
     }
 
+    private fun resolveFileToBodyResolve(file: KtFile) {
+        for (declaration in file.declarations) {
+            declaration.getOrBuildFir(firResolveState) // temporary hack, resolves declaration to BODY_RESOLVE stage
+        }
+    }
+}
+
 private data class AvailableClassifier(val classId: ClassId, val isFromStarOrPackageImport: Boolean)
 
-    private fun findFirstClassifierInScopesByName(positionScopes: List<FirScope>, targetClassName: Name): AvailableClassifier? {
+private class FirShorteningContext(val firResolveState: FirModuleResolveState) {
+
+    private val firSession: FirSession
+        get() = firResolveState.rootModuleSession
+
+    fun findFirstClassifierInScopesByName(positionScopes: List<FirScope>, targetClassName: Name): AvailableClassifier? {
         for (scope in positionScopes) {
             val classifierSymbol = scope.findFirstClassifierByName(targetClassName) ?: continue
             val classifierLookupTag = classifierSymbol.toLookupTag() as? ConeClassLikeLookupTag ?: continue
@@ -91,20 +107,14 @@ internal class KtFirReferenceShortener(
         return null
     }
 
-    private fun findFunctionsInScopes(scopes: List<FirScope>, name: Name): List<FirNamedFunctionSymbol> {
+    fun findFunctionsInScopes(scopes: List<FirScope>, name: Name): List<FirNamedFunctionSymbol> {
         return scopes.flatMap { it.getFunctions(name) }
     }
 
-    private fun findSinglePropertyInScopesByName(scopes: List<FirScope>, name: Name): FirVariableSymbol<*>? {
+    fun findSinglePropertyInScopesByName(scopes: List<FirScope>, name: Name): FirVariableSymbol<*>? {
         return scopes.asSequence().mapNotNull { it.getSinglePropertyByName(name) }.singleOrNull()
     }
 
-    private fun resolveFileToBodyResolve(file: KtFile) {
-        for (declaration in file.declarations) {
-            declaration.getOrBuildFir(firResolveState) // temporary hack, resolves declaration to BODY_RESOLVE stage
-        }
-    }
-
     private fun FirScope.findFirstClassifierByName(name: Name): FirClassifierSymbol<*>? {
         var element: FirClassifierSymbol<*>? = null
 
@@ -117,16 +127,12 @@ internal class KtFirReferenceShortener(
         return element
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
-    private fun FirScope.getSingleFunctionByName(name: Name): FirNamedFunctionSymbol? =
-        buildList { processFunctionsByName(name, this::add) }.singleOrNull()
-
     @OptIn(ExperimentalStdlibApi::class)
     private fun FirScope.getSinglePropertyByName(name: Name): FirVariableSymbol<*>? =
         buildList { processPropertiesByName(name, this::add) }.singleOrNull()
 
     @OptIn(ExperimentalStdlibApi::class)
-    private fun findScopesAtPosition(position: KtElement, newImports: List<FqName>): List<FirScope>? {
+    fun findScopesAtPosition(position: KtElement, newImports: List<FqName>): List<FirScope>? {
         val towerDataContext = firResolveState.getTowerDataContextForElement(position) ?: return null
 
         val result = buildList<FirScope> {
@@ -142,11 +148,11 @@ internal class KtFirReferenceShortener(
         val resolvedNewImports = newImports.mapNotNull { createFakeResolvedImport(it) }
         if (resolvedNewImports.isEmpty()) return null
 
-        return FirExplicitSimpleImportingScope(resolvedNewImports, firResolveState.rootModuleSession, ScopeSession())
+        return FirExplicitSimpleImportingScope(resolvedNewImports, firSession, ScopeSession())
     }
 
     private fun createFakeResolvedImport(fqNameToImport: FqName): FirResolvedImport? {
-        val packageOrClass = resolveToPackageOrClass(firResolveState.rootModuleSession.firSymbolProvider, fqNameToImport) ?: return null
+        val packageOrClass = resolveToPackageOrClass(firSession.firSymbolProvider, fqNameToImport) ?: return null
 
         val delegateImport = buildImport {
             importedFqName = fqNameToImport
@@ -159,7 +165,12 @@ internal class KtFirReferenceShortener(
         }
     }
 
-    private inner class ElementsToShortenCollector : FirVisitorVoid() {
+    fun getRegularClass(typeRef: FirTypeRef): FirRegularClass? {
+        return typeRef.toRegularClass(firSession)
+    }
+}
+
+private class ElementsToShortenCollector(private val shorteningContext: FirShorteningContext) : FirVisitorVoid() {
     val namesToImport: MutableList<FqName> = mutableListOf()
     val typesToShorten: MutableList<KtUserType> = mutableListOf()
     val qualifiersToShorten: MutableList<KtDotQualifiedExpression> = mutableListOf()
@@ -208,13 +219,13 @@ internal class KtFirReferenceShortener(
         val allClassIds = wholeClassifierId.outerClassesWithSelf
         val allTypeElements = wholeTypeElement.qualifiersWithSelf
 
-            val positionScopes = findScopesAtPosition(wholeTypeElement, namesToImport) ?: return
+        val positionScopes = shorteningContext.findScopesAtPosition(wholeTypeElement, namesToImport) ?: return
 
         for ((classId, typeElement) in allClassIds.zip(allTypeElements)) {
             // if qualifier is null, then this type have no package and thus cannot be shortened
             if (typeElement.qualifier == null) return
 
-                val firstFoundClass = findFirstClassifierInScopesByName(positionScopes, classId.shortClassName)?.classId
+            val firstFoundClass = shorteningContext.findFirstClassifierInScopesByName(positionScopes, classId.shortClassName)?.classId
 
             if (firstFoundClass == classId) {
                 addTypeToShorten(typeElement)
@@ -224,7 +235,7 @@ internal class KtFirReferenceShortener(
 
         // none class matched
         val (mostTopLevelClassId, mostTopLevelTypeElement) = allClassIds.zip(allTypeElements).last()
-            val availableClassifier = findFirstClassifierInScopesByName(positionScopes, mostTopLevelClassId.shortClassName)
+        val availableClassifier = shorteningContext.findFirstClassifierInScopesByName(positionScopes, mostTopLevelClassId.shortClassName)
 
         check(availableClassifier?.classId != mostTopLevelClassId) {
             "If this condition were true, we would have exited from the loop on the last iteration. ClassId = $mostTopLevelClassId"
@@ -267,13 +278,13 @@ internal class KtFirReferenceShortener(
     }
 
     private fun collectQualifierIfNeedsToBeShortened(wholeClassQualifier: ClassId, wholeQualifierElement: KtDotQualifiedExpression) {
-            val positionScopes = findScopesAtPosition(wholeQualifierElement, namesToImport) ?: return
+        val positionScopes = shorteningContext.findScopesAtPosition(wholeQualifierElement, namesToImport) ?: return
 
         val allClassIds = wholeClassQualifier.outerClassesWithSelf
         val allQualifiers = wholeQualifierElement.qualifiersWithSelf
 
         for ((classId, qualifier) in allClassIds.zip(allQualifiers)) {
-                val firstFoundClass = findFirstClassifierInScopesByName(positionScopes, classId.shortClassName)?.classId
+            val firstFoundClass = shorteningContext.findFirstClassifierInScopesByName(positionScopes, classId.shortClassName)?.classId
 
             if (firstFoundClass == classId) {
                 addElementToShorten(qualifier)
@@ -282,7 +293,7 @@ internal class KtFirReferenceShortener(
         }
 
         val (mostTopLevelClassId, mostTopLevelQualifier) = allClassIds.zip(allQualifiers).last()
-            val availableClassifier = findFirstClassifierInScopesByName(positionScopes, mostTopLevelClassId.shortClassName)
+        val availableClassifier = shorteningContext.findFirstClassifierInScopesByName(positionScopes, mostTopLevelClassId.shortClassName)
 
         check(availableClassifier?.classId != mostTopLevelClassId) {
             "If this condition were true, we would have exited from the loop on the last iteration. ClassId = $mostTopLevelClassId"
@@ -301,8 +312,8 @@ internal class KtFirReferenceShortener(
 
         val propertyId = (resolvedNamedReference.resolvedSymbol as? FirCallableSymbol<*>)?.callableId ?: return
 
-            val scopes = findScopesAtPosition(qualifiedProperty, namesToImport) ?: return
+        val scopes = shorteningContext.findScopesAtPosition(qualifiedProperty, namesToImport) ?: return
-            val singleAvailableProperty = findSinglePropertyInScopesByName(scopes, propertyId.callableName)
+        val singleAvailableProperty = shorteningContext.findSinglePropertyInScopesByName(scopes, propertyId.callableName)
 
         if (singleAvailableProperty?.callableId == propertyId) {
             addElementToShorten(qualifiedProperty)
@@ -318,8 +329,8 @@ internal class KtFirReferenceShortener(
         val calleeReference = functionCall.calleeReference
         val callableId = findUnambiguousReferencedCallableId(calleeReference) ?: return
 
-            val scopes = findScopesAtPosition(callExpression, namesToImport) ?: return
+        val scopes = shorteningContext.findScopesAtPosition(callExpression, namesToImport) ?: return
-            val availableCallables = findFunctionsInScopes(scopes, callableId.callableName)
+        val availableCallables = shorteningContext.findFunctionsInScopes(scopes, callableId.callableName)
 
         when {
             availableCallables.isEmpty() -> {
@@ -342,7 +353,7 @@ internal class KtFirReferenceShortener(
         // if there is no extension receiver necessary, then it can be removed
         if (functionCall.extensionReceiver is FirNoReceiverExpression) return true
 
-            val receiverType = explicitReceiver.typeRef.toRegularClass(firResolveState.rootModuleSession) ?: return true
+        val receiverType = shorteningContext.getRegularClass(explicitReceiver.typeRef) ?: return true
         return receiverType.classKind != ClassKind.OBJECT
     }
 
@@ -403,7 +414,6 @@ internal class KtFirReferenceShortener(
     private val KtDotQualifiedExpression.qualifiersWithSelf: Sequence<KtDotQualifiedExpression>
         get() = generateSequence(this) { it.receiverExpression as? KtDotQualifiedExpression }
 }
-}
 
 private class ShortenCommandImpl(
     val targetFile: KtFile,