[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 -> {

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

@@ -5,16 +5,16 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers
 
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.declarations.FirClass
 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
 
 object FirDummyCompilerLazyDeclarationResolver : FirLazyDeclarationResolver() {
     override fun startResolvingPhase(phase: FirResolvePhase) {}
     override fun finishResolvingPhase(phase: FirResolvePhase) {}
 
-    override fun lazyResolveToPhase(symbol: FirBasedSymbol<*>, toPhase: FirResolvePhase) {}
-    override fun lazyResolveToPhaseWithCallableMembers(symbol: FirClassSymbol<*>, toPhase: FirResolvePhase) {}
-    override fun lazyResolveToPhaseRecursively(symbol: FirBasedSymbol<*>, toPhase: FirResolvePhase) {}
+    override fun lazyResolveToPhase(element: FirElementWithResolveState, toPhase: FirResolvePhase) {}
+    override fun lazyResolveToPhaseWithCallableMembers(clazz: FirClass, toPhase: FirResolvePhase) {}
+    override fun lazyResolveToPhaseRecursively(element: FirElementWithResolveState, toPhase: FirResolvePhase) {}
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/symbols/FirLazyDeclarationResolver.kt

@@ -5,10 +5,10 @@
 
 package org.jetbrains.kotlin.fir.symbols
 
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.declarations.FirClass
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 
@@ -43,17 +43,17 @@ abstract class FirLazyDeclarationResolver : FirSessionComponent {
     /**
      * @see org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
      */
-    abstract fun lazyResolveToPhase(symbol: FirBasedSymbol<*>, toPhase: FirResolvePhase)
+    abstract fun lazyResolveToPhase(element: FirElementWithResolveState, toPhase: FirResolvePhase)
 
     /**
      * @see org.jetbrains.kotlin.fir.symbols.lazyResolveToPhaseWithCallableMembers
      */
-    abstract fun lazyResolveToPhaseWithCallableMembers(symbol: FirClassSymbol<*>, toPhase: FirResolvePhase)
+    abstract fun lazyResolveToPhaseWithCallableMembers(clazz: FirClass, toPhase: FirResolvePhase)
 
     /**
      * @see org.jetbrains.kotlin.fir.symbols.lazyResolveToPhaseRecursively
      */
-    abstract fun lazyResolveToPhaseRecursively(symbol: FirBasedSymbol<*>, toPhase: FirResolvePhase)
+    abstract fun lazyResolveToPhaseRecursively(element: FirElementWithResolveState, toPhase: FirResolvePhase)
 }
 
 class FirLazyResolveContractViolationException(
@@ -69,7 +69,7 @@ class FirLazyResolveContractViolationException(
 
 val FirSession.lazyDeclarationResolver: FirLazyDeclarationResolver by FirSession.sessionComponentAccessor()
 
-private val FirDeclaration.lazyDeclarationResolver get() = moduleData.session.lazyDeclarationResolver
+private val FirElementWithResolveState.lazyDeclarationResolver get() = moduleData.session.lazyDeclarationResolver
 
 /**
  * Lazy resolve [FirBasedSymbol] to [FirResolvePhase].
@@ -86,16 +86,16 @@ private val FirDeclaration.lazyDeclarationResolver get() = moduleData.session.la
  * @param toPhase the minimum phase, the declaration should be resolved to after an execution of the [lazyResolveToPhase]
  */
 fun FirBasedSymbol<*>.lazyResolveToPhase(toPhase: FirResolvePhase) {
-    fir.lazyDeclarationResolver.lazyResolveToPhase(this, toPhase)
+    fir.lazyResolveToPhase(toPhase)
 }
 
 /**
- * Lazy resolve [FirDeclaration] to [FirResolvePhase].
+ * Lazy resolve [FirElementWithResolveState] to [FirResolvePhase].
  *
  * @see lazyResolveToPhase
  */
-fun FirDeclaration.lazyResolveToPhase(toPhase: FirResolvePhase) {
-    symbol.lazyResolveToPhase(toPhase)
+fun FirElementWithResolveState.lazyResolveToPhase(toPhase: FirResolvePhase) {
+    lazyDeclarationResolver.lazyResolveToPhase(this, toPhase)
 }
 
 /**
@@ -117,7 +117,7 @@ fun FirDeclaration.lazyResolveToPhase(toPhase: FirResolvePhase) {
  * @see lazyResolveToPhase
  */
 fun FirClassSymbol<*>.lazyResolveToPhaseWithCallableMembers(toPhase: FirResolvePhase) {
-    fir.lazyDeclarationResolver.lazyResolveToPhaseWithCallableMembers(this, toPhase)
+    fir.lazyResolveToPhaseWithCallableMembers(toPhase)
 }
 
 /**
@@ -126,7 +126,7 @@ fun FirClassSymbol<*>.lazyResolveToPhaseWithCallableMembers(toPhase: FirResolveP
  * @see lazyResolveToPhaseWithCallableMembers
  */
 fun FirClass.lazyResolveToPhaseWithCallableMembers(toPhase: FirResolvePhase) {
-    symbol.lazyResolveToPhaseWithCallableMembers(toPhase)
+    lazyDeclarationResolver.lazyResolveToPhaseWithCallableMembers(this, toPhase)
 }
 
 /**
@@ -145,14 +145,14 @@ fun FirClass.lazyResolveToPhaseWithCallableMembers(toPhase: FirResolvePhase) {
  * @param toPhase the minimum phase, the declaration and all nested declarations should be resolved to after an execution of the [lazyResolveToPhase]
  */
 fun FirBasedSymbol<*>.lazyResolveToPhaseRecursively(toPhase: FirResolvePhase) {
-    fir.lazyDeclarationResolver.lazyResolveToPhaseRecursively(this, toPhase)
+    fir.lazyResolveToPhaseRecursively(toPhase)
 }
 
 /**
- * Lazy resolve [FirDeclaration] and all nested declarations to [FirResolvePhase].
+ * Lazy resolve [FirElementWithResolveState] and all nested declarations to [FirResolvePhase].
  *
  * @see lazyResolveToPhaseRecursively
  */
-fun FirDeclaration.lazyResolveToPhaseRecursively(toPhase: FirResolvePhase) {
-    symbol.lazyResolveToPhaseRecursively(toPhase)
+fun FirElementWithResolveState.lazyResolveToPhaseRecursively(toPhase: FirResolvePhase) {
+    lazyDeclarationResolver.lazyResolveToPhaseRecursively(this, toPhase)
 }

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirCompilerLazyDeclarationResolverWithPhaseChecking.kt

@@ -5,11 +5,11 @@
 
 package org.jetbrains.kotlin.test.frontend.fir.handlers
 
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.declarations.FirClass
 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.FirLazyResolveContractViolationException
-import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 
 
 class FirCompilerLazyDeclarationResolverWithPhaseChecking : FirLazyDeclarationResolver() {
@@ -20,15 +20,15 @@ class FirCompilerLazyDeclarationResolverWithPhaseChecking : FirLazyDeclarationRe
     fun getContractViolationExceptions(): List<FirLazyResolveContractViolationException> =
         exceptions
 
-    override fun lazyResolveToPhase(symbol: FirBasedSymbol<*>, toPhase: FirResolvePhase) {
+    override fun lazyResolveToPhase(element: FirElementWithResolveState, toPhase: FirResolvePhase) {
         checkIfCanLazyResolveToPhase(toPhase)
     }
 
-    override fun lazyResolveToPhaseWithCallableMembers(symbol: FirClassSymbol<*>, toPhase: FirResolvePhase) {
+    override fun lazyResolveToPhaseWithCallableMembers(clazz: FirClass, toPhase: FirResolvePhase) {
         checkIfCanLazyResolveToPhase(toPhase)
     }
 
-    override fun lazyResolveToPhaseRecursively(symbol: FirBasedSymbol<*>, toPhase: FirResolvePhase) {
+    override fun lazyResolveToPhaseRecursively(element: FirElementWithResolveState, toPhase: FirResolvePhase) {
         checkIfCanLazyResolveToPhase(toPhase)
     }