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

2021-05-03 18:03:13 +02:00
parent 959ff50fa7
commit 2e7eedd3f2
11 changed files with 163 additions and 64 deletions
@@ -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 =
-        session.firProvider.getFirClassifierContainerFile(symbol.classId)
+    private fun getFirClassifierContainerFile(symbol: FirClassLikeSymbol<*>): FirFile =
+        firProviderInterceptor?.getFirClassifierContainerFile(symbol) ?: session.firProvider.getFirClassifierContainerFile(symbol.classId)

-    open fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? =
-        session.firProvider.getFirClassifierByFqName(classId)
+    private fun getFirClassifierByFqName(classId: ClassId): FirClassLikeDeclaration<*>? =
+        if (firProviderInterceptor != null) firProviderInterceptor.getFirClassifierByFqName(classId)
+        else session.firProvider.getFirClassifierByFqName(classId)

    override fun visitElement(element: FirElement, data: Any?) {}

@@ -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 =
@@ -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
+                            )
+                        )
                    )
                )
            }
@@ -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)
@@ -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,
@@ -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.identifier
+        val typeName = type.classId.shortClassName.let {
+            if (it.isSpecial) return emptySet()
+            it.identifier
+        }

+        val result = hashSetOf<String>()
        result += typeName
        result += indexHelper.getPossibleTypeAliasExpansionNames(typeName)

@@ -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,
            )
        }

@@ -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")
    }
@@ -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.body.resolve.FirTowerDataContextCollector
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
-import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
+import org.jetbrains.kotlin.fir.resolve.transformers.FirProviderInterceptorForSupertypeResolver
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
 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())

@@ -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)
@@ -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)
+        }
+    }
+}