[FIR IDE] Implementation of KtUserType resolve

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateDepended.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.KtToFirMapping
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.containingKtFileIfAny
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
 import org.jetbrains.kotlin.psi.KtDeclaration
@@ -29,7 +30,7 @@ import org.jetbrains.kotlin.psi.KtLambdaExpression
 internal class FirModuleResolveStateDepended(
     private val originalState: FirModuleResolveStateImpl,
     val towerProviderBuiltUponElement: FirTowerContextProvider,
-    private val ktToFirMapping: Map<KtElement, FirElement>,
+    private val ktToFirMapping: KtToFirMapping?,
 ) : FirModuleResolveState() {
 
     override val project: Project get() = originalState.project
@@ -43,18 +44,14 @@ internal class FirModuleResolveStateDepended(
     override fun getOrBuildFirFor(element: KtElement): FirElement {
         val psi = originalState.elementBuilder.getPsiAsFirElementSource(element)
 
-        //TODO It return invalid elements for elements with invalid code, but try to return the most closest ones
+        ktToFirMapping?.getFirOfClosestParent(psi, this)?.let { return it }
-        var currentElement: PsiElement = psi
-        while (currentElement !is KtFile) {
-            ktToFirMapping[currentElement]?.let { return it }
-            currentElement = currentElement.parent
-        }
 
         return originalState.elementBuilder.getOrBuildFirFor(
-            element,
+            element = element,
-            originalState.firFileBuilder,
+            firFileBuilder = originalState.firFileBuilder,
-            originalState.rootModuleSession.cache,
+            moduleFileCache = originalState.rootModuleSession.cache,
-            fileStructureCache,
+            fileStructureCache = fileStructureCache,
+            state = this,
         )
     }
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateImpl.kt

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.util.FirElementFinder
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 
 internal class FirModuleResolveStateImpl(
     override val project: Project,
@@ -43,6 +44,7 @@ internal class FirModuleResolveStateImpl(
 ) : FirModuleResolveState() {
     override val rootModuleSession: FirIdeSourcesSession get() = sessionProvider.rootModuleSession
 
+
     /**
      * WARNING! This object contains scopes for all statements and declarations that were ever resolved.
      * It can grow unbounded if you never edit the files in the opened project.
@@ -57,7 +59,7 @@ internal class FirModuleResolveStateImpl(
         sessionProvider.getSession(moduleInfo)!!
 
     override fun getOrBuildFirFor(element: KtElement): FirElement =
-        elementBuilder.getOrBuildFirFor(element, firFileBuilder, rootModuleSession.cache, fileStructureCache)
+        elementBuilder.getOrBuildFirFor(element, firFileBuilder, rootModuleSession.cache, fileStructureCache, this)
 
     override fun getOrBuildFirFile(ktFile: KtFile): FirFile =
         firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache, lazyBodiesMode = false)

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/FirModuleResolveState.kt

@@ -10,15 +10,13 @@ import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
-import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtLambdaExpression
 
 abstract class FirModuleResolveState {
     abstract val project: Project

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/api/LowLevelFirApiFacadeForResolveOnAir.kt

@@ -10,15 +10,25 @@ import com.intellij.psi.util.parentOfType
 import com.intellij.psi.util.parentsOfType
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.asTowerDataElement
+import org.jetbrains.kotlin.fir.resolve.transformers.FirTypeResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
+import org.jetbrains.kotlin.fir.resolve.typeResolver
+import org.jetbrains.kotlin.fir.scopes.FirCompositeScope
 import org.jetbrains.kotlin.fir.scopes.createImportingScopes
+import org.jetbrains.kotlin.fir.toFirPsiSourceElement
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.toFirResolvedTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateDepended
@@ -27,8 +37,10 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FileTowerProv
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextAllElementsCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.KtToFirMapping
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.RawFirReplacement
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.buildFileFirAnnotation
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.buildFirUserTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
@@ -129,7 +141,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
             checkPCE = false
         )
 
-        val importingScopes = createImportingScopes(firFile, firFile.declarationSiteSession, ScopeSession(), useCaching = false)
+        val importingScopes = createImportingScopes(firFile, firFile.moduleData.session, ScopeSession(), useCaching = false)
         val fileScopeElements = importingScopes.map { it.asTowerDataElement(isLocal = false) }
         return FirTowerDataContext().addNonLocalTowerDataElements(fileScopeElements)
     }
@@ -147,7 +159,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
             ?: return FirModuleResolveStateDepended(
                 originalState,
                 FileTowerProvider(elementToAnalyze.containingKtFile, onAirGetTowerContextForFile(originalState, originalKtFile)),
-                emptyMap()
+                ktToFirMapping = null
             )
 
 
@@ -167,8 +179,8 @@ object LowLevelFirApiFacadeForResolveOnAir {
             collector = collector,
         )
 
-        val recordedMap = FirElementsRecorder.recordElementsFrom(copiedFirDeclaration, FirElementsRecorder())
+        val mapping = KtToFirMapping(copiedFirDeclaration, FirElementsRecorder())
-        return FirModuleResolveStateDepended(originalState, collector, recordedMap)
+        return FirModuleResolveStateDepended(originalState, collector, mapping)
     }
 
     private fun tryResolveAsFileAnnotation(
@@ -176,10 +188,11 @@ object LowLevelFirApiFacadeForResolveOnAir {
         state: FirModuleResolveStateImpl,
         replacement: RawFirReplacement,
         firFile: FirFile,
+        collector: FirTowerDataContextCollector? = null,
     ): FirAnnotationCall {
         val annotationCall = buildFileFirAnnotation(
-            firFile.declarationSiteSession,
+            firFile.moduleData.session,
-            firFile.declarationSiteSession.firIdeProvider.kotlinScopeProvider,
+            firFile.moduleData.session.firIdeProvider.kotlinScopeProvider,
             annotationEntry,
             replacement
         )
@@ -187,7 +200,8 @@ object LowLevelFirApiFacadeForResolveOnAir {
             firFile,
             listOf(annotationCall),
             state.rootModuleSession.cache,
-            ScopeSession()
+            ScopeSession(),
+            collector
         )
 
         return annotationCall
@@ -211,7 +225,8 @@ object LowLevelFirApiFacadeForResolveOnAir {
                     annotationEntry = annotationCall,
                     state = state,
                     replacement = replacement,
-                    firFile = originalFirFile
+                    firFile = originalFirFile,
+                    collector = collector,
                 )
             } else {
                 error("Cannot find enclosing declaration for ${replacement.from.getElementTextInContext()}")
@@ -228,7 +243,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
             firDeclarationToResolve = copiedFirDeclaration,
             moduleFileCache = state.rootModuleSession.cache,
             containerFirFile = originalFirFile,
-            provider = originalFirFile.declarationSiteSession.firIdeProvider,
+            provider = originalFirFile.moduleData.session.firIdeProvider,
             toPhase = FirResolvePhase.BODY_RESOLVE,
             checkPCE = true,
             isOnAirResolve = isOnAirResolve,
@@ -238,6 +253,28 @@ object LowLevelFirApiFacadeForResolveOnAir {
         return copiedFirDeclaration
     }
 
+    fun onAirResolveTypeInPlace(
+        place: KtElement,
+        typeReference: KtTypeReference,
+        state: FirModuleResolveState
+    ): FirResolvedTypeRef {
+        val context = state.getTowerContextProvider().getClosestAvailableParentContext(place)
+            ?: error("TowerContext not found for ${place.getElementTextInContext()}")
+
+        val session = state.rootModuleSession
+        val firTypeReference = buildFirUserTypeRef(
+            typeReference = typeReference,
+            session = session,
+            baseScopeProvider = session.firIdeProvider.kotlinScopeProvider
+        )
+
+        return FirTypeResolveTransformer(
+            session = session,
+            scopeSession = ScopeSession(),
+            initialScopes = context.towerDataElements.asReversed().mapNotNull { it.scope }
+        ).transformTypeRef(firTypeReference, null)
+    }
+
     private class TowerProviderForElementForState(private val state: FirModuleResolveState) : FirTowerContextProvider {
         override fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext? {
             return if (ktElement.isPhysical) {

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt

@@ -12,10 +12,12 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureElement
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.KtToFirMapping
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.isNonAnonymousClassOrObject
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
@@ -51,9 +53,10 @@ internal class FirElementBuilder {
         firFileBuilder: FirFileBuilder,
         moduleFileCache: ModuleFileCache,
         fileStructureCache: FileStructureCache,
+        state: FirModuleResolveState,
     ): FirElement = when (element) {
         is KtFile -> getOrBuildFirForKtFile(element, firFileBuilder, moduleFileCache)
-        else -> getOrBuildFirForNonKtFileElement(element, fileStructureCache, moduleFileCache)
+        else -> getOrBuildFirForNonKtFileElement(element, fileStructureCache, moduleFileCache, state)
     }
 
     private fun getOrBuildFirForKtFile(ktFile: KtFile, firFileBuilder: FirFileBuilder, moduleFileCache: ModuleFileCache): FirFile =
@@ -68,16 +71,17 @@ internal class FirElementBuilder {
     private fun getOrBuildFirForNonKtFileElement(
         element: KtElement,
         fileStructureCache: FileStructureCache,
-        moduleFileCache: ModuleFileCache
+        moduleFileCache: ModuleFileCache,
+        state: FirModuleResolveState,
     ): FirElement {
         require(element !is KtFile)
-        val fileStructure = fileStructureCache.getFileStructure(element.containingKtFile, moduleFileCache)
+        val firFile = element.containingKtFile
+        val fileStructure = fileStructureCache.getFileStructure(firFile, moduleFileCache)
 
         val mappings = fileStructure.getStructureElementFor(element).mappings
         val psi = getPsiAsFirElementSource(element)
-        mappings[psi]?.let { return it }
+        return mappings.getFirOfClosestParent(psi, state)
-        return psi.getFirOfClosestParent(mappings)?.second
+            ?: state.getOrBuildFirFile(firFile)
-            ?: error("FirElement is not found for:\n${element.getElementTextInContext()}")
     }
 
     @TestOnly
@@ -91,20 +95,6 @@ internal class FirElementBuilder {
     }
 }
 
-private fun KtElement.getFirOfClosestParent(cache: Map<KtElement, FirElement>): Pair<KtElement, FirElement>? {
-    var current: PsiElement? = this
-    while (current is KtElement) {
-        val mappedFir = cache[current]
-        if (mappedFir != null) {
-            return current to mappedFir
-        }
-        current = current.parent
-    }
-
-    return null
-}
-
-
 // TODO: simplify
 internal inline fun PsiElement.getNonLocalContainingOrThisDeclaration(predicate: (KtDeclaration) -> Boolean = { true }): KtNamedDeclaration? {
     var container: PsiElement? = this

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureElement.kt

@@ -5,13 +5,17 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.DeclarationCopyBuilder
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForResolveOnAir
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileDiagnosticRetriever
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.FileStructureElementDiagnostics
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
@@ -20,13 +24,49 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.psi.*
+import java.util.concurrent.ConcurrentHashMap
 
 internal sealed class FileStructureElement(val firFile: FirFile, protected val lockProvider: LockProvider<FirFile>) {
     abstract val psi: KtAnnotated
-    abstract val mappings: Map<KtElement, FirElement>
+    abstract val mappings: KtToFirMapping
     abstract val diagnostics: FileStructureElementDiagnostics
 }
 
+internal class KtToFirMapping(firElement: FirElement, recorder: FirElementsRecorder) {
+
+    private val mapping = FirElementsRecorder.recordElementsFrom(firElement, recorder)
+
+    private val userTypeMapping = ConcurrentHashMap<KtUserType, FirElement>()
+    fun getElement(ktElement: KtElement, state: FirModuleResolveState): FirElement? {
+        mapping[ktElement]?.let { return it }
+
+        val userType = when (ktElement) {
+            is KtUserType -> ktElement
+            is KtNameReferenceExpression -> ktElement as? KtUserType
+            else -> null
+        } ?: return null
+
+        //This is for not inner KtUserType
+        if (userType.parent is KtTypeReference) return null
+
+        return userTypeMapping.getOrPut(userType) {
+            val typeReference = KtPsiFactory(ktElement.project).createType(userType)
+            LowLevelFirApiFacadeForResolveOnAir.onAirResolveTypeInPlace(ktElement, typeReference, state)
+        }
+    }
+
+    fun getFirOfClosestParent(element: KtElement, state: FirModuleResolveState): FirElement? {
+        var current: PsiElement? = element
+        while (current != null && current !is KtFile) {
+            if (current is KtElement) {
+                getElement(current, state)?.let { return it }
+            }
+            current = current.parent
+        }
+        return null
+    }
+}
+
 internal sealed class ReanalyzableStructureElement<KT : KtDeclaration, S : AbstractFirBasedSymbol<*>>(
     firFile: FirFile,
     val firSymbol: S,
@@ -66,8 +106,7 @@ internal class ReanalyzableFunctionStructureElement(
     override val timestamp: Long,
     lockProvider: LockProvider<FirFile>,
 ) : ReanalyzableStructureElement<KtNamedFunction, FirFunctionSymbol<*>>(firFile, firSymbol, lockProvider) {
-    override val mappings: Map<KtElement, FirElement> =
+    override val mappings = KtToFirMapping(firSymbol.fir, recorder)
-        FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
 
     override fun reanalyze(
         newKtDeclaration: KtNamedFunction,
@@ -106,8 +145,7 @@ internal class ReanalyzablePropertyStructureElement(
     override val timestamp: Long,
     lockProvider: LockProvider<FirFile>,
 ) : ReanalyzableStructureElement<KtProperty, FirPropertySymbol>(firFile, firSymbol, lockProvider) {
-    override val mappings: Map<KtElement, FirElement> =
+    override val mappings = KtToFirMapping(firSymbol.fir, recorder)
-        FirElementsRecorder.recordElementsFrom(firSymbol.fir, recorder)
 
     override fun reanalyze(
         newKtDeclaration: KtProperty,
@@ -145,8 +183,7 @@ internal class NonReanalyzableDeclarationStructureElement(
     override val psi: KtDeclaration,
     lockProvider: LockProvider<FirFile>,
 ) : FileStructureElement(firFile, lockProvider) {
-    override val mappings: Map<KtElement, FirElement> =
+    override val mappings = KtToFirMapping(fir, recorder)
-        FirElementsRecorder.recordElementsFrom(fir, recorder)
 
     override val diagnostics = FileStructureElementDiagnostics(firFile, lockProvider, SingleNonLocalDeclarationDiagnosticRetriever(fir))
 
@@ -176,8 +213,7 @@ internal class RootStructureElement(
     override val psi: KtFile,
     lockProvider: LockProvider<FirFile>,
 ) : FileStructureElement(firFile, lockProvider) {
-    override val mappings: Map<KtElement, FirElement> =
+    override val mappings = KtToFirMapping(firFile, recorder)
-        FirElementsRecorder.recordElementsFrom(firFile, recorder)
 
     override val diagnostics = FileStructureElementDiagnostics(firFile, lockProvider, FileDiagnosticRetriever)
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FirElementsRecorder.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.references.*

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -40,6 +40,7 @@ internal class FirLazyDeclarationResolver(
         annotations: List<FirAnnotationCall>,
         moduleFileCache: ModuleFileCache,
         scopeSession: ScopeSession,
+        collector: FirTowerDataContextCollector? = null,
     ) {
         lazyResolveDeclaration(
             declaration = firFile,
@@ -49,7 +50,7 @@ internal class FirLazyDeclarationResolver(
             reresolveFile = false
         )
         firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
-            resolveFileAnnotationsWithoutLock(firFile, annotations, scopeSession)
+            resolveFileAnnotationsWithoutLock(firFile, annotations, scopeSession, collector)
         }
     }
 
@@ -61,12 +62,14 @@ internal class FirLazyDeclarationResolver(
         firFile: FirFile,
         annotations: List<FirAnnotationCall>,
         scopeSession: ScopeSession,
+        collector: FirTowerDataContextCollector? = null,
     ) {
         FirFileAnnotationsResolveTransformer(
             firFile = firFile,
             annotations = annotations,
             session = firFile.moduleData.session,
-            scopeSession = scopeSession
+            scopeSession = scopeSession,
+            firTowerDataContextCollector = collector,
         ).transformDeclaration()
     }
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/RawFirReplacement.kt

@@ -28,7 +28,7 @@ internal data class RawFirReplacement(val from: KtElement, val to: KtElement) {
         private var replacementApplied = false
 
         private fun ensureReplacementIsValid() {
-            require(isApplicableForReplacement(from)) {
+            require(from == to || isApplicableForReplacement(from)) {
                 "Replacement is possible for applicable type but given ${from::class.simpleName}"
             }
             require(from::class == to::class) {
@@ -44,7 +44,7 @@ internal data class RawFirReplacement(val from: KtElement, val to: KtElement) {
         }
 
         fun ensureApplied() {
-            check(replacementApplied) {
+            check(from == to || replacementApplied) {
                 "Replacement requested but was not applied for ${from::class.simpleName}"
             }
         }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/RawFirUserTypeRefBuilder.kt

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
+
+import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.builder.RawFirBuilder
+import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
+import org.jetbrains.kotlin.fir.types.FirUserTypeRef
+import org.jetbrains.kotlin.psi.KtTypeReference
+
+internal fun buildFirUserTypeRef(
+    typeReference: KtTypeReference,
+    session: FirSession,
+    baseScopeProvider: FirScopeProvider
+): FirUserTypeRef {
+    val builder = object : RawFirBuilder(session, baseScopeProvider) {
+        fun build(): FirUserTypeRef = Visitor().visitTypeReference(typeReference, Unit) as FirUserTypeRef
+    }
+    builder.context.packageFqName = typeReference.containingKtFile.packageFqName
+    return builder.build()
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt

@@ -6,15 +6,10 @@
 package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.tree.TokenSet
-import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.fir.*
-import com.intellij.psi.util.parentOfType
-import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticPropertySymbol
 import org.jetbrains.kotlin.fir.resolve.diagnostics.ConeAmbiguityError
@@ -29,8 +24,6 @@ import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.idea.fir.getCandidateSymbols
 import org.jetbrains.kotlin.idea.fir.isImplicitFunctionCall
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForResolveOnAir
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
@@ -169,7 +162,7 @@ internal object FirReferenceResolveHelper {
         val fir = expression.getOrBuildFir(analysisSession.firResolveState)
         val session = analysisSession.firResolveState.rootModuleSession
         return when (fir) {
-            is FirResolvedTypeRef -> getSymbolsForResolvedTypeRef(fir, expression, session, symbolBuilder, analysisSession.firResolveState)
+            is FirResolvedTypeRef -> getSymbolsForResolvedTypeRef(fir, expression, session, symbolBuilder)
             is FirResolvedQualifier ->
                 getSymbolsForResolvedQualifier(fir, expression, session, symbolBuilder, analysisSession)
             is FirAnnotationCall -> getSymbolsForAnnotationCall(fir, session, symbolBuilder)
@@ -349,41 +342,20 @@ internal object FirReferenceResolveHelper {
         }
     }
 
-    private fun tryResolvePartiallyCorrectReference(
-        expression: KtSimpleNameExpression,
-        resolveState: FirModuleResolveState
-    ): FirResolvedTypeRef? {
-        val expressionUserType = expression.parent as? KtUserType ?: return null
-        val typeReference = expressionUserType.parentOfType<KtTypeReference>(withSelf = false) ?: return null
-        return LowLevelFirApiFacadeForResolveOnAir.onAirResolveElement(
-            state = resolveState,
-            place = typeReference,
-            elementToResolve = KtPsiFactory(expression.project).createType(expressionUserType),
-        ) as? FirResolvedTypeRef
-    }
-
     private fun getSymbolsForResolvedTypeRef(
         fir: FirResolvedTypeRef,
         expression: KtSimpleNameExpression,
         session: FirSession,
         symbolBuilder: KtSymbolByFirBuilder,
-        resolveState: FirModuleResolveState,
     ): Collection<KtSymbol> {
-        if (expression.isPartOfUserTypeRefQualifier()) {
-            val typeQualifier = findPossibleTypeQualifier(expression, fir)?.toTargetPsi(session, symbolBuilder)
-            val typeOrPackageQualifier =
-                typeQualifier ?: getPackageSymbolFor(expression, symbolBuilder, forQualifiedType = true)
 
-            if (typeOrPackageQualifier != null) return listOf(typeOrPackageQualifier)
+        val isPossiblyPackage = fir is FirErrorTypeRef && expression.isPartOfUserTypeRefQualifier()
 
-            val resolvedPartType = tryResolvePartiallyCorrectReference(
+        val resultSymbol =
-                expression,
+            if (isPossiblyPackage) getPackageSymbolFor(expression, symbolBuilder, forQualifiedType = true)
-                resolveState
+            else fir.toTargetSymbol(session, symbolBuilder)
-            )?.toTargetSymbol(session, symbolBuilder)
 
-            return listOfNotNull(resolvedPartType)
+        return listOfNotNull(resultSymbol)
-        }
-        return listOfNotNull(fir.toTargetSymbol(session, symbolBuilder))
     }
 
     private fun getSymbolsForResolvedQualifier(

idea/testData/resolve/references/ClassNameBeforeDot.kt

@@ -1,4 +1,3 @@
-// IGNORE_FIR
 
 class Foo