[FIR IDE] Correct in-body supertype resolve of fake nodes

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSupertypesResolution.kt

@@ -35,10 +35,19 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
-class FirSupertypeResolverProcessor(session: FirSession, scopeSession: ScopeSession) : FirTransformerBasedResolveProcessor(session, scopeSession) {
+class FirSupertypeResolverProcessor(session: FirSession, scopeSession: ScopeSession) :
+    FirTransformerBasedResolveProcessor(session, scopeSession) {
     override val transformer = FirSupertypeResolverTransformer(session, scopeSession)
 }
 
+/**
+ * Interceptor needed by IDE to resolve in-air created declarations.
+ */
+interface FirProviderInterceptorForSupertypeResolver {
+    fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile
+    fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>?
+}
+
 open class FirSupertypeResolverTransformer(
     final override val session: FirSession,
     scopeSession: ScopeSession
@@ -65,12 +74,14 @@ fun <F : FirClassLikeDeclaration<F>> F.runSupertypeResolvePhaseForLocalClass(
     scopeSession: ScopeSession,
     currentScopeList: List<FirScope>,
     localClassesNavigationInfo: LocalClassesNavigationInfo,
+    firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
 ): F {
     val supertypeComputationSession = SupertypeComputationSession()
     val supertypeResolverVisitor = FirSupertypeResolverVisitor(
         session, supertypeComputationSession, scopeSession,
         currentScopeList.toPersistentList(),
-        localClassesNavigationInfo
+        localClassesNavigationInfo,
+        firProviderInterceptor
     )
 
     this.accept(supertypeResolverVisitor, null)
@@ -172,20 +183,22 @@ fun FirRegularClass.resolveSupertypesInTheAir(session: FirSession): List<FirType
         .resolveSpecificClassLikeSupertypes(this, superTypeRefs)
 }
 
-open class FirSupertypeResolverVisitor(
+class FirSupertypeResolverVisitor(
     private val session: FirSession,
     private val supertypeComputationSession: SupertypeComputationSession,
     private val scopeSession: ScopeSession,
     private val scopeForLocalClass: PersistentList<FirScope>? = null,
-    private val localClassesNavigationInfo: LocalClassesNavigationInfo? = null
+    private val localClassesNavigationInfo: LocalClassesNavigationInfo? = null,
+    private val firProviderInterceptor: FirProviderInterceptorForSupertypeResolver? = null,
 ) : FirDefaultVisitor<Unit, Any?>() {
     private val supertypeGenerationExtensions = session.extensionService.supertypeGenerators
 
-    open fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile =
+    private fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile =
-        session.firProvider.getFirClassifierContainerFile(symbol.classId)
+        firProviderInterceptor?.getFirClassifierContainerFile(symbol) ?: session.firProvider.getFirClassifierContainerFile(symbol.classId)
 
-    open fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? =
+    private fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? =
-        session.firProvider.getFirClassifierByFqName(classId)
+        if (firProviderInterceptor != null) firProviderInterceptor.getFirClassifierByFqName(classId)
+        else session.firProvider.getFirClassifierByFqName(classId)
 
     override fun visitElement(element: FirElement, data: Any?) {}
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
 import org.jetbrains.kotlin.fir.scopes.FirCompositeScope
@@ -31,7 +32,8 @@ open class FirBodyResolveTransformer(
     scopeSession: ScopeSession,
     val returnTypeCalculator: ReturnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve(),
     outerBodyResolveContext: BodyResolveContext? = null,
-    val firTowerDataContextCollector: FirTowerDataContextCollector? = null
+    val firTowerDataContextCollector: FirTowerDataContextCollector? = null,
+    val firProviderInterceptor: FirProviderInterceptorForSupertypeResolver? = null,
 ) : FirAbstractBodyResolveTransformer(phase) {
 
     final override val context: BodyResolveContext =

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -351,7 +351,13 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
     override fun transformRegularClass(regularClass: FirRegularClass, data: ResolutionMode): FirStatement {
         if (regularClass.isLocal && regularClass !in context.targetedLocalClasses) {
-            return regularClass.runAllPhasesForLocalClass(transformer, components, data, transformer.firTowerDataContextCollector)
+            return regularClass.runAllPhasesForLocalClass(
+                transformer,
+                components,
+                data,
+                transformer.firTowerDataContextCollector,
+                transformer.firProviderInterceptor
+            )
         }
 
         doTransformTypeParameters(regularClass)
@@ -360,7 +366,13 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
     override fun transformTypeAlias(typeAlias: FirTypeAlias, data: ResolutionMode): FirDeclaration {
         if (typeAlias.isLocal && typeAlias !in context.targetedLocalClasses) {
-            return typeAlias.runAllPhasesForLocalClass(transformer, components, data, transformer.firTowerDataContextCollector)
+            return typeAlias.runAllPhasesForLocalClass(
+                transformer,
+                components,
+                data,
+                transformer.firTowerDataContextCollector,
+                transformer.firProviderInterceptor
+            )
         }
         doTransformTypeParameters(typeAlias)
         typeAlias.transformAnnotations(transformer, data)
@@ -400,7 +412,13 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         data: ResolutionMode
     ): FirStatement {
         if (anonymousObject !in context.targetedLocalClasses) {
-            return anonymousObject.runAllPhasesForLocalClass(transformer, components, data, transformer.firTowerDataContextCollector)
+            return anonymousObject.runAllPhasesForLocalClass(
+                transformer,
+                components,
+                data,
+                transformer.firTowerDataContextCollector,
+                transformer.firProviderInterceptor
+            )
         }
         dataFlowAnalyzer.enterClass()
         if (anonymousObject.typeRef !is FirResolvedTypeRef) {
@@ -858,7 +876,12 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                 valueParameter.transformReturnTypeRef(
                     StoreType,
                     valueParameter.returnTypeRef.resolvedTypeFromPrototype(
-                        ConeKotlinErrorType(ConeSimpleDiagnostic("No type for parameter", DiagnosticKind.ValueParameterWithNoTypeAnnotation))
+                        ConeKotlinErrorType(
+                            ConeSimpleDiagnostic(
+                                "No type for parameter",
+                                DiagnosticKind.ValueParameterWithNoTypeAnnotation
+                            )
+                        )
                     )
                 )
             }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirTowerDataContextCollector.kt

@@ -10,6 +10,9 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 
+/**
+ * Collector needed to collect TowerDataContext in IDE for in-air resolved fir elements.
+ */
 interface FirTowerDataContextCollector {
     fun addFileContext(file: FirFile, context: FirTowerDataContext)
     fun addStatementContext(statement: FirStatement, context: FirTowerDataContext)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/LocalClassesResolution.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
 import org.jetbrains.kotlin.fir.resolve.transformers.runStatusResolveForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runSupertypeResolvePhaseForLocalClass
 import org.jetbrains.kotlin.fir.resolve.transformers.runTypeResolvePhaseForLocalClass
@@ -18,7 +19,8 @@ fun <F : FirClassLikeDeclaration<F>> F.runAllPhasesForLocalClass(
     transformer: FirAbstractBodyResolveTransformer,
     components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
     resolutionMode: ResolutionMode,
-    firTowerDataContextCollector: FirTowerDataContextCollector?
+    firTowerDataContextCollector: FirTowerDataContextCollector?,
+    firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
 ): F {
     if (this.resolvePhase > FirResolvePhase.RAW_FIR) return this
     if (this is FirRegularClass) {
@@ -36,7 +38,8 @@ fun <F : FirClassLikeDeclaration<F>> F.runAllPhasesForLocalClass(
         components.session,
         components.scopeSession,
         components.createCurrentScopeList(),
-        localClassesNavigationInfo
+        localClassesNavigationInfo,
+        firProviderInterceptor
     )
     runTypeResolvePhaseForLocalClass(
         components.session,

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt

@@ -283,9 +283,12 @@ private class TypeNamesProvider(private val indexHelper: IndexHelper) {
     fun KtAnalysisSession.findAllNames(type: KtType): Set<String> {
         if (type !is KtClassType) return emptySet()
 
-        val result = hashSetOf<String>()
+        val typeName = type.classId.shortClassName.let {
-        val typeName = type.classId.shortClassName.identifier
+            if (it.isSpecial) return emptySet()
+            it.identifier
+        }
 
+        val result = hashSetOf<String>()
         result += typeName
         result += indexHelper.getPossibleTypeAliasExpansionNames(typeName)
 

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

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataC
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.SingleElementTowerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirProviderInterceptorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
@@ -70,7 +71,12 @@ object LowLevelFirApiFacadeForDependentCopy {
         require(!elementToResolve.isPhysical)
 
         val collector = FirTowerDataContextAllElementsCollector()
-        val declaration = runBodyResolve(state, replacement = RawFirReplacement(place, elementToResolve))
+        val declaration = runBodyResolve(
+            state = state,
+            replacement = RawFirReplacement(place, elementToResolve),
+            collector = null,
+            useFirProviderInterceptor = true
+        )
 
         val expressionLocator = object : FirVisitorVoid() {
             var result: FirElement? = null
@@ -158,7 +164,8 @@ object LowLevelFirApiFacadeForDependentCopy {
         val copiedFirDeclaration = runBodyResolve(
             originalState,
             collector = collector,
-            replacement = RawFirReplacement(sameDeclarationInOriginalFile, dependencyNonLocalDeclaration)
+            replacement = RawFirReplacement(sameDeclarationInOriginalFile, dependencyNonLocalDeclaration),
+            useFirProviderInterceptor = true
         )
 
         val recordedMap = FirElementsRecorder.recordElementsFrom(copiedFirDeclaration, FirElementsRecorder())
@@ -169,6 +176,7 @@ object LowLevelFirApiFacadeForDependentCopy {
         state: FirModuleResolveStateImpl,
         replacement: RawFirReplacement<T>,
         collector: FirTowerDataContextCollector? = null,
+        useFirProviderInterceptor: Boolean = false
     ): FirDeclaration {
         val copiedFirDeclaration = DeclarationCopyBuilder.createDeclarationCopy(
             state = state,
@@ -177,6 +185,13 @@ object LowLevelFirApiFacadeForDependentCopy {
 
         val originalFirFile = state.getOrBuildFirFor(replacement.from.containingKtFile) as FirFile
 
+        val firProviderInterceptor =
+            if (useFirProviderInterceptor) FirProviderInterceptorForIDE.createForFirElement(
+                session = originalFirFile.declarationSiteSession,
+                firFile = originalFirFile,
+                element = copiedFirDeclaration
+            ) else null
+
         state.firFileBuilder.runCustomResolveWithPCECheck(originalFirFile, state.rootModuleSession.cache) {
             state.firLazyDeclarationResolver.runLazyResolveWithoutLock(
                 copiedFirDeclaration,
@@ -187,7 +202,8 @@ object LowLevelFirApiFacadeForDependentCopy {
                 toPhase = FirResolvePhase.BODY_RESOLVE,
                 towerDataContextCollector = collector,
                 checkPCE = true,
-                lastNonLazyPhase = FirResolvePhase.IMPORTS
+                lastNonLazyPhase = FirResolvePhase.IMPORTS,
+                firProviderInterceptor = firProviderInterceptor,
             )
         }
 

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDesignation
@@ -121,7 +122,8 @@ internal class FirLazyDeclarationResolver(
         toPhase: FirResolvePhase,
         towerDataContextCollector: FirTowerDataContextCollector? = null,
         checkPCE: Boolean,
-        lastNonLazyPhase: FirResolvePhase = LAST_NON_LAZY_PHASE
+        lastNonLazyPhase: FirResolvePhase = LAST_NON_LAZY_PHASE,
+        firProviderInterceptor: FirProviderInterceptorForSupertypeResolver? = null
     ) {
         if (fromPhase >= toPhase) return
         val nonLazyPhase = minOf(toPhase, lastNonLazyPhase)
@@ -157,7 +159,8 @@ internal class FirLazyDeclarationResolver(
                 currentPhase,
                 scopeSession,
                 towerDataContextCollector,
-                designation
+                firProviderInterceptor,
+                designation,
             )
         }
     }
@@ -167,6 +170,7 @@ internal class FirLazyDeclarationResolver(
         phase: FirResolvePhase,
         scopeSession: ScopeSession,
         towerDataContextCollector: FirTowerDataContextCollector?,
+        firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
         designation: FirDeclarationDesignation
     ) {
         if (designation.fullDesignation.all { it.resolvePhase >= phase }) {
@@ -177,7 +181,8 @@ internal class FirLazyDeclarationResolver(
             designation,
             containerFirFile,
             scopeSession,
-            towerDataContextCollector
+            towerDataContextCollector,
+            firProviderInterceptor,
         )
 
         firFileBuilder.firPhaseRunner.runPhaseWithCustomResolve(phase) {
@@ -189,13 +194,14 @@ internal class FirLazyDeclarationResolver(
         designation: FirDeclarationDesignation,
         containerFirFile: FirFile,
         scopeSession: ScopeSession,
-        towerDataContextCollector: FirTowerDataContextCollector?
+        towerDataContextCollector: FirTowerDataContextCollector?,
+        firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
     ): FirLazyTransformerForIDE = when (this) {
         FirResolvePhase.SUPER_TYPES -> FirDesignatedSupertypeResolverTransformerForIDE(
             designation,
-            containerFirFile,
             containerFirFile.moduleData.session,
-            scopeSession
+            scopeSession,
+            firProviderInterceptor,
         )
         FirResolvePhase.SEALED_CLASS_INHERITORS -> FirLazyTransformerForIDE.EMPTY
         FirResolvePhase.TYPES -> FirDesignatedTypeResolverTransformerForIDE(
@@ -226,7 +232,8 @@ internal class FirLazyDeclarationResolver(
             designation,
             containerFirFile.moduleData.session,
             scopeSession,
-            towerDataContextCollector
+            towerDataContextCollector,
+            firProviderInterceptor
         )
         else -> error("Non-lazy phase $this")
     }

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

@@ -11,10 +11,8 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
 
@@ -23,7 +21,8 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
     designation: FirDeclarationDesignation,
     session: FirSession,
     scopeSession: ScopeSession,
-    towerDataContextCollector: FirTowerDataContextCollector? = null
+    towerDataContextCollector: FirTowerDataContextCollector? = null,
+    firProviderInterceptor: FirProviderInterceptorForSupertypeResolver? = null,
 ) : FirLazyTransformerForIDE, FirBodyResolveTransformer(
     session,
     phase = FirResolvePhase.BODY_RESOLVE,
@@ -35,7 +34,8 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
         ImplicitBodyResolveComputationSession(),
         ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
     ),
-    firTowerDataContextCollector = towerDataContextCollector
+    firTowerDataContextCollector = towerDataContextCollector,
+    firProviderInterceptor = firProviderInterceptor,
 ) {
     private val ideDeclarationTransformer = IDEDeclarationTransformer(designation.toDesignationIterator())
 

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

@@ -7,54 +7,35 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.trasformers
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.FirApplySupertypesTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSupertypeResolverVisitor
 import org.jetbrains.kotlin.fir.resolve.transformers.SupertypeComputationSession
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
-import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.idea.fir.low.level.api.trasformers.FirLazyTransformerForIDE.Companion.ensurePhase
-import org.jetbrains.kotlin.name.ClassId
 
 internal class FirDesignatedSupertypeResolverTransformerForIDE(
     private val designation: FirDeclarationDesignation,
-    private val firFile: FirFile,
     private val session: FirSession,
     private val scopeSession: ScopeSession,
-    ) : FirLazyTransformerForIDE {
+    private val firProviderInterceptor: FirProviderInterceptorForSupertypeResolver?,
+) : FirLazyTransformerForIDE {
 
     private val supertypeComputationSession = SupertypeComputationSession()
 
-    private inner class SupertypeResolver(
-        private val symbolSet: Set<FirClassLikeSymbol<*>>,
-        private val classIdToElementMap: Map<ClassId, FirClassLikeDeclaration<*>>
-    ) : FirSupertypeResolverVisitor(session, supertypeComputationSession, scopeSession) {
-        override fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile =
-            if (symbolSet.contains(symbol)) firFile else session.firProvider.getFirClassifierContainerFile(symbol)
-
-        override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? =
-            classIdToElementMap[classId] ?: session.firProvider.getFirClassifierByFqName(classId)
-    }
-
     override fun transformDeclaration() {
         designation.ensurePhase(FirResolvePhase.SUPER_TYPES, exceptLast = true)
-        val nodeInfoCollector = object : FirVisitorVoid() {
-            val symbolSet = mutableSetOf<FirClassLikeSymbol<*>>()
-            val classIdToElementMap = mutableMapOf<ClassId, FirClassLikeDeclaration<*>>()
-            override fun visitElement(element: FirElement) {
-                if (element is FirClassLikeDeclaration<*>) {
-                    symbolSet.add(element.symbol)
-                    classIdToElementMap[element.symbol.classId] = element
-                }
-                element.acceptChildren(this)
-            }
-        }
-        nodeInfoCollector.visitElement(designation.declaration)
 
-        val resolver = SupertypeResolver(nodeInfoCollector.symbolSet, nodeInfoCollector.classIdToElementMap)
+        val resolver = FirSupertypeResolverVisitor(
+            session = session,
+            supertypeComputationSession = supertypeComputationSession,
+            scopeSession = scopeSession,
+            scopeForLocalClass = null,
+            localClassesNavigationInfo = null,
+            firProviderInterceptor = firProviderInterceptor,
+        )
         designation.declaration.accept(resolver, null)
         val applySupertypesTransformer = FirApplySupertypesTransformer(supertypeComputationSession)
         designation.declaration.transform<FirElement, Void?>(applySupertypesTransformer, null)

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

@@ -0,0 +1,48 @@
+/*
+ * 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.trasformers
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.resolve.firProvider
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.name.ClassId
+
+internal class FirProviderInterceptorForIDE private constructor(
+    private val firFile: FirFile,
+    private val session: FirSession,
+    private val symbolSet: Set<FirClassLikeSymbol<*>>,
+    private val classIdToElementMap: Map<ClassId, FirClassLikeDeclaration<*>>): FirProviderInterceptorForSupertypeResolver {
+
+    override fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile =
+        if (symbolSet.contains(symbol)) firFile else session.firProvider.getFirClassifierContainerFile(symbol)
+
+    override fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? =
+        classIdToElementMap[classId] ?: session.firProvider.getFirClassifierByFqName(classId)
+
+    companion object {
+        fun createForFirElement(session: FirSession, firFile: FirFile, element: FirElement): FirProviderInterceptorForSupertypeResolver {
+            val nodeInfoCollector = object : FirVisitorVoid() {
+                val symbolSet = mutableSetOf<FirClassLikeSymbol<*>>()
+                val classIdToElementMap = mutableMapOf<ClassId, FirClassLikeDeclaration<*>>()
+                override fun visitElement(element: FirElement) {
+                    if (element is FirClassLikeDeclaration<*>) {
+                        symbolSet.add(element.symbol)
+                        classIdToElementMap[element.symbol.classId] = element
+                    }
+                    element.acceptChildren(this)
+                }
+            }
+            nodeInfoCollector.visitElement(element)
+
+            return FirProviderInterceptorForIDE(firFile, session, nodeInfoCollector.symbolSet, nodeInfoCollector.classIdToElementMap)
+        }
+    }
+}