[FIR] make FirLazyDeclarationResolver FirElementWithResolveState-based

This is required to be able to call lazy resolve for FirFileAnnotationsContainer

^KT-63042

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/LLFirLazyDeclarationResolver.kt

@@ -6,32 +6,31 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.ThreadSafeMutableState
+import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 
 @ThreadSafeMutableState
 internal class LLFirLazyDeclarationResolver : FirLazyDeclarationResolver() {
     override fun startResolvingPhase(phase: FirResolvePhase) {}
     override fun finishResolvingPhase(phase: FirResolvePhase) {}
 
-    override fun lazyResolveToPhase(symbol: FirBasedSymbol<*>, toPhase: FirResolvePhase) {
-        val fir = symbol.fir
-        val session = fir.moduleData.session
+    override fun lazyResolveToPhase(element: FirElementWithResolveState, toPhase: FirResolvePhase) {
+        val session = element.moduleData.session
         if (session !is LLFirResolvableModuleSession) return
         val moduleComponents = session.moduleComponents
         moduleComponents.firModuleLazyDeclarationResolver.lazyResolve(
-            target = fir,
+            target = element,
             scopeSession = moduleComponents.scopeSessionProvider.getScopeSession(),
             toPhase = toPhase,
         )
     }
 
-    override fun lazyResolveToPhaseWithCallableMembers(symbol: FirClassSymbol<*>, toPhase: FirResolvePhase) {
-        val fir = symbol.fir as? FirRegularClass ?: return
+    override fun lazyResolveToPhaseWithCallableMembers(clazz: FirClass, toPhase: FirResolvePhase) {
+        val fir = clazz as? FirRegularClass ?: return
         val session = fir.moduleData.session
         if (session !is LLFirResolvableModuleSession) return
         val moduleComponents = session.moduleComponents
@@ -42,13 +41,12 @@ internal class LLFirLazyDeclarationResolver : FirLazyDeclarationResolver() {
         )
     }
 
-    override fun lazyResolveToPhaseRecursively(symbol: FirBasedSymbol<*>, toPhase: FirResolvePhase) {
-        val fir = symbol.fir
-        val session = fir.moduleData.session
+    override fun lazyResolveToPhaseRecursively(element: FirElementWithResolveState, toPhase: FirResolvePhase) {
+        val session = element.moduleData.session
         if (session !is LLFirResolvableModuleSession) return
         val moduleComponents = session.moduleComponents
         moduleComponents.firModuleLazyDeclarationResolver.lazyResolveRecursively(
-            target = fir,
+            target = element,
             scopeSession = moduleComponents.scopeSessionProvider.getScopeSession(),
             toPhase = toPhase,
         )

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructure.kt

@@ -202,14 +202,7 @@ internal class FileStructure private constructor(
         container is KtFile -> {
             val firFile = moduleComponents.firFileBuilder.buildRawFirFileWithCaching(ktFile)
             firFile.lazyResolveToPhase(FirResolvePhase.IMPORTS)
-
-            firFile.annotationsContainer?.let {
-                moduleComponents.firModuleLazyDeclarationResolver.lazyResolve(
-                    target = it,
-                    scopeSession = moduleComponents.scopeSessionProvider.getScopeSession(),
-                    FirResolvePhase.BODY_RESOLVE,
-                )
-            }
+            firFile.annotationsContainer?.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
 
             RootStructureElement(firFile, moduleComponents)
         }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyDeclarationResolveTestCase.kt

@@ -51,13 +51,8 @@ abstract class AbstractFirLazyDeclarationResolveTestCase : AbstractLowLevelApiLa
     ): Pair<FirElementWithResolveState, ((FirResolvePhase) -> Unit)> = when {
         Directives.RESOLVE_FILE_ANNOTATIONS in moduleStructure.allDirectives -> {
             val annotationContainer = firResolveSession.getOrBuildFirFile(ktFile).annotationsContainer!!
-            val session = annotationContainer.moduleData.session as LLFirResolvableModuleSession
             annotationContainer to fun(phase: FirResolvePhase) {
-                session.moduleComponents.firModuleLazyDeclarationResolver.lazyResolve(
-                    annotationContainer,
-                    session.getScopeSession(),
-                    phase,
-                )
+                annotationContainer.lazyResolveToPhase(phase)
             }
         }
         Directives.RESOLVE_FILE in moduleStructure.allDirectives -> {