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

Ensure the contextual module is passed wherever possible.

^KT-57559 Fixed

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/CliFe10AnalysisFacade.kt

@@ -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)
     }
 }

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10AnalysisSession.kt

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

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10CliAnalysisSessionProvider.kt

@@ -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
 

analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/components/KtFe10SymbolContainingDeclarationProvider.kt

@@ -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.*
-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.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) }
+        val descriptor = symbol.getDescriptor()
-            ?: symbol.psi
 
-        return psiForModule?.getKtModule(analysisSession.analysisContext.resolveSession.project)
+        val symbolPsi = descriptor?.let(DescriptorToSourceUtils::getContainingFile) ?: symbol.psi
-            ?: symbol.getDescriptor()?.getFakeContainingKtModule()
+        if (symbolPsi != null) {
-            ?: (symbol as? KtBackingFieldSymbol)?.owningProperty?.let { getContainingModule(it) }
+            return analysisSession.getModule(symbolPsi)
-            ?: TODO(symbol::class.java.name)
+        }
+
+        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? {
+    private fun getFakeContainingKtModule(descriptor: DescriptorWithContainerSource): KtModule {
-        return when (this) {
+        val libraryPath = Paths.get((descriptor.containerSource as JvmPackagePartSource).knownJvmBinaryClass?.containingLibrary!!)
-            is DescriptorWithContainerSource -> {
+        return object : KtLibraryModule {
-                val libraryPath = Paths.get((containerSource as JvmPackagePartSource).knownJvmBinaryClass?.containingLibrary!!)
+            override val libraryName: String = libraryPath.fileName.toString().substringBeforeLast(".")
-                object : KtLibraryModule {
+            override val librarySources: KtLibrarySourceModule? = null
-                    override val libraryName: String = libraryPath.fileName.toString().substringBeforeLast(".")
+            override fun getBinaryRoots(): Collection<Path> = listOf(libraryPath)
-                    override val librarySources: KtLibrarySourceModule? = null
+            override val directRegularDependencies: List<KtModule> = emptyList()
-                    override fun getBinaryRoots(): Collection<Path> = listOf(libraryPath)
+            override val directDependsOnDependencies: List<KtModule> = emptyList()
-                    override val directRegularDependencies: List<KtModule> = emptyList()
+            override val transitiveDependsOnDependencies: List<KtModule> = emptyList()
-                    override val directDependsOnDependencies: List<KtModule> = emptyList()
+            override val directFriendDependencies: List<KtModule> = emptyList()
-                    override val transitiveDependsOnDependencies: List<KtModule> = emptyList()
+            override val contentScope: GlobalSearchScope = ProjectScope.getLibrariesScope(project)
-                    override val directFriendDependencies: List<KtModule> = emptyList()
+            override val platform: TargetPlatform
-                    override val contentScope: GlobalSearchScope = ProjectScope.getLibrariesScope(project)
+                get() = descriptor.platform!!
-                    override val platform: TargetPlatform
+            override val analyzerServices: PlatformDependentAnalyzerServices
-                        get() = this@getFakeContainingKtModule.platform!!
+                get() = JvmPlatformAnalyzerServices
-                    override val analyzerServices: PlatformDependentAnalyzerServices
+            override val project: Project
-                        get() = JvmPlatformAnalyzerServices
+                get() = analysisSession.analysisContext.resolveSession.project
-                    override val project: Project
-                        get() = analysisSession.analysisContext.resolveSession.project
 
-                }
-            }
-
-            else -> null
         }
     }
 }

analysis/analysis-api-fe10/tests/org/jetbrains/kotlin/analysis/api/fe10/test/configurator/AnalysisApiFe10TestServiceRegistrar.kt

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

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSessionProvider.kt

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

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt

@@ -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) }
             }
         }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirPsiJavaClassSymbol.kt

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

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbolProvider.kt

@@ -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

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/utils/errorUtils.kt

@@ -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

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/symbols/AbstractSymbolRestoreFromDifferentModuleTest.kt

@@ -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.DebugSymbolRenderer
-import org.jetbrains.kotlin.analysis.api.symbols.pointers.KtSymbolPointer
+import org.jetbrains.kotlin.analysis.api.symbols.KtDeclarationSymbol
-import org.jetbrains.kotlin.analysis.utils.printer.prettyPrint
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
-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.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")

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/KtAnalysisSession.kt

@@ -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)
+}

analysis/analysis-internal-utils/src/org/jetbrains/kotlin/analysis/utils/errors/exceptionWithAttachmentBuilderHelpers.kt

@@ -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)
     }
 }
 

analysis/kt-references/src/org/jetbrains/kotlin/idea/references/KtSimpleNameReference.kt

@@ -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")

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

@@ -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)
+}

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

@@ -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("place", place, firResolveSession::getModule)
-            withPsiEntry("elementToResolve", elementToResolve)
+            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("matchingPsi", dependencyNonLocalDeclaration, originalFirResolveSession::getModule)
-                withPsiEntry("originalFile", originalKtFile)
+                withPsiEntry("originalFile", originalKtFile, originalFirResolveSession::getModule)
             }
 
         recordOriginalDeclaration(

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirFileBuilder.kt

@@ -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(
+internal class LLFirFileBuilder(val moduleComponents: LLFirModuleResolveComponents) {
-    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()
             }
         }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt

@@ -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

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

@@ -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

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

@@ -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

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

@@ -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()
+        val moduleForFile = projectStructureProvider.getModule(this, session.ktModule)
-        return this.resolveToFirSymbol(firResolveSession).fir
+        val sessionForFile = moduleForFile.getFirResolveSession(project)
+        return this.resolveToFirSymbol(sessionForFile).fir
     }
 
     override fun getOwnersOfDeclaration(declaration: FirDeclaration): List<FirBasedSymbol<*>>? {

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

@@ -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

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

@@ -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()
+        val module = getModule(element)
-        return getResolvableSessionFor(ktModule).moduleComponents
+        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 module = getModule(ktDeclaration)
-        val firSession = getSessionFor(ktModule)
+        val session = getSessionFor(module)
-        val searcher = FirDeclarationForCompiledElementSearcher(firSession.symbolProvider)
+        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 }
             }
 

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

@@ -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 }
             }
         }
 

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

@@ -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) {

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/compiler/based/SealedClassesInheritorsCaclulatorPreAnalysisHandler.kt

@@ -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 projectStructureProvider = project.getService(ProjectStructureProvider::class.java)
-            val ktModule = ktFiles.map(ktModuleProvider::getKtModuleForKtElement).distinct().single()
+            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) }

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

@@ -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

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

@@ -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)
     }
 

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

@@ -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)
 }

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/SymbolKotlinAsJavaSupport.kt

@@ -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)) {
+    override fun declarationLocation(file: KtFile): DeclarationLocation? {
-        is KtSourceModule -> DeclarationLocation.ProjectSources
+        return when (projectStructureProvider.getModule(file, contextualModule = null)) {
-        is KtLibraryModule -> DeclarationLocation.LibraryClasses
+            is KtSourceModule -> DeclarationLocation.ProjectSources
-        is KtLibrarySourceModule -> DeclarationLocation.LibrarySources
+            is KtLibraryModule -> DeclarationLocation.LibraryClasses
-        else -> null
+            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(
+    override fun createInstanceOfLightFacade(facadeFqName: FqName, files: List<KtFile>): KtLightClassForFacade {
-        facadeFqName: FqName,
+        val module = projectStructureProvider.getModule(files.first(), contextualModule = null)
-        files: List<KtFile>,
+        return SymbolLightClassForFacade(facadeFqName, files, module)
-    ): KtLightClassForFacade = SymbolLightClassForFacade(facadeFqName, files, files.first().getKtModule(project))
+    }
 
     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

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/symbolLightClassUtils.kt

@@ -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