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

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()