FIR IDE: run lazy resolve on declaration when getting phased fir

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers
 
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSymbolOwner
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
@@ -28,16 +27,9 @@ fun <D> AbstractFirBasedSymbol<D>.phasedFir(
             "Incorrect resolvePhase: actual: $availablePhase, expected: $requiredPhase\n For: ${fir.render()}"
         }
 
-        val containingFile = when (this) {
-            is FirCallableSymbol<*> -> provider.getFirCallableContainerFile(this)
-            is FirClassLikeSymbol<*> -> provider.getFirClassifierContainerFile(this)
-            else -> null
-        }
-            ?: throw AssertionError("Cannot get container file by symbol: $this (${result.render()})")
-
         val resolver = fir.session.phasedFirFileResolver
             ?: error("phasedFirFileResolver should be defined when working with FIR in phased mode")
-        resolver.resolveFile(containingFile, fromPhase = availablePhase, toPhase = requiredPhase)
+        resolver.resolveDeclaration(fir, fromPhase = availablePhase, toPhase = requiredPhase)
     }
     return result
 }

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

@@ -7,11 +7,12 @@ package org.jetbrains.kotlin.fir.resolve.transformers
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 
 abstract class PhasedFirFileResolver : FirSessionComponent {
-    abstract fun resolveFile(firFile: FirFile, fromPhase: FirResolvePhase, toPhase: FirResolvePhase)
+    abstract fun resolveDeclaration(declaration: FirDeclaration, fromPhase: FirResolvePhase, toPhase: FirResolvePhase)
 }
 
 internal val FirSession.phasedFirFileResolver: PhasedFirFileResolver? by FirSession.nullableSessionComponentAccessor()

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

@@ -5,17 +5,19 @@
 
 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
 
 internal class IdePhasedFirFileResolver(
-    private val firFileBuilder: FirFileBuilder,
+    private val firFileBuilder: FirElementBuilder,
     private val cache: ModuleFileCache
 ) : PhasedFirFileResolver() {
-    override fun resolveFile(firFile: FirFile, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
-        firFileBuilder.runResolveWithLock(firFile, cache, fromPhase, toPhase)
+    override fun resolveDeclaration(declaration: FirDeclaration, fromPhase: FirResolvePhase, toPhase: FirResolvePhase) {
+        firFileBuilder.lazyResolveDeclarationNoPCECheck(declaration, cache, toPhase)
     }
 }

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

@@ -68,6 +68,18 @@ internal class FirElementBuilder(
         }
     }
 
+    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,
@@ -154,6 +166,30 @@ internal class FirElementBuilder(
         }
     }
 
+    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<*>,

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

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 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
@@ -61,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(firBuilder, cache)
+                val phasedFirFileResolver = IdePhasedFirFileResolver(FirElementBuilder(firFileBuilder), cache)
 
                 registerCommonComponents()
                 registerResolveComponents()