FIR IDE: always pass resolve state as parameter when getting fir by psi

2020-07-03 00:29:07 +03:00
parent 0c13a7f89a
commit f98ff2c25e
11 changed files with 77 additions and 66 deletions
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.idea.fir.low.level.api

+import com.google.common.collect.MapMaker
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
@@ -31,15 +32,15 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtTypeReference

-internal interface FirModuleResolveState {
-    val sessionProvider: FirProjectSessionProvider
+abstract class FirModuleResolveState {
+    internal abstract val sessionProvider: FirProjectSessionProvider

-    fun getSession(psi: KtElement): FirSession {
+    internal fun getSession(psi: KtElement): FirSession {
        val moduleInfo = psi.getModuleInfo()
        return getSession(psi.project, moduleInfo)
    }

-    fun getSession(project: Project, moduleInfo: IdeaModuleInfo): FirSession {
+    internal fun getSession(project: Project, moduleInfo: IdeaModuleInfo): FirSession {
        sessionProvider.getSession(moduleInfo)?.let { return it }
        val lock = when (moduleInfo) {
            is ModuleSourceInfo -> moduleInfo.module
@@ -58,20 +59,20 @@ internal interface FirModuleResolveState {
        }
    }

-    operator fun get(psi: KtElement): FirElement?
+    internal abstract operator fun get(psi: KtElement): FirElement?

-    fun getDiagnostics(psi: KtElement): List<Diagnostic>
+    internal abstract fun getDiagnostics(psi: KtElement): List<Diagnostic>

-    fun hasDiagnosticsForFile(file: KtFile): Boolean
+    internal abstract fun hasDiagnosticsForFile(file: KtFile): Boolean

-    fun record(psi: KtElement, fir: FirElement)
+    internal abstract fun record(psi: KtElement, fir: FirElement)

-    fun record(psi: KtElement, diagnostic: Diagnostic)
+    internal abstract fun record(psi: KtElement, diagnostic: Diagnostic)

-    fun setDiagnosticsForFile(file: KtFile, fir: FirFile, diagnostics: Iterable<FirDiagnostic<*>> = emptyList())
+    internal abstract fun setDiagnosticsForFile(file: KtFile, fir: FirFile, diagnostics: Iterable<FirDiagnostic<*>> = emptyList())
 }

-internal class FirModuleResolveStateImpl(override val sessionProvider: FirProjectSessionProvider) : FirModuleResolveState {
+internal class FirModuleResolveStateImpl(override val sessionProvider: FirProjectSessionProvider) : FirModuleResolveState() {
    private val cache = mutableMapOf<KtElement, FirElement>()

    private val diagnosticCache = mutableMapOf<KtElement, MutableList<Diagnostic>>()
@@ -7,17 +7,23 @@ package org.jetbrains.kotlin.idea.fir.low.level.api

 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.psi.KtElement

 object LowLevelFirApiFacade {
-    fun getOrBuildFirFor(element: KtElement, phase: FirResolvePhase): FirElement =
-        element.getOrBuildFir(element.firResolveState(), phase)
+    fun getResolveStateFor(element: KtElement): FirModuleResolveState =
+        element.firResolveState()

-    fun getDiagnosticsFor(element: KtElement): Collection<Diagnostic> {
+    fun getSessionFor(element: KtElement, resolveState: FirModuleResolveState): FirSession =
+        resolveState.getSession(element)
+
+    fun getOrBuildFirFor(element: KtElement, resolveState: FirModuleResolveState, phase: FirResolvePhase): FirElement =
+        element.getOrBuildFir(resolveState, phase)
+
+    fun getDiagnosticsFor(element: KtElement, resolveState: FirModuleResolveState): Collection<Diagnostic> {
        val file = element.containingKtFile
-        val state = element.firResolveState()
-        file.getOrBuildFirWithDiagnostics(state)
-        return state.getDiagnostics(element)
+        file.getOrBuildFirWithDiagnostics(resolveState)
+        return resolveState.getDiagnostics(element)
    }
 }
@@ -5,10 +5,7 @@

 package org.jetbrains.kotlin.idea.fir

-import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
@@ -16,28 +13,32 @@ import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
-import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import kotlin.reflect.KClass

 fun KtElement.getOrBuildFir(
+    resolveState: FirModuleResolveState,
    phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
-) = LowLevelFirApiFacade.getOrBuildFirFor(this, phase)
+) = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState, phase)

 inline fun <reified E : FirElement> KtElement.getOrBuildFirSafe(
+    resolveState: FirModuleResolveState,
    phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
-) = LowLevelFirApiFacade.getOrBuildFirFor(this, phase) as? E
+) = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState, phase) as? E

 inline fun <reified E : FirElement> KtElement.getOrBuildFirOfType(
+    resolveState: FirModuleResolveState,
    phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE
 ): E {
-    val fir = LowLevelFirApiFacade.getOrBuildFirFor(this, phase)
+    val fir = LowLevelFirApiFacade.getOrBuildFirFor(this, resolveState, phase)
    if (fir is E) return fir
    throw InvalidFirElementTypeException(this, E::class, fir::class)
 }

+
 class InvalidFirElementTypeException(
    ktElement: KtElement,
    expectedFirClass: KClass<out FirElement>,
@@ -45,10 +46,6 @@ class InvalidFirElementTypeException(
 ) : IllegalStateException("For $ktElement with text `${ktElement.text}` the $expectedFirClass expected, but $actualFirClass found")


-val KtElement.session
-    get() = parentOfType<KtDeclaration>()?.getOrBuildFirSafe<FirDeclaration>()?.session
-        ?: containingKtFile.getOrBuildFirSafe<FirFile>()!!.session
-
 fun FirFunctionCall.isImplicitFunctionCall(): Boolean {
    if (dispatchReceiver !is FirQualifiedAccessExpression) return false
    val resolvedCalleeSymbol = (calleeReference as? FirResolvedNamedReference)?.resolvedSymbol
@@ -34,13 +34,17 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*

 class KtFirAnalysisSession(
-    element: KtElement
+    private val element: KtElement
 ) : KtAnalysisSession(element.project) {
    internal val token get() = validityToken

+    internal val firResolveState = LowLevelFirApiFacade.getResolveStateFor(element)
+
+    internal val firSession get() = LowLevelFirApiFacade.getSessionFor(element, firResolveState)
+
    internal val firSymbolBuilder = KtSymbolByFirBuilder(
-        element.session.firSymbolProvider,
-        ConeTypeCheckerContext(isErrorTypeEqualsToAnything = true, isStubTypeEqualsToAnything = true, element.session),
+        firSession.firSymbolProvider,
+        ConeTypeCheckerContext(isErrorTypeEqualsToAnything = true, isStubTypeEqualsToAnything = true, firSession),
        element.project,
        validityToken
    )
@@ -48,7 +52,8 @@ class KtFirAnalysisSession(
    override val symbolProvider: KtSymbolProvider =
        KtFirSymbolProvider(
            this,
-            element.session.firSymbolProvider,
+            firSession.firSymbolProvider,
+            firResolveState,
            firSymbolBuilder
        )

@@ -58,17 +63,16 @@ class KtFirAnalysisSession(

    override fun getSmartCastedToTypes(expression: KtExpression): Collection<KtType>? = withValidityAssertion {
        // TODO filter out not used smartcasts
-        expression.getOrBuildFirSafe<FirExpressionWithSmartcast>()?.typesFromSmartCast?.map { it.asTypeInfo() }
+        expression.getOrBuildFirSafe<FirExpressionWithSmartcast>(firResolveState)?.typesFromSmartCast?.map { it.asTypeInfo() }
    }

    @OptIn(ExperimentalStdlibApi::class)
    override fun getImplicitReceiverSmartCasts(expression: KtExpression): Collection<ImplicitReceiverSmartCast> = withValidityAssertion {
        // TODO filter out not used smartcasts
-        val qualifiedExpression = expression.getOrBuildFirSafe<FirQualifiedAccessExpression>() ?: return emptyList()
+        val qualifiedExpression = expression.getOrBuildFirSafe<FirQualifiedAccessExpression>(firResolveState) ?: return emptyList()
        if (qualifiedExpression.dispatchReceiver !is FirExpressionWithSmartcast
            && qualifiedExpression.extensionReceiver !is FirExpressionWithSmartcast
        ) return emptyList()
-        val session = expression.session
        buildList {
            (qualifiedExpression.dispatchReceiver as? FirExpressionWithSmartcast)?.let { smartCasted ->
                ImplicitReceiverSmartCast(
@@ -87,34 +91,34 @@ class KtFirAnalysisSession(


    override fun getReturnTypeForKtDeclaration(declaration: KtDeclaration): KtType = withValidityAssertion {
-        val firDeclaration = declaration.getOrBuildFirOfType<FirCallableDeclaration<*>>()
+        val firDeclaration = declaration.getOrBuildFirOfType<FirCallableDeclaration<*>>(firResolveState)
        firDeclaration.returnTypeRef.coneType.asTypeInfo()
    }

    override fun getKtExpressionType(expression: KtExpression): KtType = withValidityAssertion {
-        expression.getOrBuildFirOfType<FirExpression>().typeRef.coneType.asTypeInfo()
+        expression.getOrBuildFirOfType<FirExpression>(firResolveState).typeRef.coneType.asTypeInfo()
    }

    override fun isSubclassOf(klass: KtClassOrObject, superClassId: ClassId): Boolean {
        assertIsValid()
        var result = false
-        forEachSuperClass(klass.getOrBuildFirSafe() ?: return false) { type ->
-            result = result || type.firClassLike(klass.session)?.symbol?.classId == superClassId
+        forEachSuperClass(klass.getOrBuildFirSafe(firResolveState) ?: return false) { type ->
+            result = result || type.firClassLike(firSession)?.symbol?.classId == superClassId
        }
        return result
    }

    override fun getDiagnosticsForElement(element: KtElement): Collection<Diagnostic> = withValidityAssertion {
-        LowLevelFirApiFacade.getDiagnosticsFor(element)
+        LowLevelFirApiFacade.getDiagnosticsFor(element, firResolveState)
    }

    override fun resolveCall(call: KtBinaryExpression): CallInfo? = withValidityAssertion {
-        val firCall = call.getOrBuildFirSafe<FirFunctionCall>() ?: return null
+        val firCall = call.getOrBuildFirSafe<FirFunctionCall>(firResolveState) ?: return null
        resolveCall(firCall, call)
    }

    override fun resolveCall(call: KtCallExpression): CallInfo? = withValidityAssertion {
-        val firCall = when (val fir = call.getOrBuildFir()) {
+        val firCall = when (val fir = call.getOrBuildFir(firResolveState)) {
            is FirFunctionCall -> fir
            is FirSafeCallExpression -> fir.regularQualifiedAccess as? FirFunctionCall
            else -> null
@@ -123,7 +127,7 @@ class KtFirAnalysisSession(
    }

    private fun resolveCall(firCall: FirFunctionCall, callExpression: KtExpression): CallInfo? {
-        val session = callExpression.session
+        val session = LowLevelFirApiFacade.getSessionFor(callExpression, firResolveState)
        val resolvedFunctionSymbol = firCall.calleeReference.toTargetSymbol(session, firSymbolBuilder)
        val resolvedCalleeSymbol = (firCall.calleeReference as? FirResolvedNamedReference)?.resolvedSymbol
        return when {
@@ -168,7 +172,6 @@ class KtFirAnalysisSession(
    private fun ConeKotlinType.asTypeInfo() = firSymbolBuilder.buildKtType(this)

    companion object {
-
        private val kotlinFunctionInvokeCallableIds = (0..23).flatMapTo(hashSetOf()) { arity ->
            listOf(
                CallableId(KotlinBuiltIns.getFunctionClassId(arity), Name.identifier("invoke")),
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.idea.fir.getOrBuildFirOfType
+import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.frontend.api.ValidityOwner
 import org.jetbrains.kotlin.idea.frontend.api.ValidityOwnerByValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
@@ -22,48 +23,49 @@ import org.jetbrains.kotlin.psi.*
 internal class KtFirSymbolProvider(
    override val token: ValidityOwner,
    firSymbolProvider: FirSymbolProvider,
+    private val resolveState: FirModuleResolveState,
    private val firSymbolBuilder: KtSymbolByFirBuilder
 ) : KtSymbolProvider(), ValidityOwnerByValidityToken {
    private val firSymbolProvider by weakRef(firSymbolProvider)

    override fun getParameterSymbol(psi: KtParameter): KtParameterSymbol = withValidityAssertion {
-        firSymbolBuilder.buildParameterSymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildParameterSymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getFunctionSymbol(psi: KtNamedFunction): KtFunctionSymbol = withValidityAssertion {
-        firSymbolBuilder.buildFunctionSymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildFunctionSymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getConstructorSymbol(psi: KtConstructor<*>): KtConstructorSymbol = withValidityAssertion {
-        firSymbolBuilder.buildConstructorSymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildConstructorSymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getTypeParameterSymbol(psi: KtTypeParameter): KtTypeParameterSymbol = withValidityAssertion {
-        firSymbolBuilder.buildTypeParameterSymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildTypeParameterSymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getTypeAliasSymbol(psi: KtTypeAlias): KtTypeAliasSymbol = withValidityAssertion {
-        firSymbolBuilder.buildTypeAliasSymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildTypeAliasSymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getEnumEntrySymbol(psi: KtEnumEntry): KtEnumEntrySymbol = withValidityAssertion {
-        firSymbolBuilder.buildEnumEntrySymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildEnumEntrySymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getAnonymousFunctionSymbol(psi: KtNamedFunction): KtAnonymousFunctionSymbol = withValidityAssertion {
-        firSymbolBuilder.buildAnonymousFunctionSymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildAnonymousFunctionSymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getAnonymousFunctionSymbol(psi: KtLambdaExpression): KtAnonymousFunctionSymbol = withValidityAssertion {
-        firSymbolBuilder.buildAnonymousFunctionSymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildAnonymousFunctionSymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getVariableSymbol(psi: KtProperty): KtVariableSymbol = withValidityAssertion {
-        firSymbolBuilder.buildVariableSymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildVariableSymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol = withValidityAssertion {
-        firSymbolBuilder.buildClassSymbol(psi.getOrBuildFirOfType())
+        firSymbolBuilder.buildClassSymbol(psi.getOrBuildFirOfType(resolveState))
    }

    override fun getClassOrObjectSymbolByClassId(classId: ClassId): KtClassOrObjectSymbol? = withValidityAssertion {
@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
 import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.fir.*
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirPackageSymbol
@@ -28,7 +29,6 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
-import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis

 internal object FirReferenceResolveHelper {
    fun FirResolvedTypeRef.toTargetSymbol(session: FirSession, symbolBuilder: KtSymbolByFirBuilder): KtSymbol? {
@@ -102,11 +102,12 @@ internal object FirReferenceResolveHelper {

    internal fun resolveSimpleNameReference(
        ref: KtFirSimpleNameReference,
-        symbolBuilder: KtSymbolByFirBuilder
+        analysisSession: KtFirAnalysisSession
    ): Collection<KtSymbol> {
        val expression = ref.expression
-        val fir = expression.getOrBuildFir()
-        val session = expression.session
+        val symbolBuilder = analysisSession.firSymbolBuilder
+        val fir = expression.getOrBuildFir(analysisSession.firResolveState)
+        val session = analysisSession.firSession
        when (fir) {
            is FirResolvable -> {
                val calleeReference =
@@ -136,7 +137,7 @@ internal object FirReferenceResolveHelper {
                        parent = parent.parent as? KtDotQualifiedExpression
                        continue
                    }
-                    val parentFir = selectorExpression.getOrBuildFir()
+                    val parentFir = selectorExpression.getOrBuildFir(analysisSession.firResolveState)
                    if (parentFir is FirQualifiedAccess) {
                        return listOfNotNull(classId.toTargetPsi(session, symbolBuilder, parentFir.calleeReference))
                    }
@@ -208,7 +209,7 @@ internal object FirReferenceResolveHelper {
                        parent = parent.parent as? KtDotQualifiedExpression
                        continue
                    }
-                    val parentFir = selectorExpression.getOrBuildFir()
+                    val parentFir = selectorExpression.getOrBuildFir(analysisSession.firResolveState)
                    if (parentFir is FirResolvedQualifier) {
                        var classId = parentFir.classId
                        while (unresolvedCounter > 0) {
@@ -18,7 +18,7 @@ internal class KtFirSimpleNameReference(
 ) : KtSimpleNameReference(expression), KtFirReference {
    override fun resolveToSymbols(analysisSession: KtAnalysisSession): Collection<KtSymbol> {
        check(analysisSession is KtFirAnalysisSession)
-        return FirReferenceResolveHelper.resolveSimpleNameReference(this, analysisSession.firSymbolBuilder)
+        return FirReferenceResolveHelper.resolveSimpleNameReference(this, analysisSession)
    }

    override fun doCanBeReferenceTo(candidateTarget: PsiElement): Boolean {
@@ -22,7 +22,7 @@ class KtFirDestructuringDeclarationReference(

    override fun resolveToSymbols(analysisSession: KtAnalysisSession): Collection<KtSymbol> {
        check(analysisSession is KtFirAnalysisSession)
-        val fir = expression.getOrBuildFirSafe<FirProperty>() ?: return emptyList()
+        val fir = expression.getOrBuildFirSafe<FirProperty>(analysisSession.firResolveState) ?: return emptyList()
        val componentFunctionSymbol = (fir.initializer as? FirComponentCall)?.getCalleeSymbol() ?: return emptyList()
        return listOfNotNull(componentFunctionSymbol.fir.buildSymbol(analysisSession.firSymbolBuilder))
    }
@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.psi.KtForExpression
 open class KtFirForLoopInReference(expression: KtForExpression) : KtForLoopInReference(expression), KtFirReference {
    override fun resolveToSymbols(analysisSession: KtAnalysisSession): Collection<KtSymbol> {
        check(analysisSession is KtFirAnalysisSession)
-        val firLoop = expression.getOrBuildFirSafe<FirWhileLoop>() ?: return emptyList()
+        val firLoop = expression.getOrBuildFirSafe<FirWhileLoop>(analysisSession.firResolveState) ?: return emptyList()
        val condition = firLoop.condition as? FirFunctionCall
        val iterator = run {
            val callee = (condition?.explicitReceiver as? FirQualifiedAccessExpression)?.calleeReference
@@ -23,7 +23,7 @@ class KtFirPropertyDelegationMethodsReference(
 ) : KtPropertyDelegationMethodsReference(element), KtFirReference {
    override fun resolveToSymbols(analysisSession: KtAnalysisSession): Collection<KtSymbol> {
        check(analysisSession is KtFirAnalysisSession)
-        val property = (expression.parent as? KtElement)?.getOrBuildFirSafe<FirProperty>() ?: return emptyList()
+        val property = (expression.parent as? KtElement)?.getOrBuildFirSafe<FirProperty>(analysisSession.firResolveState) ?: return emptyList()
        if (property.delegate == null) return emptyList()
        val getValueSymbol = (property.getter?.singleStatementOfType<FirReturnExpression>()?.result as? FirFunctionCall)?.getCalleeSymbol()
        val setValueSymbol = property.setter?.singleStatementOfType<FirFunctionCall>()?.getCalleeSymbol()
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.fir
 import com.intellij.openapi.util.io.FileUtil
 import com.intellij.rt.execution.junit.FileComparisonFailure
 import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtFunction
@@ -18,9 +19,9 @@ abstract class AbstractKtDeclarationAndFirDeclarationEqualityChecker : KotlinLig
    protected fun doTest(path: String) {
        val file = File(path)
        val ktFile = myFixture.configureByText(file.name, FileUtil.loadFile(file)) as KtFile
-
+        val resolveState = LowLevelFirApiFacade.getResolveStateFor(ktFile)
        ktFile.forEachDescendantOfType<KtFunction> { ktFunction ->
-            val firFunction = ktFunction.getOrBuildFir() as FirFunction<*>
+            val firFunction = ktFunction.getOrBuildFir(resolveState) as FirFunction<*>
            if (!KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(ktFunction, firFunction)) {
                throw FileComparisonFailure(
                    /* message= */          null,