[LL API] Migrate 'getKtModule()' usages to 'ProjectStructureProvider'

Ensure the contextual module is passed wherever possible. ^KT-57559 Fixed
2023-04-18 17:25:34 +09:00
parent c90d094af6
commit e9c077220f
31 changed files with 187 additions and 133 deletions
@@ -11,7 +11,7 @@ import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.container.ComponentProvider
 import org.jetbrains.kotlin.container.get
@@ -34,7 +34,7 @@ import org.jetbrains.kotlin.resolve.lazy.ResolveSession
 import org.jetbrains.kotlin.types.checker.KotlinTypeRefiner
 import org.jetbrains.kotlin.util.CancellationChecker

-class CliFe10AnalysisFacade(private val project: Project) : Fe10AnalysisFacade {
+class CliFe10AnalysisFacade : Fe10AnalysisFacade {
    override fun getResolveSession(element: KtElement): ResolveSession {
        return getHandler(element).resolveSession ?: error("Resolution is not performed")
    }
@@ -68,7 +68,7 @@ class CliFe10AnalysisFacade(private val project: Project) : Fe10AnalysisFacade {
    }

    private fun getHandler(useSiteElement: KtElement): KtFe10AnalysisHandlerExtension {
-        val ktModule = useSiteElement.getKtModule(project)
+        val ktModule = ProjectStructureProvider.getModule(useSiteElement, contextualModule = null)
        return KtFe10AnalysisHandlerExtension.getInstance(ktModule.project, ktModule)
    }
 }
@@ -16,7 +16,7 @@ import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolProvider
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolProviderByJavaPsi
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile

@@ -28,7 +28,7 @@ class KtFe10AnalysisSession(
 ) : KtAnalysisSession(analysisContext.token) {
    constructor(contextElement: KtElement, token: KtLifetimeToken) : this(
        Fe10AnalysisContext(Fe10AnalysisFacade.getInstance(contextElement.project), contextElement, token),
-        contextElement.getKtModule()
+        ProjectStructureProvider.getModule(contextElement, contextualModule = null)
    )


@@ -9,11 +9,8 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
-import org.jetbrains.kotlin.analysis.api.descriptors.symbols.base.KtFe10Symbol
-import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenFactory
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.psi.KtElement
 import java.lang.UnsupportedOperationException

@@ -14,17 +14,14 @@ import org.jetbrains.kotlin.analysis.api.descriptors.components.base.Fe10KtAnaly
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.KtFe10DescDefaultBackingFieldSymbol
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.getDescriptor
 import org.jetbrains.kotlin.analysis.api.descriptors.symbols.descriptorBased.base.toKtSymbol
+import org.jetbrains.kotlin.analysis.api.getModule
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.symbols.KtBackingFieldSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtDeclarationSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolKind
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithKind
-import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
-import org.jetbrains.kotlin.analysis.project.structure.KtLibrarySourceModule
-import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
@@ -54,39 +51,42 @@ internal class KtFe10SymbolContainingDeclarationProvider(

    // TODO this is a dummy and incorrect implementation just to satisfy some tests
    override fun getContainingModule(symbol: KtSymbol): KtModule {
-        val psiForModule = symbol.getDescriptor()?.let { DescriptorToSourceUtils.getContainingFile(it) }
-            ?: symbol.psi
+        val descriptor = symbol.getDescriptor()

-        return psiForModule?.getKtModule(analysisSession.analysisContext.resolveSession.project)
-            ?: symbol.getDescriptor()?.getFakeContainingKtModule()
-            ?: (symbol as? KtBackingFieldSymbol)?.owningProperty?.let { getContainingModule(it) }
-            ?: TODO(symbol::class.java.name)
+        val symbolPsi = descriptor?.let(DescriptorToSourceUtils::getContainingFile) ?: symbol.psi
+        if (symbolPsi != null) {
+            return analysisSession.getModule(symbolPsi)
+        }
+
+        if (descriptor is DescriptorWithContainerSource) {
+            return getFakeContainingKtModule(descriptor)
+        }
+
+        if (symbol is KtBackingFieldSymbol) {
+            return getContainingModule(symbol.owningProperty)
+        }
+
+        TODO(symbol::class.java.name)
    }

-    private fun DeclarationDescriptor.getFakeContainingKtModule(): KtModule? {
-        return when (this) {
-            is DescriptorWithContainerSource -> {
-                val libraryPath = Paths.get((containerSource as JvmPackagePartSource).knownJvmBinaryClass?.containingLibrary!!)
-                object : KtLibraryModule {
-                    override val libraryName: String = libraryPath.fileName.toString().substringBeforeLast(".")
-                    override val librarySources: KtLibrarySourceModule? = null
-                    override fun getBinaryRoots(): Collection<Path> = listOf(libraryPath)
-                    override val directRegularDependencies: List<KtModule> = emptyList()
-                    override val directDependsOnDependencies: List<KtModule> = emptyList()
-                    override val transitiveDependsOnDependencies: List<KtModule> = emptyList()
-                    override val directFriendDependencies: List<KtModule> = emptyList()
-                    override val contentScope: GlobalSearchScope = ProjectScope.getLibrariesScope(project)
-                    override val platform: TargetPlatform
-                        get() = this@getFakeContainingKtModule.platform!!
-                    override val analyzerServices: PlatformDependentAnalyzerServices
-                        get() = JvmPlatformAnalyzerServices
-                    override val project: Project
-                        get() = analysisSession.analysisContext.resolveSession.project
+    private fun getFakeContainingKtModule(descriptor: DescriptorWithContainerSource): KtModule {
+        val libraryPath = Paths.get((descriptor.containerSource as JvmPackagePartSource).knownJvmBinaryClass?.containingLibrary!!)
+        return object : KtLibraryModule {
+            override val libraryName: String = libraryPath.fileName.toString().substringBeforeLast(".")
+            override val librarySources: KtLibrarySourceModule? = null
+            override fun getBinaryRoots(): Collection<Path> = listOf(libraryPath)
+            override val directRegularDependencies: List<KtModule> = emptyList()
+            override val directDependsOnDependencies: List<KtModule> = emptyList()
+            override val transitiveDependsOnDependencies: List<KtModule> = emptyList()
+            override val directFriendDependencies: List<KtModule> = emptyList()
+            override val contentScope: GlobalSearchScope = ProjectScope.getLibrariesScope(project)
+            override val platform: TargetPlatform
+                get() = descriptor.platform!!
+            override val analyzerServices: PlatformDependentAnalyzerServices
+                get() = JvmPlatformAnalyzerServices
+            override val project: Project
+                get() = analysisSession.analysisContext.resolveSession.project

-                }
-            }
-
-            else -> null
        }
    }
 }
@@ -37,7 +37,7 @@ object AnalysisApiFe10TestServiceRegistrar : AnalysisApiTestServiceRegistrar() {
    override fun registerProjectModelServices(project: MockProject, testServices: TestServices) {
        project.apply {
            registerService(KtAnalysisSessionProvider::class.java, KtFe10AnalysisSessionProvider(project))
-            registerService(Fe10AnalysisFacade::class.java, CliFe10AnalysisFacade(project))
+            registerService(Fe10AnalysisFacade::class.java, CliFe10AnalysisFacade())
            registerService(ModuleVisibilityManager::class.java, CliModuleVisibilityManagerImpl(enabled = true))

            registerService(ReadWriteAccessChecker::class.java, ReadWriteAccessCheckerDescriptorsImpl())
@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenFactory
 import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
 import org.jetbrains.kotlin.psi.KtElement
 import java.util.concurrent.ConcurrentMap
@@ -36,7 +36,8 @@ class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider
    }

    override fun getAnalysisSession(useSiteKtElement: KtElement, factory: KtLifetimeTokenFactory): KtAnalysisSession {
-        return getAnalysisSessionByUseSiteKtModule(useSiteKtElement.getKtModule(project), factory)
+        val module = ProjectStructureProvider.getModule(useSiteKtElement, contextualModule = null)
+        return getAnalysisSessionByUseSiteKtModule(module, factory)
    }

    override fun getAnalysisSessionByUseSiteKtModule(useSiteKtModule: KtModule, factory: KtLifetimeTokenFactory): KtAnalysisSession {
@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayOfSymbol
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayTypeToArrayOfCall
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirFunctionSymbol
+import org.jetbrains.kotlin.analysis.api.getModule
 import org.jetbrains.kotlin.analysis.api.impl.base.components.AbstractKtCallResolver
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.signatures.KtCallableSignature
@@ -1290,7 +1291,7 @@ internal class KtFirCallResolver(
    @KtAnalysisApiInternals
    override fun unresolvedKtCallError(psi: KtElement): KtErrorCallInfo {
        LOG.logErrorWithAttachment("${psi::class.simpleName} should always resolve to a KtCallInfo") {
-            withPsiEntry("psi", psi)
+            withPsiEntry("psi", psi, analysisSession::getModule)
            provideAdditionalAttachmentToUnresolvedCall(psi, this)
        }

@@ -1309,7 +1310,7 @@ internal class KtFirCallResolver(
                "Error during resolving call ${element::class.java.name}",
                exception = e,
            ) {
-                withPsiEntry("psi", element)
+                withPsiEntry("psi", element, analysisSession::getModule)
                element.getOrBuildFir(firResolveSession)?.let { withFirEntry("fir", it) }
            }
        }
@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.analysis.api.base.KtContextReceiver
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.annotations.KtFirAnnotationListForDeclaration
 import org.jetbrains.kotlin.analysis.api.fir.utils.cached
+import org.jetbrains.kotlin.analysis.api.getModule
 import org.jetbrains.kotlin.analysis.api.impl.base.annotations.KtEmptyAnnotationsList
 import org.jetbrains.kotlin.analysis.api.impl.base.symbols.toKtClassKind
 import org.jetbrains.kotlin.analysis.api.lifetime.withValidityAssertion
@@ -21,7 +22,6 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.analysis.api.symbols.KtTypeParameterSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolKind
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.firClassByPsiClassProvider
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.analysis.utils.classIdIfNonLocal
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibility
@@ -116,7 +116,7 @@ internal class KtFirPsiJavaClassSymbol(
    override val hasLazyFirSymbol: Boolean get() = true

    override val firSymbol: FirRegularClassSymbol by cached {
-        val module = psi.getKtModule(analysisSession.project)
+        val module = analysisSession.getModule(psi)
        val provider = analysisSession.firResolveSession.getSessionFor(module).firClassByPsiClassProvider
        val firClassSymbol = provider.getFirClass(psi)

@@ -7,8 +7,8 @@ package org.jetbrains.kotlin.analysis.api.fir.symbols

 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.components.KtFirAnalysisSessionComponent
+import org.jetbrains.kotlin.analysis.api.getModule
 import org.jetbrains.kotlin.analysis.api.symbols.*
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbolOfType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
@@ -72,7 +72,7 @@ internal class KtFirSymbolProvider(
            is FirAnonymousFunctionSymbol -> firSymbolBuilder.functionLikeBuilder.buildAnonymousFunctionSymbol(firSymbol)
            else -> buildErrorWithAttachment("Unexpected ${firSymbol::class}") {
                withFirSymbolEntry("firSymbol", firSymbol)
-                withPsiEntry("function", psi)
+                withPsiEntry("function", psi, analysisSession::getModule)
            }
        }
    }
@@ -141,8 +141,9 @@ internal class KtFirSymbolProvider(
        return when (val firClassLike = resolveToFirSymbolOfType<FirClassLikeSymbol<*>>(firResolveSession)) {
            is FirTypeAliasSymbol -> firClassLike.fullyExpandedClass(analysisSession.useSiteSession)
                ?: buildErrorWithAttachment("${firClassLike.fir::class} should be expanded to the expected type alias") {
+                    val errorElement = this@resolveToFirClassLikeSymbol
                    withFirSymbolEntry("firClassLikeSymbol", firClassLike)
-                    withPsiEntry("ktClassOrObject", this@resolveToFirClassLikeSymbol)
+                    withPsiEntry("ktClassOrObject", errorElement, analysisSession::getModule)
                }
            is FirAnonymousObjectSymbol -> firClassLike
            is FirRegularClassSymbol -> firClassLike
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 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.
 */

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.api.fir.utils

 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSymbol
+import org.jetbrains.kotlin.analysis.api.getModule
 import org.jetbrains.kotlin.analysis.api.symbols.DebugSymbolRenderer
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withFirEntry
@@ -18,7 +19,10 @@ fun ExceptionAttachmentBuilder.withSymbolAttachment(name: String, symbol: KtSymb
    with(analysisSession) {
        withEntry(name, symbol) { DebugSymbolRenderer(renderExtra = true).render(it) }
    }
-    withPsiEntry("${name}Psi", symbol.psi)
+
+    val psi = symbol.psi
+    val psiModule = psi?.let(analysisSession::getModule)
+    withPsiEntry("${name}Psi", psi, psiModule)

    if (symbol is KtFirSymbol<*>) {
        val symbolFir = symbol.firSymbol.fir
@@ -5,19 +5,18 @@

 package org.jetbrains.kotlin.analysis.api.impl.base.test.cases.symbols

-import org.jetbrains.kotlin.analysis.api.renderer.declarations.KtDeclarationRenderer
 import org.jetbrains.kotlin.analysis.api.renderer.declarations.impl.KtDeclarationRendererForDebug
-import org.jetbrains.kotlin.analysis.api.symbols.*
-import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
-import org.jetbrains.kotlin.analysis.utils.printer.prettyPrint
-import org.jetbrains.kotlin.test.services.TestServices
-import org.jetbrains.kotlin.test.services.assertions
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.api.symbols.DebugSymbolRenderer
+import org.jetbrains.kotlin.analysis.api.symbols.KtDeclarationSymbol
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedTest
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
+import org.jetbrains.kotlin.analysis.utils.printer.prettyPrint
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.test.services.TestModuleStructure
+import org.jetbrains.kotlin.test.services.TestServices
+import org.jetbrains.kotlin.test.services.assertions

 abstract class AbstractSymbolRestoreFromDifferentModuleTest : AbstractAnalysisApiBasedTest() {
    private val defaultRenderer = KtDeclarationRendererForDebug.WITH_QUALIFIED_NAMES
@@ -33,6 +32,9 @@ abstract class AbstractSymbolRestoreFromDifferentModuleTest : AbstractAnalysisAp
                caretTag = "restoreAt"
            ).single().first

+        val declarationModule = ProjectStructureProvider.getModule(declaration, contextualModule = null)
+        val restoreAtModule = ProjectStructureProvider.getModule(restoreAt, contextualModule = null)
+
        val (debugRendered, prettyRendered, pointer) = analyseForTest(declaration) {
            val symbol = declaration.getSymbol()
            val pointer = symbol.createPointer()
@@ -46,19 +48,19 @@ abstract class AbstractSymbolRestoreFromDifferentModuleTest : AbstractAnalysisAp
        }

        val actualDebug = prettyPrint {
-            appendLine("Inital from ${declaration.getKtModule().moduleDescription}:")
+            appendLine("Inital from ${declarationModule.moduleDescription}:")
            appendLine(debugRendered)
            appendLine()
-            appendLine("Restored in ${restoreAt.getKtModule().moduleDescription}:")
+            appendLine("Restored in ${restoreAtModule.moduleDescription}:")
            appendLine(debugRenderedRestored ?: NOT_RESTORED)
        }
        testServices.assertions.assertEqualsToTestDataFileSibling(actualDebug)

        val actualPretty = prettyPrint {
-            appendLine("Inital from ${declaration.getKtModule().moduleDescription}:")
+            appendLine("Inital from ${declarationModule.moduleDescription}:")
            appendLine(prettyRendered)
            appendLine()
-            appendLine("Restored in ${restoreAt.getKtModule().moduleDescription}:")
+            appendLine("Restored in ${restoreAtModule.moduleDescription}:")
            appendLine(prettyRenderedRestored ?: NOT_RESTORED)
        }
        testServices.assertions.assertEqualsToTestDataFileSibling(actualPretty, extension = ".pretty.txt")
@@ -5,12 +5,14 @@

 package org.jetbrains.kotlin.analysis.api

+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.api.components.*
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeOwner
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile

@@ -178,3 +180,7 @@ public abstract class KtAnalysisSession(final override val token: KtLifetimeToke
        get() = typesCreatorImpl
    protected abstract val typesCreatorImpl: KtTypeCreator
 }
+
+public fun KtAnalysisSession.getModule(element: PsiElement): KtModule {
+    return ProjectStructureProvider.getModule(element, useSiteModule)
+}
@@ -8,11 +8,14 @@ package org.jetbrains.kotlin.analysis.utils.errors
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiInvalidElementAccessException
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.analysis.utils.printer.getElementTextInContext
 import org.jetbrains.kotlin.psi.KtElement

-public fun ExceptionAttachmentBuilder.withPsiEntry(name: String, psi: PsiElement?) {
+public fun ExceptionAttachmentBuilder.withPsiEntry(name: String, psi: PsiElement?, moduleFactory: (PsiElement) -> KtModule) {
+    return withPsiEntry(name, psi, psi?.let(moduleFactory))
+}
+
+public fun ExceptionAttachmentBuilder.withPsiEntry(name: String, psi: PsiElement?, module: KtModule?) {
    withEntry(name, psi) { psiElement ->
        when {
            !psiElement.isValid -> "INVALID PSI ${PsiInvalidElementAccessException.findOutInvalidationReason(psiElement)}"
@@ -21,7 +24,7 @@ public fun ExceptionAttachmentBuilder.withPsiEntry(name: String, psi: PsiElement
        }
    }
    if (psi != null) {
-        withKtModuleEntry("${name}KtModule", psi.getKtModule())
+        withKtModuleEntry("${name}KtModule", module)
    }
 }

@@ -11,6 +11,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.project.structure.KtCompilerPluginsProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtCompilerPluginsProvider.CompilerPluginType
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -20,7 +21,6 @@ import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 import org.jetbrains.kotlin.types.expressions.OperatorConventions.ASSIGN_METHOD
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule

 abstract class KtSimpleNameReference(expression: KtSimpleNameExpression) : KtSimpleReference<KtSimpleNameExpression>(expression) {
    // Extension point used by deprecated android extensions.
@@ -94,7 +94,7 @@ abstract class KtSimpleNameReference(expression: KtSimpleNameExpression) : KtSim
    abstract fun getImportAlias(): KtImportAlias?

    private fun isAssignmentResolved(project: Project, binaryExpression: KtBinaryExpression): Boolean {
-        val sourceModule = binaryExpression.getKtModule(element.project) as? KtSourceModule ?: return false
+        val sourceModule = ProjectStructureProvider.getModule(binaryExpression, contextualModule = null) as? KtSourceModule ?: return false
        val reference = binaryExpression.operationReference.reference ?: return false
        val pluginPresenceService = project.getService(KtCompilerPluginsProvider::class.java)
            ?: error("KtAssignResolutionPresenceService is not available as a service")
@@ -6,9 +6,11 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.api

 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
@@ -57,3 +59,7 @@ abstract class LLFirResolveSession {

    abstract fun getTowerContextProvider(ktFile: KtFile): FirTowerContextProvider
 }
+
+fun LLFirResolveSession.getModule(element: PsiElement): KtModule {
+    return ProjectStructureProvider.getModule(element, useSiteKtModule)
+}
@@ -26,7 +26,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
 import org.jetbrains.kotlin.analysis.utils.errors.withPsiEntry
 import org.jetbrains.kotlin.analysis.utils.printer.getElementTextInContext
@@ -105,8 +104,8 @@ object LowLevelFirApiFacadeForResolveOnAir {

        declaration.accept(expressionLocator)
        return expressionLocator.result ?: errorWithFirSpecificEntries("Resolved on-air element was not found in containing declaration") {
-            withPsiEntry("place", place)
-            withPsiEntry("elementToResolve", elementToResolve)
+            withPsiEntry("place", place, firResolveSession::getModule)
+            withPsiEntry("elementToResolve", elementToResolve, firResolveSession::getModule)
        }
    }

@@ -114,7 +113,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
        firResolveSession: LLFirResolveSession,
        ktFile: KtFile,
    ): FirTowerContextProvider {
-        val ktModule = ktFile.getKtModule(firResolveSession.project)
+        val ktModule = firResolveSession.getModule(ktFile)
        val session = firResolveSession.getSessionFor(ktModule) as LLFirResolvableModuleSession
        val moduleComponents = session.moduleComponents

@@ -153,7 +152,8 @@ object LowLevelFirApiFacadeForResolveOnAir {
        firResolveSession: LLFirResolvableResolveSession,
        file: KtFile,
    ): FirTowerDataContext {
-        val session = firResolveSession.getSessionFor(file.getKtModule(firResolveSession.project)) as LLFirResolvableModuleSession
+        val module = firResolveSession.getModule(file)
+        val session = firResolveSession.getSessionFor(module) as LLFirResolvableModuleSession
        val moduleComponents = session.moduleComponents

        val firFile = moduleComponents.firFileBuilder.buildRawFirFileWithCaching(file)
@@ -188,8 +188,8 @@ object LowLevelFirApiFacadeForResolveOnAir {

        val sameDeclarationInOriginalFile = PsiTreeUtil.findSameElementInCopy(dependencyNonLocalDeclaration, originalKtFile)
            ?: buildErrorWithAttachment("Cannot find original function matching") {
-                withPsiEntry("matchingPsi", dependencyNonLocalDeclaration)
-                withPsiEntry("originalFile", originalKtFile)
+                withPsiEntry("matchingPsi", dependencyNonLocalDeclaration, originalFirResolveSession::getModule)
+                withPsiEntry("originalFile", originalKtFile, originalFirResolveSession::getModule)
            }

        recordOriginalDeclaration(
@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder

 import org.jetbrains.kotlin.analysis.api.impl.barebone.annotations.ThreadSafe
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.utils.errors.checkWithAttachmentBuilder
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
@@ -18,16 +18,19 @@ import org.jetbrains.kotlin.psi.KtFile
 * Responsible for building [FirFile] by [KtFile]
 */
@ThreadSafe
-internal class LLFirFileBuilder(
-    val moduleComponents: LLFirModuleResolveComponents,
-) {
+internal class LLFirFileBuilder(val moduleComponents: LLFirModuleResolveComponents) {
+    private val projectStructureProvider by lazy { ProjectStructureProvider.getInstance(moduleComponents.session.project) }
+
    fun buildRawFirFileWithCaching(ktFile: KtFile): FirFile = moduleComponents.cache.fileCached(ktFile) {
-        checkWithAttachmentBuilder(ktFile.getKtModule() == moduleComponents.module, { "Modules are inconsistent" }) {
+        val contextualModule = moduleComponents.module
+        val actualFileModule = projectStructureProvider.getModule(ktFile, contextualModule)
+
+        checkWithAttachmentBuilder(actualFileModule == contextualModule, { "Modules are inconsistent" }) {
            withEntry("file", ktFile.name)
-            withEntry("file module", ktFile.getKtModule()) {
+            withEntry("file module", actualFileModule) {
                it.toString()
            }
-            withEntry("components module", moduleComponents.module) {
+            withEntry("components module", contextualModule) {
                it.toString()
            }
        }
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
 import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
 import org.jetbrains.kotlin.analysis.utils.errors.withPsiEntry
@@ -218,7 +219,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
            val typeParameters = mutableListOf<FirTypeParameterRef>()
            context.appendOuterTypeParameters(ignoreLastLevel = false, typeParameters)
            val containingClass = this.containingClass ?: buildErrorWithAttachment("Constructor outside of class") {
-                withPsiEntry("constructor", constructor)
+                withPsiEntry("constructor", constructor, baseSession.llFirModuleData.ktModule)
            }
            val selfType = classOrObject.toDelegatedSelfType(typeParameters, containingClass.symbol)
            val superTypeCallEntry = classOrObject.superTypeListEntries.lastIsInstanceOrNull<KtSuperTypeCallEntry>()
@@ -230,7 +231,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
            val params = extractContructorConversionParams(classOrObject, constructor)
            val delegatedTypeRef = (originalDeclaration as FirConstructor).delegatedConstructor?.constructedTypeRef
                ?: buildErrorWithAttachment("Secondary constructor without delegated call") {
-                    withPsiEntry("constructor", constructor)
+                    withPsiEntry("constructor", constructor, baseSession.llFirModuleData.ktModule)
                }
            return constructor.toFirConstructor(
                delegatedTypeRef,
@@ -264,7 +265,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(

        override fun visitEnumEntry(enumEntry: KtEnumEntry, data: Unit?): FirElement {
            val owner = containingClass ?: buildErrorWithAttachment("Enum entry outside of class") {
-                withPsiEntry("enumEntry", enumEntry)
+                withPsiEntry("enumEntry", enumEntry, baseSession.llFirModuleData.ktModule)
            }
            val classOrObject = owner.psi as KtClassOrObject
            val primaryConstructor = classOrObject.primaryConstructor
@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.caches.NullableCaffeineCa
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirSymbolProviderNameCacheBase
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.createDeclarationProvider
 import org.jetbrains.kotlin.fir.FirSession
@@ -105,7 +104,7 @@ internal class LLFirCombinedKotlinSymbolProvider private constructor(
        val callableId = CallableId(packageFqName, name)

        declarationProvider.getTopLevelCallableFiles(callableId)
-            .groupBy { it.getKtModule(project) }
+            .groupBy { getModule(it) }
            .forEach { (ktModule, ktFiles) ->
                // If `ktModule` cannot be found in the map, `ktFiles` cannot be processed by any of the available providers, because none
                // of them belong to the correct module. We can skip in that case, because iterating through all providers wouldn't lead to
@@ -18,7 +18,6 @@ import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.analysis.utils.classIdIfNonLocal

 class LLFirFirClassByPsiClassProvider(private val session: LLFirSession) : FirSessionComponent {
@@ -37,7 +36,7 @@ class LLFirFirClassByPsiClassProvider(private val session: LLFirSession) : FirSe
            }
        ) {
            withEntry("virtualFilePath", psiClass.containingFile.virtualFile?.path)
-            withPsiEntry("psiClass", psiClass)
+            withPsiEntry("psiClass", psiClass, session.ktModule)
        }

        if (psiClass.qualifiedName == null) {
@@ -71,7 +70,7 @@ class LLFirFirClassByPsiClassProvider(private val session: LLFirSession) : FirSe
        val provider = session.nullableJavaSymbolProvider ?: session.symbolProvider
        val symbol = provider.getClassLikeSymbolByClassId(classId)
            ?: buildErrorWithAttachment("No classifier found") {
-                withPsiEntry("psiClass", psiClass)
+                withPsiEntry("psiClass", psiClass, session.ktModule)
                withEntry("classId", classId) { it.asString() }
            }
        return symbol as FirRegularClassSymbol
@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSourcesSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinAnnotationsResolver
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.fir.FirElement
@@ -44,6 +45,7 @@ internal class LLFirIdePredicateBasedProvider(
    private val annotationsResolver: KotlinAnnotationsResolver,
    private val declarationProvider: KotlinDeclarationProvider,
 ) : FirPredicateBasedProvider() {
+    private val projectStructureProvider by lazy { ProjectStructureProvider.getInstance(session.project) }

    private val registeredPluginAnnotations: FirRegisteredPluginAnnotations
        get() = session.registeredPluginAnnotations
@@ -75,8 +77,9 @@ internal class LLFirIdePredicateBasedProvider(
            this !is KtProperty
        ) return null

-        val firResolveSession = this.getFirResolveSession()
-        return this.resolveToFirSymbol(firResolveSession).fir
+        val moduleForFile = projectStructureProvider.getModule(this, session.ktModule)
+        val sessionForFile = moduleForFile.getFirResolveSession(project)
+        return this.resolveToFirSymbol(sessionForFile).fir
    }

    override fun getOwnersOfDeclaration(declaration: FirDeclaration): List<FirBasedSymbol<*>>? {
@@ -40,7 +40,13 @@ abstract class LLFirSelectingCombinedSymbolProvider<PROVIDER : FirSymbolProvider
    /**
     * Cache [ProjectStructureProvider] to avoid service access when getting [KtModule]s.
     */
-    private val projectStructureProvider: ProjectStructureProvider = project.getService(ProjectStructureProvider::class.java)
+    private val projectStructureProvider: ProjectStructureProvider = ProjectStructureProvider.getInstance(project)
+
+    private val contextualModule = session.llFirModuleData.ktModule
+
+    protected fun getModule(element: PsiElement): KtModule {
+        return projectStructureProvider.getModule(element, contextualModule)
+    }

    /**
     * Selects the element with the highest module precedence in [candidates], returning the element and the provider to which resolution
@@ -60,7 +66,7 @@ abstract class LLFirSelectingCombinedSymbolProvider<PROVIDER : FirSymbolProvider

        for (candidate in candidates) {
            val element = getElement(candidate) ?: continue
-            val ktModule = projectStructureProvider.getKtModuleForKtElement(element)
+            val ktModule = getModule(element)

            // If `ktModule` cannot be found in the map, `candidate` cannot be processed by any of the available providers, because none of
            // them belong to the correct module. We can skip in that case, because iterating through all providers wouldn't lead to any
@@ -10,6 +10,7 @@ import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getModule
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.canBePartOfParentDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
@@ -23,7 +24,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecific
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
 import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
 import org.jetbrains.kotlin.analysis.utils.errors.withPsiEntry
@@ -93,15 +93,15 @@ internal abstract class LLFirResolvableResolveSession(
    }

    protected fun getModuleComponentsForElement(element: KtElement): LLFirModuleResolveComponents {
-        val ktModule = element.getKtModule()
-        return getResolvableSessionFor(ktModule).moduleComponents
+        val module = getModule(element)
+        return getResolvableSessionFor(module).moduleComponents
    }

    override fun resolveToFirSymbol(
        ktDeclaration: KtDeclaration,
        phase: FirResolvePhase
    ): FirBasedSymbol<*> {
-        val module = ktDeclaration.getKtModule()
+        val module = getModule(ktDeclaration)
        retryOnInvalidSession {
            return when (getModuleKind(module)) {
                ModuleKind.RESOLVABLE_MODULE -> findSourceFirSymbol(ktDeclaration, module).also { resolveFirToPhase(it.fir, phase) }
@@ -115,9 +115,9 @@ internal abstract class LLFirResolvableResolveSession(
            "This method will only work on compiled declarations, but this declaration is not compiled: ${ktDeclaration.getElementTextInContext()}"
        }

-        val ktModule = ktDeclaration.getKtModule(project)
-        val firSession = getSessionFor(ktModule)
-        val searcher = FirDeclarationForCompiledElementSearcher(firSession.symbolProvider)
+        val module = getModule(ktDeclaration)
+        val session = getSessionFor(module)
+        val searcher = FirDeclarationForCompiledElementSearcher(session.symbolProvider)
        val firDeclaration = searcher.findNonLocalDeclaration(ktDeclaration)
        return firDeclaration.symbol
    }
@@ -133,7 +133,7 @@ internal abstract class LLFirResolvableResolveSession(

        val nonLocalDeclaration = ktDeclaration.getNonLocalContainingOrThisDeclaration()
            ?: buildErrorWithAttachment("Declaration should have non-local container") {
-                withPsiEntry("ktDeclaration", ktDeclaration)
+                withPsiEntry("ktDeclaration", ktDeclaration, ::getModule)
                withEntry("module", module) { it.moduleDescription }
            }

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.util
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirModuleWithDependenciesSymbolProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirBuiltinsAndCloneableSession
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.utils.errors.ExceptionAttachmentBuilder
 import org.jetbrains.kotlin.analysis.utils.errors.withClassEntry
 import org.jetbrains.kotlin.fir.declarations.*
@@ -26,6 +26,11 @@ import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 * Allows to search for FIR declarations by compiled [KtDeclaration]s.
 */
 internal class FirDeclarationForCompiledElementSearcher(private val symbolProvider: FirSymbolProvider) {
+    private val projectStructureProvider by lazy {
+        val project = symbolProvider.session.llFirModuleData.ktModule.project
+        ProjectStructureProvider.getInstance(project)
+    }
+
    fun findNonLocalDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
        return when (ktDeclaration) {
            is KtEnumEntry -> findNonLocalEnumEntry(ktDeclaration)
@@ -81,7 +86,10 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
        if (classCandidate == null) {
            errorWithFirSpecificEntries("We should be able to find a symbol for $classId", psi = declaration) {
                withEntry("classId", classId) { it.asString() }
-                withEntry("ktModule", declaration.getKtModule()) { it.moduleDescription }
+
+                val contextualModule = symbolProvider.session.llFirModuleData.ktModule
+                val moduleForFile = projectStructureProvider.getModule(declaration, contextualModule)
+                withEntry("ktModule", moduleForFile) { it.moduleDescription }
            }
        }

@@ -7,7 +7,8 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.util

 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
-import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.utils.errors.ExceptionAttachmentBuilder
 import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
 import org.jetbrains.kotlin.analysis.utils.errors.withKtModuleEntry
@@ -36,7 +37,7 @@ fun ExceptionAttachmentBuilder.withFirEntry(name: String, fir: FirElement) {
    if (fir is FirElementWithResolveState) {
        withKtModuleEntry("${name}KtModule", fir.llFirModuleData.ktModule)
    }
-    withPsiEntry("${name}Psi", fir.psi)
+    withPsiEntry("${name}Psi", fir.psi, (fir as? FirElementWithResolveState)?.llFirSession?.ktModule)
 }


@@ -65,7 +66,7 @@ fun errorWithFirSpecificEntries(
        }

        if (psi != null) {
-            withPsiEntry("psi", psi)
+            withPsiEntry("psi", psi, ProjectStructureProvider.getModule(psi, contextualModule = null))
        }

        if (coneType != null) {
@@ -48,8 +48,8 @@ class SealedClassesInheritorsCaclulatorPreAnalysisHandler(
        for ((testModule, ktFiles) in ktFilesByModule) {
            if (ktFiles.isEmpty()) continue
            val project = testServices.compilerConfigurationProvider.getProject(testModule)
-            val ktModuleProvider = project.getService(ProjectStructureProvider::class.java)
-            val ktModule = ktFiles.map(ktModuleProvider::getKtModuleForKtElement).distinct().single()
+            val projectStructureProvider = project.getService(ProjectStructureProvider::class.java)
+            val ktModule = ktFiles.map { projectStructureProvider.getModule(it, contextualModule = null) }.distinct().single()

            val tmpFirResolveSession = LLFirResolveSessionService.getInstance(project).getFirResolveSessionNoCaching(ktModule)
            val firFiles = ktFiles.map { it.getOrBuildFirFile(tmpFirResolveSession) }
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostic

 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.getModule
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.BeforeElementDiagnosticCollectionHandler
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.beforeElementDiagnosticCollectionHandler
@@ -22,7 +23,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLeve
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirOutOfContentRootTestConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.useFirSessionConfigurator
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
@@ -45,7 +45,8 @@ abstract class AbstractFirContextCollectionTest : AbstractLowLevelApiSingleFileT
        resolveWithClearCaches(ktFile) { firResolveSession ->
            check(firResolveSession.isSourceSession)

-            val session = firResolveSession.getSessionFor(ktFile.getKtModule()) as LLFirResolvableModuleSession
+            val module = firResolveSession.getModule(ktFile)
+            val session = firResolveSession.getSessionFor(module) as LLFirResolvableModuleSession
            val handler = session.beforeElementDiagnosticCollectionHandler as BeforeElementTestDiagnosticCollectionHandler

            val fileStructureCache = session.moduleComponents.fileStructureCache
@@ -17,7 +17,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableM
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirOutOfContentRootTestConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import org.jetbrains.kotlin.test.KotlinTestUtils
@@ -102,9 +102,10 @@ abstract class AbstractFileStructureTest : AbstractLowLevelApiSingleFileTest() {
    }

    private fun KtFile.getFileStructure(): FileStructure {
-        val moduleFirResolveSession = getFirResolveSession()
+        val module = ProjectStructureProvider.getModule(this, contextualModule = null)
+        val moduleFirResolveSession = module.getFirResolveSession(project)
        check(moduleFirResolveSession.isSourceSession)
-        val session = moduleFirResolveSession.getSessionFor(getKtModule()) as LLFirResolvableModuleSession
+        val session = moduleFirResolveSession.getSessionFor(module) as LLFirResolvableModuleSession
        return session.moduleComponents.fileStructureCache.getFileStructure(this)
    }

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirNotUnderContentRootResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirSourceResolveSession
-import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.test.framework.services.environmentManager
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
@@ -40,7 +40,7 @@ internal fun FirDeclaration.name(): String = symbol.name()

 internal inline fun <R> resolveWithClearCaches(context: KtElement, action: (LLFirResolveSession) -> R): R {
    val project = context.project
-    val module = context.getKtModule(project)
+    val module = ProjectStructureProvider.getModule(context, contextualModule = null)
    val resolveSession = LLFirResolveSessionService.getInstance(project).getFirResolveSessionNoCaching(module)
    return action(resolveSession)
 }
@@ -35,6 +35,8 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtScript

 class SymbolKotlinAsJavaSupport(project: Project) : KotlinAsJavaSupportBase<KtModule>(project) {
+    private val projectStructureProvider by lazy { ProjectStructureProvider.getInstance(project) }
+
    override fun findClassOrObjectDeclarationsInPackage(
        packageFqName: FqName,
        searchScope: GlobalSearchScope
@@ -92,16 +94,21 @@ class SymbolKotlinAsJavaSupport(project: Project) : KotlinAsJavaSupportBase<KtMo
            .map { fqn.child(it) }

    override fun createInstanceOfLightScript(script: KtScript): KtLightClass {
-        return SymbolLightClassForScript(script, script.getKtModule(project))
+        val module = ProjectStructureProvider.getModule(script, contextualModule = null)
+        return SymbolLightClassForScript(script, module)
    }

-    override fun KtFile.findModule(): KtModule = getKtModule(project)
+    override fun KtFile.findModule(): KtModule {
+        return projectStructureProvider.getModule(this, contextualModule = null)
+    }

-    override fun declarationLocation(file: KtFile): DeclarationLocation? = when (file.getKtModule(project)) {
-        is KtSourceModule -> DeclarationLocation.ProjectSources
-        is KtLibraryModule -> DeclarationLocation.LibraryClasses
-        is KtLibrarySourceModule -> DeclarationLocation.LibrarySources
-        else -> null
+    override fun declarationLocation(file: KtFile): DeclarationLocation? {
+        return when (projectStructureProvider.getModule(file, contextualModule = null)) {
+            is KtSourceModule -> DeclarationLocation.ProjectSources
+            is KtLibraryModule -> DeclarationLocation.LibraryClasses
+            is KtLibrarySourceModule -> DeclarationLocation.LibrarySources
+            else -> null
+        }
    }

    override fun createInstanceOfDecompiledLightClass(classOrObject: KtClassOrObject): KtLightClass? {
@@ -109,7 +116,8 @@ class SymbolKotlinAsJavaSupport(project: Project) : KotlinAsJavaSupportBase<KtMo
    }

    override fun createInstanceOfLightClass(classOrObject: KtClassOrObject): KtLightClass? {
-        return createSymbolLightClassNoCache(classOrObject, classOrObject.getKtModule(project))
+        val module = projectStructureProvider.getModule(classOrObject, contextualModule = null)
+        return createSymbolLightClassNoCache(classOrObject, module)
    }

    override fun createInstanceOfDecompiledLightFacade(facadeFqName: FqName, files: List<KtFile>): KtLightClassForFacade? {
@@ -128,10 +136,10 @@ class SymbolKotlinAsJavaSupport(project: Project) : KotlinAsJavaSupportBase<KtMo
        return project.createAllLibrariesModificationTracker()
    }

-    override fun createInstanceOfLightFacade(
-        facadeFqName: FqName,
-        files: List<KtFile>,
-    ): KtLightClassForFacade = SymbolLightClassForFacade(facadeFqName, files, files.first().getKtModule(project))
+    override fun createInstanceOfLightFacade(facadeFqName: FqName, files: List<KtFile>): KtLightClassForFacade {
+        val module = projectStructureProvider.getModule(files.first(), contextualModule = null)
+        return SymbolLightClassForFacade(facadeFqName, files, module)
+    }

    override val KtModule.contentSearchScope: GlobalSearchScope get() = this.contentScope

@@ -168,7 +176,10 @@ class SymbolKotlinAsJavaSupport(project: Project) : KotlinAsJavaSupportBase<KtMo

    override fun getFakeLightClass(classOrObject: KtClassOrObject): KtFakeLightClass = SymbolBasedFakeLightClass(classOrObject)

-    private fun KtElement.isFromSourceOrLibraryBinary(project: Project): Boolean = getKtModule(project).isFromSourceOrLibraryBinary()
+    private fun KtElement.isFromSourceOrLibraryBinary(): Boolean {
+        val module = projectStructureProvider.getModule(this, contextualModule = null)
+        return module.isFromSourceOrLibraryBinary()
+    }

    private fun KtModule.isFromSourceOrLibraryBinary() = when (this) {
        is KtSourceModule -> true
@@ -11,6 +11,7 @@ import com.intellij.psi.PsiModifier
 import com.intellij.psi.PsiReferenceList
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.annotations.hasAnnotation
+import org.jetbrains.kotlin.analysis.api.getModule
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithTypeParameters
@@ -21,7 +22,6 @@ import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.api.types.KtTypeMappingMode
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
-import org.jetbrains.kotlin.analysis.project.structure.getKtModuleOfTypeSafe
 import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
 import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
 import org.jetbrains.kotlin.asJava.builder.LightMemberOriginForDeclaration
@@ -537,7 +537,7 @@ internal fun KtSymbolWithMembers.createInnerClasses(
    }

    val jvmDefaultMode = classOrObject
-        ?.getKtModuleOfTypeSafe<KtSourceModule>()
+        ?.let { getModule(it) as? KtSourceModule }
        ?.languageVersionSettings
        ?.getFlag(JvmAnalysisFlags.jvmDefaultMode)
        ?: JvmDefaultMode.DEFAULT