LL API: simplify FirModuleResolveState interface

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSessionProviderStorage
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
@@ -26,10 +27,10 @@ internal class FirIdeResolveStateService(project: Project) {
         project.createProjectWideOutOfBlockModificationTracker(),
         ProjectRootModificationTracker.getInstance(project),
     ) {
-        ConcurrentHashMap<KtModule, FirModuleResolveStateImpl>()
+        ConcurrentHashMap<KtModule, FirSourceModuleResolveState>()
     }
 
-    fun getResolveState(module: KtModule): FirModuleResolveStateImpl =
+    fun getResolveState(module: KtModule): FirSourceModuleResolveState =
         stateCache.computeIfAbsent(module) { createResolveStateFor(module, sessionProviderStorage) }
 
     companion object {
@@ -40,13 +41,13 @@ internal class FirIdeResolveStateService(project: Project) {
             module: KtModule,
             sessionProviderStorage: FirIdeSessionProviderStorage,
             configureSession: (FirIdeSession.() -> Unit)? = null,
-        ): FirModuleResolveStateImpl {
+        ): FirSourceModuleResolveState {
             if (module !is KtSourceModule) {
                 error("Creating FirModuleResolveState is not yet supported for $module")
             }
             val sessionProvider = sessionProviderStorage.getSessionProvider(module, configureSession)
             val firFileBuilder = sessionProvider.rootModuleSession.firFileBuilder
-            return FirModuleResolveStateImpl(
+            return FirSourceModuleResolveState(
                 sessionProviderStorage.project,
                 module,
                 sessionProvider,

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.KtToFirMapping
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.containingKtFileIfAny
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalKtFile
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
@@ -27,7 +28,7 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 
 internal class FirModuleResolveStateDepended(
-    val originalState: FirModuleResolveStateImpl,
+    val originalState: FirSourceModuleResolveState,
     val towerProviderBuiltUponElement: FirTowerContextProvider,
     private val ktToFirMapping: KtToFirMapping?,
 ) : FirModuleResolveState() {
@@ -78,15 +79,7 @@ internal class FirModuleResolveStateDepended(
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> =
         TODO("Diagnostics are not implemented for depended state")
 
-    override fun findSourceFirSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
+    override fun resolveToFirSymbol(ktDeclaration: KtDeclaration, phase: FirResolvePhase): FirBasedSymbol<*> {
-        return originalState.findSourceFirSymbol(ktDeclaration)
+        return originalState.resolveToFirSymbol(ktDeclaration, phase)
-    }
-
-    override fun findSourceFirSymbol(ktDeclaration: KtLambdaExpression): FirBasedSymbol<*> {
-        return originalState.findSourceFirSymbol(ktDeclaration)
-    }
-
-    override fun findSourceFirCompiledSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
-        return originalState.findSourceFirCompiledSymbol(ktDeclaration)
     }
 }

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

@@ -53,37 +53,7 @@ abstract class FirModuleResolveState {
 
     internal abstract fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic>
 
-    abstract fun findSourceFirSymbol(
+    abstract fun resolveToFirSymbol(ktDeclaration: KtDeclaration, phase: FirResolvePhase): FirBasedSymbol<*>
-        ktDeclaration: KtDeclaration,
-    ): FirBasedSymbol<*>
-
-    abstract fun findSourceFirSymbol(
-        ktDeclaration: KtLambdaExpression,
-    ): FirBasedSymbol<*>
-
-    /**
-     * Looks for compiled non-local [ktDeclaration] declaration by querying its classId/callableId from the SymbolProvider.
-     *
-     * Works only if [ktDeclaration] is compiled (i.e. comes from .class file).
-     */
-    abstract fun findSourceFirCompiledSymbol(
-        ktDeclaration: KtDeclaration
-    ): FirBasedSymbol<*>
-
-
-    open fun resolveToFirSymbol(
-        ktDeclaration: KtDeclaration,
-        resolveState: FirModuleResolveState,
-        phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
-    ): FirBasedSymbol<*> {
-        return when (val module = ktDeclaration.getKtModule()) {
-            is KtSourceModule -> resolveState.findSourceFirSymbol(ktDeclaration).also { resolveState.resolveFirToPhase(it.fir, phase) }
-            is KtLibraryModule -> resolveState.findSourceFirCompiledSymbol(ktDeclaration)
-            else -> error("unsupported module $module")
-        }
-    }
 
     internal abstract fun resolveFirToPhase(declaration: FirDeclaration, toPhase: FirResolvePhase)
-
-
 }

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

@@ -43,7 +43,7 @@ fun KtDeclaration.resolveToFirSymbol(
     resolveState: FirModuleResolveState,
     phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
 ): FirBasedSymbol<*> {
-    return resolveState.resolveToFirSymbol(this, resolveState, phase)
+    return resolveState.resolveToFirSymbol(this, phase)
 }
 
 /**

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

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.analysis.api.impl.barebone.parentOfType
 import org.jetbrains.kotlin.analysis.api.impl.barebone.parentsOfType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.DeclarationCopyBuilder.withBodyFrom
 import org.jetbrains.kotlin.analysis.low.level.api.fir.FirModuleResolveStateDepended
-import org.jetbrains.kotlin.analysis.low.level.api.fir.FirModuleResolveStateImpl
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FileTowerProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerDataContextAllElementsCollector
@@ -67,7 +67,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         place: T,
         elementToResolve: T,
     ): FirElement {
-        require(state is FirModuleResolveStateImpl)
+        require(state is FirSourceModuleResolveState)
         require(place.isPhysical)
 
         val declaration = runBodyResolveOnAir(
@@ -95,7 +95,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         state: FirModuleResolveState,
         place: KtElement,
     ): FirTowerContextProvider {
-        require(state is FirModuleResolveStateImpl)
+        require(state is FirSourceModuleResolveState)
 
         return if (place is KtFile) {
             FileTowerProvider(place, onAirGetTowerContextForFile(state, place))
@@ -116,7 +116,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
     }
 
     private fun onAirGetTowerContextForFile(
-        state: FirModuleResolveStateImpl,
+        state: FirSourceModuleResolveState,
         file: KtFile,
     ): FirTowerDataContext {
         require(file.isPhysical)
@@ -144,7 +144,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         originalKtFile: KtFile,
         elementToAnalyze: KtElement
     ): FirModuleResolveState {
-        require(originalState is FirModuleResolveStateImpl)
+        require(originalState is FirSourceModuleResolveState)
         require(elementToAnalyze !is KtFile) { "KtFile for dependency element not supported" }
         require(!elementToAnalyze.isPhysical) { "Depended state should be build only for non-physical elements" }
 
@@ -178,7 +178,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
 
     private fun tryResolveAsFileAnnotation(
         annotationEntry: KtAnnotationEntry,
-        state: FirModuleResolveStateImpl,
+        state: FirSourceModuleResolveState,
         replacement: RawFirReplacement,
         firFile: FirFile,
         collector: FirTowerDataContextCollector? = null,
@@ -202,7 +202,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
     }
 
     private fun runBodyResolveOnAir(
-        state: FirModuleResolveStateImpl,
+        state: FirSourceModuleResolveState,
         replacement: RawFirReplacement,
         onAirCreatedDeclaration: Boolean,
         collector: FirTowerDataContextCollector? = null,

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

@@ -0,0 +1,10 @@
+/*
+ * Copyright 2010-2022 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.analysis.low.level.api.fir.state
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
+
+abstract class FirPhysicalModuleResolveState : FirModuleResolveState()

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

@@ -1,13 +1,12 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2022 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.analysis.low.level.api.fir
+package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.DiagnosticsCollector
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirElementBuilder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
@@ -22,6 +21,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirDeclarationForCom
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getElementTextInContext
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
+import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
@@ -38,13 +38,13 @@ import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.psi.*
 
-internal class FirModuleResolveStateImpl(
+internal class FirSourceModuleResolveState(
     override val project: Project,
-    override val module: KtModule,
+    override val module: KtSourceModule,
     private val sessionProvider: FirIdeSessionProvider,
     val firFileBuilder: FirFileBuilder,
     val firLazyDeclarationResolver: FirLazyDeclarationResolver,
-) : FirModuleResolveState() {
+) : FirPhysicalModuleResolveState() {
     override val rootModuleSession: FirIdeSourcesSession get() = sessionProvider.rootModuleSession
 
     /**
@@ -82,30 +82,28 @@ internal class FirModuleResolveStateImpl(
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> =
         diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
 
-    override fun findSourceFirSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
+    override fun resolveToFirSymbol(
+        ktDeclaration: KtDeclaration,
+        phase: FirResolvePhase
+    ): FirBasedSymbol<*> {
+        return when (val module = ktDeclaration.getKtModule()) {
+            is KtSourceModule -> findSourceFirSymbol(ktDeclaration).also { resolveFirToPhase(it.fir, phase) }
+            is KtLibraryModule -> findFirCompiledSymbol(ktDeclaration)
+            else -> error("unsupported module $module")
+        }
+    }
+
+    private fun findSourceFirSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
         return findSourceFirDeclarationByExpression(ktDeclaration.originalDeclaration ?: ktDeclaration).symbol
     }
 
-    override fun findSourceFirSymbol(ktDeclaration: KtLambdaExpression): FirBasedSymbol<*> {
+    private fun findFirCompiledSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
-        return findSourceFirDeclarationByExpression(ktDeclaration).symbol
-    }
-
-    override fun findSourceFirCompiledSymbol(ktDeclaration: KtDeclaration): FirBasedSymbol<*> {
         require(ktDeclaration.containingKtFile.isCompiled) {
             "This method will only work on compiled declarations, but this declaration is not compiled: ${ktDeclaration.getElementTextInContext()}"
         }
 
         val searcher = FirDeclarationForCompiledElementSearcher(rootModuleSession.symbolProvider)
-
+        val firDeclaration = searcher.findNonLocalDeclaration(ktDeclaration)
-        val firDeclaration = when (ktDeclaration) {
-            is KtEnumEntry -> searcher.findNonLocalEnumEntry(ktDeclaration)
-            is KtClassOrObject -> searcher.findNonLocalClass(ktDeclaration)
-            is KtConstructor<*> -> searcher.findConstructorOfNonLocalClass(ktDeclaration)
-            is KtNamedFunction -> searcher.findNonLocalFunction(ktDeclaration)
-            is KtProperty -> searcher.findNonLocalProperty(ktDeclaration)
-
-            else -> error("Unsupported compiled declaration of type ${ktDeclaration::class}: ${ktDeclaration.getElementTextInContext()}")
-        }
         return firDeclaration.symbol
     }
 

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

@@ -21,8 +21,19 @@ import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
  * Allows to search for FIR declarations by compiled [KtDeclaration]s.
  */
 internal class FirDeclarationForCompiledElementSearcher(private val symbolProvider: FirSymbolProvider) {
+    fun findNonLocalDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
+        return when (ktDeclaration) {
+            is KtEnumEntry -> findNonLocalEnumEntry(ktDeclaration)
+            is KtClassOrObject -> findNonLocalClass(ktDeclaration)
+            is KtConstructor<*> -> findConstructorOfNonLocalClass(ktDeclaration)
+            is KtNamedFunction -> findNonLocalFunction(ktDeclaration)
+            is KtProperty -> findNonLocalProperty(ktDeclaration)
 
-    fun findNonLocalEnumEntry(declaration: KtEnumEntry): FirEnumEntry {
+            else -> error("Unsupported compiled declaration of type ${ktDeclaration::class}: ${ktDeclaration.getElementTextInContext()}")
+        }
+    }
+
+    private fun findNonLocalEnumEntry(declaration: KtEnumEntry): FirEnumEntry {
         require(!declaration.isLocal)
         val classId = declaration.containingClassOrObject?.getClassId()
             ?: error("Non-local class should have classId. The class is ${declaration.getElementTextInContext()}")
@@ -35,7 +46,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
         } as FirEnumEntry
     }
 
-    fun findNonLocalClass(declaration: KtClassOrObject): FirClassLikeDeclaration {
+    private fun findNonLocalClass(declaration: KtClassOrObject): FirClassLikeDeclaration {
         require(!declaration.isLocal)
         val classId = declaration.getClassId()
             ?: error("Non-local class should have classId. The class is ${declaration.getElementTextInContext()}")
@@ -46,7 +57,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
         return classCandidate.fir
     }
 
-    fun findConstructorOfNonLocalClass(declaration: KtConstructor<*>): FirConstructor {
+    private fun findConstructorOfNonLocalClass(declaration: KtConstructor<*>): FirConstructor {
         val containingClass = declaration.containingClassOrObject
             ?: error("Constructor must have outer class: ${declaration.getElementTextInContext()}")
 
@@ -62,7 +73,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
         return constructorCandidate.fir
     }
 
-    fun findNonLocalFunction(declaration: KtNamedFunction): FirFunction {
+    private fun findNonLocalFunction(declaration: KtNamedFunction): FirFunction {
         require(!declaration.isLocal)
 
         val functionCandidate =
@@ -73,7 +84,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
         return functionCandidate.fir
     }
 
-    fun findNonLocalProperty(declaration: KtProperty): FirProperty {
+    private fun findNonLocalProperty(declaration: KtProperty): FirProperty {
         require(!declaration.isLocal)
 
         val propertyCandidate =

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirRenderer
@@ -53,7 +54,7 @@ abstract class AbstractFirLazyDeclarationResolveTest : AbstractLowLevelApiSingle
         val rendererOption = FirRenderer.RenderMode.WithDeclarationAttributes.copy(renderDeclarationResolvePhase = true)
         val resultBuilder = StringBuilder()
         resolveWithClearCaches(ktFile) { firModuleResolveState ->
-            check(firModuleResolveState is FirModuleResolveStateImpl)
+            check(firModuleResolveState is FirSourceModuleResolveState)
             val declarationToResolve = firModuleResolveState
                 .getOrBuildFirFile(ktFile)
                 .findResolveMe()
@@ -67,7 +68,7 @@ abstract class AbstractFirLazyDeclarationResolveTest : AbstractLowLevelApiSingle
         }
 
         resolveWithClearCaches(ktFile) { firModuleResolveState ->
-            check(firModuleResolveState is FirModuleResolveStateImpl)
+            check(firModuleResolveState is FirSourceModuleResolveState)
             val firFile = firModuleResolveState.getOrBuildFirFile(ktFile)
             firFile.ensureResolved(FirResolvePhase.BODY_RESOLVE)
             resultBuilder.append("\nFILE RAW TO BODY:\n")

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LowLevelFirApiFacadeForResolveOnAir
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.render
@@ -32,7 +33,7 @@ abstract class AbstractFirOnAirResolveTest : AbstractLowLevelApiSingleFileTest()
         check(place::class == onAir::class)
 
         resolveWithClearCaches(ktFile) { resolveState ->
-            check(resolveState is FirModuleResolveStateImpl)
+            check(resolveState is FirSourceModuleResolveState)
             val firElement = LowLevelFirApiFacadeForResolveOnAir.onAirResolveElement(resolveState, place, onAir)
             val rendered = firElement.render(FirRenderer.RenderMode.WithResolvePhases)
             testServices.assertions.assertEqualsToTestDataFileSibling(rendered)

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

@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStack
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
-import org.jetbrains.kotlin.analysis.low.level.api.fir.FirModuleResolveStateImpl
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getDiagnostics
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
@@ -41,7 +41,7 @@ abstract class AbstractFirContextCollectionTest : AbstractLowLevelApiSingleFileT
                 register(BeforeElementDiagnosticCollectionHandler::class, handler)
             }
         ) { resolveState ->
-            check(resolveState is FirModuleResolveStateImpl)
+            check(resolveState is FirSourceModuleResolveState)
 
             val fileStructure = resolveState.fileStructureCache.getFileStructure(ktFile, resolveState.rootModuleSession.cache)
             val allStructureElements = fileStructure.getAllStructureElements()

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/AbstractFileStructureTest.kt

@@ -9,7 +9,7 @@ import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
 import com.intellij.psi.impl.source.tree.LeafPsiElement
-import org.jetbrains.kotlin.analysis.low.level.api.fir.FirModuleResolveStateImpl
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.FirSourceModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.psi.*
@@ -85,7 +85,7 @@ abstract class AbstractFileStructureTest : AbstractLowLevelApiSingleFileTest() {
     }
 
     private fun KtFile.getFileStructure(): FileStructure {
-        val moduleResolveState = getResolveState() as FirModuleResolveStateImpl
+        val moduleResolveState = getResolveState() as FirSourceModuleResolveState
         return moduleResolveState.fileStructureCache.getFileStructure(
             ktFile = this,
             moduleFileCache = moduleResolveState.rootModuleSession.cache