FIR IDE: move lazy resolve to FirLazyDeclarationResolver

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

@@ -10,6 +10,7 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.idea.util.psiModificationTrackerBasedCachedValue
 import java.util.concurrent.ConcurrentHashMap
 import org.jetbrains.kotlin.idea.caches.project.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeLibrariesSession
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
@@ -43,6 +44,7 @@ internal class FirIdeResolveStateService(private val project: Project) {
             librariesSession,
             sessionProvider,
             sourcesSession.firFileBuilder,
+            FirLazyDeclarationResolver(sourcesSession.firFileBuilder),
             sourcesSession.cache
         )
     }

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

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBui
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.PsiToFirCache
 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.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.psi.KtElement
@@ -63,10 +64,11 @@ internal open class FirModuleResolveStateImpl(
     override val firIdeLibrariesSession: FirSession,
     private val sessionProvider: FirIdeSessionProvider,
     val firFileBuilder: FirFileBuilder,
+    val firLazyDeclarationResolver: FirLazyDeclarationResolver,
     val fileCache: ModuleFileCache,
 ) : FirModuleResolveState() {
     val psiToFirCache = PsiToFirCache(fileCache)
-    val elementBuilder = FirElementBuilder(firFileBuilder)
+    val elementBuilder = FirElementBuilder(firFileBuilder, firLazyDeclarationResolver)
     private val diagnosticsCollector = DiagnosticsCollector(firFileBuilder, fileCache)
 
     override fun getSessionFor(moduleInfo: IdeaModuleInfo): FirSession =
@@ -83,7 +85,7 @@ internal open class FirModuleResolveStateImpl(
     }
 
     override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
-        elementBuilder.lazyResolveDeclarationWithPCECheck(declaration, fileCache, toPhase)
+        firLazyDeclarationResolver.lazyResolveDeclaration(declaration, fileCache, toPhase, checkPCE = true)
         return declaration
     }
 
@@ -97,7 +99,14 @@ internal open class FirModuleResolveStateImpl(
         toPhase: FirResolvePhase,
         towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>
     ) {
-        elementBuilder.runLazyResolveForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextForStatement)
+        firLazyDeclarationResolver.runLazyResolveWithoutLock(
+            firFunction,
+            containerFirFile,
+            firIdeProvider,
+            toPhase,
+            towerDataContextForStatement,
+            checkPCE = false
+        )
     }
 }
 

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

@@ -6,18 +6,16 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
 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.resolve.transformers.PhasedFirFileResolver
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBuilder
-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.lazy.resolve.FirLazyDeclarationResolver
 
 internal class IdePhasedFirFileResolver(
-    private val firFileBuilder: FirElementBuilder,
+    private val lazyDeclarationResolver: FirLazyDeclarationResolver,
     private val cache: ModuleFileCache
 ) : PhasedFirFileResolver() {
     override fun resolveDeclaration(declaration: FirDeclaration, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
-        firFileBuilder.lazyResolveDeclarationNoPCECheck(declaration, cache, toPhase)
+        lazyDeclarationResolver.lazyResolveDeclaration(declaration, cache, toPhase, checkPCE = false)
     }
 }

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

@@ -23,7 +23,12 @@ internal class DiagnosticsCollector(
     fun getDiagnosticsFor(element: KtElement): List<Diagnostic> {
         val ktFile = element.containingKtFile
         val diagnostics = diagnosticsForFile.computeIfAbsent(ktFile) {
-            val firFile = firFileBuilder.getFirFileResolvedToPhaseWithCaching(ktFile, cache, toPhase = FirResolvePhase.BODY_RESOLVE)
+            val firFile = firFileBuilder.getFirFileResolvedToPhaseWithCaching(
+                ktFile,
+                cache,
+                toPhase = FirResolvePhase.BODY_RESOLVE,
+                checkPCE = true
+            )
             DiagnosticsForFile.collectDiagnosticsForFile(firFile)
         }
         return diagnostics.getDiagnosticsFor(element)

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

@@ -6,25 +6,16 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
 
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.render
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
-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.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
+import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
 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.providers.FirIdeProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.FirElementFinder
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
-import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
-import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
 
 /**
  * Maps [KtElement] to [FirElement]
@@ -33,6 +24,7 @@ import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
 @ThreadSafe
 internal class FirElementBuilder(
     private val firFileBuilder: FirFileBuilder,
+    private val firLazyDeclarationResolver: FirLazyDeclarationResolver,
 ) {
     fun getOrBuildFirFor(
         element: KtElement,
@@ -51,181 +43,14 @@ internal class FirElementBuilder(
             null -> firFile
             else -> error("Unsupported: ${container.text}")
         }
-
-        runLazyResolveWithPCECheck(containerFir, moduleFileCache, firFile, firFile.session.firIdeProvider, toPhase)
+        firLazyDeclarationResolver.lazyResolveDeclaration(containerFir, moduleFileCache, toPhase, checkPCE = true)
 
         return psiToFirCache.getFir(element, containerFir, firFile)
     }
-
-    fun lazyResolveDeclarationWithPCECheck(
-        declaration: FirDeclaration,
-        moduleFileCache: ModuleFileCache,
-        toPhase: FirResolvePhase
-    ) {
-        if (declaration.resolvePhase < toPhase) {
-            val firFile = declaration.getContainingFile() ?: error("FirFile was not found for\n${declaration.render()}")
-            runLazyResolveWithPCECheck(declaration, moduleFileCache, firFile, firFile.session.firIdeProvider, toPhase)
-        }
-    }
-
-    fun lazyResolveDeclarationNoPCECheck(
-        declaration: FirDeclaration,
-        moduleFileCache: ModuleFileCache,
-        toPhase: FirResolvePhase
-    ) {
-        if (declaration.resolvePhase < toPhase) {
-            val firFile = declaration.getContainingFile() ?: error("FirFile was not found for\n${declaration.render()}")
-            runLazyResolveWithoutPCECheck(declaration, moduleFileCache, firFile, firFile.session.firIdeProvider, toPhase)
-        }
-    }
-
-
-    private fun runLazyResolvePhase(
-        firDeclarationToResolve: FirDeclaration,
-        containerFirFile: FirFile,
-        provider: FirProvider,
-        toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>?,
-    ) {
-        val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider)
-
-        val designation = mutableListOf<FirDeclaration>(containerFirFile)
-        if (nonLocalDeclarationToResolve !is FirFile) {
-            val id = when (nonLocalDeclarationToResolve) {
-                is FirCallableDeclaration<*> -> {
-                    nonLocalDeclarationToResolve.symbol.callableId.classId
-                }
-                is FirRegularClass -> {
-                    nonLocalDeclarationToResolve.symbol.classId
-                }
-                else -> error("Unsupported: ${nonLocalDeclarationToResolve.render()}")
-            }
-            val outerClasses = generateSequence(id) { classId ->
-                classId.outerClassId
-            }.mapTo(mutableListOf()) { provider.getFirClassifierByFqName(it)!! }
-            designation += outerClasses.asReversed()
-            if (nonLocalDeclarationToResolve is FirCallableDeclaration<*>) {
-                designation += nonLocalDeclarationToResolve
-            }
-        }
-        if (designation.all { it.resolvePhase >= toPhase }) {
-            return
-        }
-        val scopeSession = ScopeSession()
-        val transformer = FirDesignatedBodyResolveTransformerForIDE(
-            designation.iterator(), containerFirFile.session,
-            scopeSession,
-            implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
-            towerDataContextForStatement
-        )
-        containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
-    }
-
-    /**
-     * Should be invoked under lock
-     */
-    private fun runLazyResolveWithoutLock(
-        firDeclarationToResolve: FirDeclaration,
-        containerFirFile: FirFile,
-        provider: FirProvider,
-        toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
-    ) {
-        val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
-        if (firDeclarationToResolve.resolvePhase < nonLazyPhase) {
-            firFileBuilder.runResolveWithoutLockNoPCECheck(
-                containerFirFile,
-                fromPhase = firDeclarationToResolve.resolvePhase,
-                toPhase = nonLazyPhase
-            )
-        }
-        if (toPhase <= nonLazyPhase) return
-        runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextForStatement)
-    }
-
-    fun runLazyResolveWithPCECheck(
-        firDeclarationToResolve: FirDeclaration,
-        cache: ModuleFileCache,
-        containerFirFile: FirFile,
-        provider: FirProvider,
-        toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
-    ) {
-        val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
-        if (firDeclarationToResolve.resolvePhase < nonLazyPhase) {
-            firFileBuilder.runResolveWithPCECheck(
-                containerFirFile,
-                fromPhase = firDeclarationToResolve.resolvePhase,
-                toPhase = nonLazyPhase,
-                cache = cache,
-            )
-        }
-        if (toPhase <= nonLazyPhase) return
-        firFileBuilder.runCustomResolveWithPCECheck(containerFirFile, cache) {
-            runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextForStatement)
-        }
-    }
-
-    private fun runLazyResolveWithoutPCECheck(
-        firDeclarationToResolve: FirDeclaration,
-        cache: ModuleFileCache,
-        containerFirFile: FirFile,
-        provider: FirProvider,
-        toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
-    ) {
-        val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
-        if (firDeclarationToResolve.resolvePhase < nonLazyPhase) {
-            firFileBuilder.runResolveWithLock(
-                containerFirFile,
-                fromPhase = firDeclarationToResolve.resolvePhase,
-                toPhase = nonLazyPhase,
-                cache = cache,
-            )
-        }
-        if (toPhase <= nonLazyPhase) return
-        firFileBuilder.runCustomResolveUnderLock(containerFirFile, cache) {
-            runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextForStatement)
-        }
-    }
-
-
-    //TODO for completion only
-    fun runLazyResolveForCompletion(
-        firFunction: FirFunction<*>,
-        containerFirFile: FirFile,
-        firIdeProvider: FirIdeProvider,
-        toPhase: FirResolvePhase,
-        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>
-    ) {
-        runLazyResolveWithoutLock(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextForStatement)
-    }
-}
-
-private fun FirDeclaration.getNonLocalDeclarationToResolve(provider: FirProvider): FirDeclaration {
-    if (this is FirFile) return this
-    val ktDeclaration = psi as? KtDeclaration ?: error("FirDeclaration should have a PSI of type KtDeclaration")
-    if (!KtPsiUtil.isLocal(ktDeclaration)) return this
-    return when (val nonLocalPsi = ktDeclaration.getNonLocalContainingDeclarationWithFqName()) {
-        is KtClassOrObject -> provider.getFirClassifierByFqName(
-            nonLocalPsi.classIdIfNonLocal()
-                ?: error("Container classId should not be null for non-local declaration")
-        ) ?: error("Could not find class ${nonLocalPsi.classIdIfNonLocal()}")
-        is KtProperty, is KtNamedFunction -> {
-            val containerClass = nonLocalPsi.containingClassOrObject
-                ?: error("Container class should not be null for non-local declaration")
-            val containerClassId = containerClass.classIdIfNonLocal()
-                ?: error("Container classId should not be null for non-local declaration")
-            val containingFir = provider.getFirClassifierByFqName(containerClassId) as? FirRegularClass
-                ?: error("Could not find class $containerClassId")
-            containingFir.declarations.first { it.psi === nonLocalPsi }
-        }
-        else -> error("Invalid container ${nonLocalPsi?.let { it::class } ?: "null"}")
-    }
 }
 
 
-private fun KtElement.getNonLocalContainingDeclarationWithFqName(): KtNamedDeclaration? {
+fun KtElement.getNonLocalContainingDeclarationWithFqName(): KtNamedDeclaration? {
     var container = parent
     while (container != null && container !is KtFile) {
         if (container is KtNamedDeclaration

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/FirFileBuilder.kt

@@ -39,13 +39,14 @@ internal class FirFileBuilder(
     fun getFirFileResolvedToPhaseWithCaching(
         ktFile: KtFile,
         cache: ModuleFileCache,
-        @Suppress("SameParameterValue") toPhase: FirResolvePhase
+        @Suppress("SameParameterValue") toPhase: FirResolvePhase,
+        checkPCE: Boolean
     ): FirFile {
         val firFile = buildRawFirFileWithCaching(ktFile, cache)
         if (toPhase > FirResolvePhase.RAW_FIR) {
             cache.firFileLockProvider.withLock(firFile) {
                 //add lock for implit type resolve phase & super type
-                runResolveWithoutLockNoPCECheck(firFile, fromPhase = firFile.resolvePhase, toPhase = toPhase)
+                runResolveWithoutLock(firFile, fromPhase = firFile.resolvePhase, toPhase = toPhase, checkPCE = checkPCE)
             }
         }
         return firFile
@@ -62,38 +63,25 @@ internal class FirFileBuilder(
         return lock.lockWithPCECheck(LOCKING_INTERVAL_MS) { resolve() }
     }
 
-    fun runResolveWithLock(firFile: FirFile, cache: ModuleFileCache, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
-        cache.firFileLockProvider.withLock(firFile) {
-            runResolveWithoutLockNoPCECheck(firFile, fromPhase, toPhase)
-        }
-    }
-
-    fun runResolveWithPCECheck(firFile: FirFile, cache: ModuleFileCache, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
-        val lock = cache.firFileLockProvider.getLockFor(firFile)
-        lock.lockWithPCECheck(LOCKING_INTERVAL_MS) {
-            val scopeSession = ScopeSession()
-            var currentPhase = fromPhase
-            while (currentPhase < toPhase) {
-                checkCanceled()
-                currentPhase = currentPhase.next
-                firPhaseRunner.runPhase(firFile, currentPhase, scopeSession)
-            }
-        }
-    }
-
-
-    fun runResolveWithoutLockNoPCECheck(firFile: FirFile, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
+    fun runResolveWithoutLock(
+        firFile: FirFile,
+        fromPhase: FirResolvePhase,
+        toPhase: FirResolvePhase,
+        checkPCE: Boolean
+    ) {
         assert(fromPhase <= toPhase) {
             "Trying to resolve file ${firFile.name} from $fromPhase to $toPhase"
         }
         val scopeSession = ScopeSession()
         var currentPhase = fromPhase
         while (currentPhase < toPhase) {
+            if (checkPCE) checkCanceled()
             currentPhase = currentPhase.next
             firPhaseRunner.runPhase(firFile, currentPhase, scopeSession)
         }
     }
 
+
     companion object {
         private const val LOCKING_INTERVAL_MS = 500L
     }

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

@@ -0,0 +1,140 @@
+/*
+ * Copyright 2010-2020 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.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
+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.idea.fir.low.level.api.element.builder.FirDesignatedBodyResolveTransformerForIDE
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingDeclarationWithFqName
+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.providers.firIdeProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
+import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+
+internal class FirLazyDeclarationResolver(
+    private val firFileBuilder: FirFileBuilder
+) {
+    fun lazyResolveDeclaration(
+        declaration: FirDeclaration,
+        moduleFileCache: ModuleFileCache,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
+        checkPCE: Boolean = false
+    ) {
+        if (declaration.resolvePhase >= toPhase) return
+        val firFile = declaration.getContainingFile()
+            ?: error("FirFile was not found for\n${declaration.render()}")
+        val provider = firFile.session.firIdeProvider
+        if (checkPCE) {
+            firFileBuilder.runCustomResolveWithPCECheck(firFile, moduleFileCache) {
+                runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextForStatement, checkPCE = true)
+            }
+        } else {
+            firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
+                executeWithoutPCE {
+                    runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextForStatement, checkPCE = false)
+                }
+            }
+        }
+    }
+
+    fun runLazyResolveWithoutLock(
+        firDeclarationToResolve: FirDeclaration,
+        containerFirFile: FirFile,
+        provider: FirProvider,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>? = null,
+        checkPCE: Boolean
+    ) {
+        val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS)
+        if (firDeclarationToResolve.resolvePhase < nonLazyPhase) {
+            firFileBuilder.runResolveWithoutLock(
+                containerFirFile,
+                fromPhase = firDeclarationToResolve.resolvePhase,
+                toPhase = nonLazyPhase,
+                checkPCE = checkPCE
+            )
+        }
+        if (toPhase <= nonLazyPhase) return
+        if (checkPCE) checkCanceled()
+        runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextForStatement)
+    }
+
+    private fun runLazyResolvePhase(
+        firDeclarationToResolve: FirDeclaration,
+        containerFirFile: FirFile,
+        provider: FirProvider,
+        toPhase: FirResolvePhase,
+        towerDataContextForStatement: MutableMap<FirStatement, FirTowerDataContext>?,
+    ) {
+        val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider)
+
+        val designation = mutableListOf<FirDeclaration>(containerFirFile)
+        if (nonLocalDeclarationToResolve !is FirFile) {
+            val id = when (nonLocalDeclarationToResolve) {
+                is FirCallableDeclaration<*> -> {
+                    nonLocalDeclarationToResolve.symbol.callableId.classId
+                }
+                is FirRegularClass -> {
+                    nonLocalDeclarationToResolve.symbol.classId
+                }
+                else -> error("Unsupported: ${nonLocalDeclarationToResolve.render()}")
+            }
+            val outerClasses = generateSequence(id) { classId ->
+                classId.outerClassId
+            }.mapTo(mutableListOf()) { provider.getFirClassifierByFqName(it)!! }
+            designation += outerClasses.asReversed()
+            if (nonLocalDeclarationToResolve is FirCallableDeclaration<*>) {
+                designation += nonLocalDeclarationToResolve
+            }
+        }
+        if (designation.all { it.resolvePhase >= toPhase }) {
+            return
+        }
+        val scopeSession = ScopeSession()
+        val transformer = FirDesignatedBodyResolveTransformerForIDE(
+            designation.iterator(), containerFirFile.session,
+            scopeSession,
+            implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+            towerDataContextForStatement
+        )
+        containerFirFile.transform<FirFile, ResolutionMode>(transformer, ResolutionMode.ContextDependent)
+    }
+}
+
+
+private fun FirDeclaration.getNonLocalDeclarationToResolve(provider: FirProvider): FirDeclaration {
+    if (this is FirFile) return this
+    val ktDeclaration = psi as? KtDeclaration ?: error("FirDeclaration should have a PSI of type KtDeclaration")
+    if (!KtPsiUtil.isLocal(ktDeclaration)) return this
+    return when (val nonLocalPsi = ktDeclaration.getNonLocalContainingDeclarationWithFqName()) {
+        is KtClassOrObject -> provider.getFirClassifierByFqName(
+            nonLocalPsi.classIdIfNonLocal()
+                ?: error("Container classId should not be null for non-local declaration")
+        ) ?: error("Could not find class ${nonLocalPsi.classIdIfNonLocal()}")
+        is KtProperty, is KtNamedFunction -> {
+            val containerClass = nonLocalPsi.containingClassOrObject
+                ?: error("Container class should not be null for non-local declaration")
+            val containerClassId = containerClass.classIdIfNonLocal()
+                ?: error("Container classId should not be null for non-local declaration")
+            val containingFir = provider.getFirClassifierByFqName(containerClassId) as? FirRegularClass
+                ?: error("Could not find class $containerClassId")
+            containingFir.declarations.first { it.psi === nonLocalPsi }
+        }
+        else -> error("Invalid container ${nonLocalPsi?.let { it::class } ?: "null"}")
+    }
+}

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

@@ -25,9 +25,9 @@ import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
 import org.jetbrains.kotlin.fir.resolve.transformers.PhasedFirFileResolver
 import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.IdePhasedFirFileResolver
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCacheImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
@@ -62,7 +62,7 @@ internal class FirIdeSourcesSession private constructor(
             val searchScope = ModuleWithDependentsScope(project, dependentModules.map { it.module })
             return FirIdeSourcesSession(moduleInfo, sessionProvider, searchScope, firBuilder).apply {
                 val cache = ModuleFileCacheImpl(this)
-                val phasedFirFileResolver = IdePhasedFirFileResolver(FirElementBuilder(firFileBuilder), cache)
+                val phasedFirFileResolver = IdePhasedFirFileResolver(FirLazyDeclarationResolver(firFileBuilder), cache)
 
                 registerCommonComponents()
                 registerResolveComponents()

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.resolve.firProvider
 fun FirDeclaration.getContainingFile(): FirFile? {
     val provider = session.firProvider
     return when (this) {
+        is FirFile -> this
         is FirCallableDeclaration<*> -> provider.getFirCallableContainerFile(symbol)
         is FirClassLikeDeclaration<*> -> provider.getFirClassifierContainerFile(symbol)
         else -> error("Unsupported declaration ${this::class.java}")