[LL API] Move out traits of 'LLFirResolveSession' to separate classes

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

@@ -14,10 +14,12 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirElemen
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.KtToFirMapping
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLDiagnosticProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLScopeSessionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.TowerProviderForElementForState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirScopeSessionProvider
-import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
 import org.jetbrains.kotlin.analysis.utils.caches.*
 import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.FirElement
@@ -30,7 +32,6 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
 
 /**
  * An [LLFirResolveSession] which *depends* on an [originalFirResolveSession], but can provide its own FIR elements and symbols selectively.
@@ -44,28 +45,14 @@ internal class LLFirResolveSessionDepended(
     val originalFirResolveSession: LLFirResolvableResolveSession,
     val towerProviderBuiltUponElement: FirTowerContextProvider,
     private val ktToFirMapping: KtToFirMapping?,
-) : LLFirResolveSession() {
-    override val project: Project get() = originalFirResolveSession.project
-    override val useSiteKtModule: KtModule get() = originalFirResolveSession.useSiteKtModule
-    override val useSiteFirSession get() = originalFirResolveSession.useSiteFirSession
-
-    private val scopeSessionProviderCache = SoftCachedMap.create<FirSession, LLFirScopeSessionProvider>(
-        project,
-        SoftCachedMap.Kind.SOFT_KEYS_SOFT_VALUES,
-        listOf(
-            PsiModificationTracker.MODIFICATION_COUNT,
-            ProjectRootModificationTracker.getInstance(project)
-        )
-    )
-
-    override fun getScopeSessionFor(firSession: FirSession): ScopeSession {
-        return scopeSessionProviderCache
-            .getOrPut(firSession) { LLFirScopeSessionProvider.create(project, invalidationTrackers = emptyList()) }
-            .getScopeSession()
-    }
-
-    override fun getSessionFor(module: KtModule): LLFirSession =
-        originalFirResolveSession.getSessionFor(module)
+) : LLFirResolveSession(
+    originalFirResolveSession.moduleProvider,
+    originalFirResolveSession.moduleKindProvider,
+    originalFirResolveSession.sessionProvider,
+    LLDependedScopeSessionProvider(originalFirResolveSession.project)
+) {
+    override val diagnosticProvider: LLDiagnosticProvider
+        get() = LLDependedDiagnosticProvider
 
     override fun getOrBuildFirFor(element: KtElement): FirElement? {
         val psi = FirElementBuilder.getPsiAsFirElementSource(element) ?: return null
@@ -80,12 +67,6 @@ internal class LLFirResolveSessionDepended(
         originalFirResolveSession.resolveFirToPhase(declaration, toPhase)
     }
 
-    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> =
-        TODO("Diagnostics are not implemented for depended state")
-
-    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> =
-        TODO("Diagnostics are not implemented for depended state")
-
     override fun resolveToFirSymbol(ktDeclaration: KtDeclaration, phase: FirResolvePhase): FirBasedSymbol<*> {
         val declarationToResolve = ktDeclaration.originalDeclaration ?: ktDeclaration
         ktToFirMapping?.getElement(declarationToResolve)?.let { it as? FirDeclaration }?.symbol?.let { return it }
@@ -96,3 +77,30 @@ internal class LLFirResolveSessionDepended(
     override fun getTowerContextProvider(ktFile: KtFile): FirTowerContextProvider =
         TowerProviderForElementForState(this)
 }
+
+private object LLDependedDiagnosticProvider : LLDiagnosticProvider {
+    override fun collectDiagnostics(file: KtFile, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
+        throw NotImplementedError("Diagnostics are not implemented for depended state")
+    }
+
+    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
+        throw NotImplementedError("Diagnostics are not implemented for depended state")
+    }
+}
+
+private class LLDependedScopeSessionProvider(private val project: Project) : LLScopeSessionProvider {
+    private val scopeSessionProviderCache = SoftCachedMap.create<FirSession, LLFirScopeSessionProvider>(
+        project,
+        SoftCachedMap.Kind.SOFT_KEYS_SOFT_VALUES,
+        listOf(
+            PsiModificationTracker.MODIFICATION_COUNT,
+            ProjectRootModificationTracker.getInstance(project)
+        )
+    )
+
+    override fun getScopeSession(session: LLFirSession): ScopeSession {
+        return scopeSessionProviderCache
+            .getOrPut(session) { LLFirScopeSessionProvider.create(project, invalidationTrackers = emptyList()) }
+            .getScopeSession()
+    }
+}

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

@@ -9,8 +9,14 @@ 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.low.level.api.fir.state.LLDiagnosticProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLModuleKindProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLModuleProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLScopeSessionProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLSessionProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
+import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
 import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
@@ -26,16 +32,31 @@ import org.jetbrains.kotlin.psi.KtFile
 /**
  * An entry point for a FIR Low Level API resolution. Represents a project view from a use-site [KtModule].
  */
-abstract class LLFirResolveSession {
-    abstract val project: Project
+abstract class LLFirResolveSession(
+    val moduleProvider: LLModuleProvider,
+    val moduleKindProvider: LLModuleKindProvider,
+    val sessionProvider: LLSessionProvider,
+    val scopeSessionProvider: LLScopeSessionProvider
+) {
+    abstract val diagnosticProvider: LLDiagnosticProvider
 
-    abstract val useSiteFirSession: LLFirSession
+    val useSiteKtModule: KtModule
+        get() = moduleProvider.useSiteModule
 
-    abstract val useSiteKtModule: KtModule
+    val project: Project
+        get() = useSiteKtModule.project
 
-    abstract fun getSessionFor(module: KtModule): LLFirSession
+    val useSiteFirSession: LLFirSession
+        get() = sessionProvider.useSiteSession
 
-    abstract fun getScopeSessionFor(firSession: FirSession): ScopeSession
+    fun getSessionFor(module: KtModule): LLFirSession {
+        return sessionProvider.getSession(module)
+    }
+
+    fun getScopeSessionFor(firSession: FirSession): ScopeSession {
+        requireIsInstance<LLFirSession>(firSession)
+        return scopeSessionProvider.getScopeSession(firSession)
+    }
 
     /**
      * Build fully resolved FIR node for a requested element.
@@ -51,9 +72,13 @@ abstract class LLFirResolveSession {
      */
     internal abstract fun getOrBuildFirFile(ktFile: KtFile): FirFile
 
-    internal abstract fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic>
+    internal fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
+        return diagnosticProvider.getDiagnostics(element, filter)
+    }
 
-    internal abstract fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic>
+    internal fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> {
+        return diagnosticProvider.collectDiagnostics(ktFile, filter)
+    }
 
     abstract fun resolveToFirSymbol(ktDeclaration: KtDeclaration, phase: FirResolvePhase): FirBasedSymbol<*>
 

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

@@ -0,0 +1,51 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.state
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
+import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+
+interface LLDiagnosticProvider {
+    /**
+     * Returns all compiler diagnostics for the [file], matching the [filter].
+     */
+    fun collectDiagnostics(file: KtFile, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic>
+
+    /**
+     * Returns all compiler diagnostics for the specific [element], matching the [filter].
+     * This function is not recursive; diagnostics for nested elements are not returned.
+     */
+    fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic>
+}
+
+internal object LLEmptyDiagnosticProvider : LLDiagnosticProvider {
+    override fun collectDiagnostics(file: KtFile, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
+        return emptyList()
+    }
+
+    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
+        return emptyList()
+    }
+}
+
+internal class LLSourceDiagnosticProvider(
+    private val moduleProvider: LLModuleProvider,
+    private val sessionProvider: LLSessionProvider
+) : LLDiagnosticProvider {
+    override fun collectDiagnostics(file: KtFile, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
+        val module = moduleProvider.getModule(file)
+        val moduleComponents = sessionProvider.getResolvableSession(module).moduleComponents
+        return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(file, filter)
+    }
+
+    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
+        val module = moduleProvider.getModule(element)
+        val moduleComponents = sessionProvider.getResolvableSession(module).moduleComponents
+        return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
+    }
+}

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

@@ -5,27 +5,25 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirLibraryOrLibrarySourceResolvableModuleSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
 
 internal class LLFirLibraryOrLibrarySourceResolvableResolveSession(
     useSiteKtModule: KtModule,
     useSiteSessionFactory: (KtModule) -> LLFirSession
-) : LLFirResolvableResolveSession(useSiteKtModule, useSiteSessionFactory) {
-    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> =
-        emptyList()
-
-    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> =
-        emptyList()
-
-    override fun getModuleKind(module: KtModule): ModuleKind {
-        LLFirLibraryOrLibrarySourceResolvableModuleSession.checkIsValidKtModule(module)
-        return if (module == useSiteKtModule) ModuleKind.RESOLVABLE_MODULE else ModuleKind.BINARY_MODULE
-    }
+) : LLFirResolvableResolveSession(
+    useSiteKtModule = useSiteKtModule,
+    moduleKindProvider = LLLibraryModuleKindProvider(useSiteKtModule),
+    useSiteSessionFactory = useSiteSessionFactory
+) {
+    override val diagnosticProvider: LLDiagnosticProvider
+        get() = LLEmptyDiagnosticProvider
 }
 
+private class LLLibraryModuleKindProvider(private val useSiteModule: KtModule) : LLModuleKindProvider {
+    override fun getKind(module: KtModule): KtModuleKind {
+        LLFirLibraryOrLibrarySourceResolvableModuleSession.checkIsValidKtModule(module)
+        return if (module == useSiteModule) KtModuleKind.RESOLVABLE_MODULE else KtModuleKind.BINARY_MODULE
+    }
+}

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

@@ -5,28 +5,26 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
 
 internal class LLFirNotUnderContentRootResolveSession(
     useSiteKtModule: KtModule,
     useSiteSessionFactory: (KtModule) -> LLFirSession
-) : LLFirResolvableResolveSession(useSiteKtModule, useSiteSessionFactory) {
-    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> =
-        emptyList()
-
-    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> =
-        emptyList()
-
-    override fun getModuleKind(module: KtModule): ModuleKind {
-        return when (module) {
-            useSiteKtModule -> ModuleKind.RESOLVABLE_MODULE
-            else -> ModuleKind.BINARY_MODULE
-        }
-    }
+) : LLFirResolvableResolveSession(
+    useSiteKtModule = useSiteKtModule,
+    moduleKindProvider = LLNotUnderContentRootModuleKindProvider(useSiteKtModule),
+    useSiteSessionFactory = useSiteSessionFactory
+) {
+    override val diagnosticProvider: LLDiagnosticProvider
+        get() = LLEmptyDiagnosticProvider
 }
 
+private class LLNotUnderContentRootModuleKindProvider(private val useSiteModule: KtModule) : LLModuleKindProvider {
+    override fun getKind(module: KtModule): KtModuleKind {
+        return when (module) {
+            useSiteModule -> KtModuleKind.RESOLVABLE_MODULE
+            else -> KtModuleKind.BINARY_MODULE
+        }
+    }
+}

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

@@ -5,33 +5,22 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
-import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.ModificationTracker
-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.getNonLocalContainingDeclaration
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirDeclarationForCompiledElementSearcher
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
-import org.jetbrains.kotlin.analysis.utils.errors.requireIsInstance
 import org.jetbrains.kotlin.analysis.utils.errors.withPsiEntry
-import org.jetbrains.kotlin.psi.psiUtil.getElementTextWithContext
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousObjectExpression
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -40,44 +29,19 @@ import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.psiUtil.getElementTextWithContext
+import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 
 internal abstract class LLFirResolvableResolveSession(
-    final override val useSiteKtModule: KtModule,
-    private val useSiteSessionFactory: (KtModule) -> LLFirSession
-) : LLFirResolveSession() {
-    final override val project: Project
-        get() = useSiteKtModule.project
-
-    private val useSiteFirSessionCached = CachedValuesManager.getManager(project).createCachedValue {
-        val session = useSiteSessionFactory(useSiteKtModule)
-        CachedValueProvider.Result.create(session, session.createValidityTracker())
-    }
-
-    final override val useSiteFirSession: LLFirSession
-        get() = useSiteFirSessionCached.value
-
-    override fun getSessionFor(module: KtModule): LLFirSession {
-        return getSession(module, preferBinary = true)
-    }
-
-    private fun getResolvableSessionFor(module: KtModule): LLFirResolvableModuleSession {
-        return getSession(module, preferBinary = false) as LLFirResolvableModuleSession
-    }
-
-    private fun getSession(module: KtModule, preferBinary: Boolean): LLFirSession {
-        if (module == useSiteFirSession.ktModule) {
-            return useSiteFirSession
-        }
-
-        val cache = LLFirSessionCache.getInstance(module.project)
-        return cache.getSession(module, preferBinary)
-    }
-
-    override fun getScopeSessionFor(firSession: FirSession): ScopeSession {
-        requireIsInstance<LLFirSession>(firSession)
-        return firSession.getScopeSession()
-    }
-
+    useSiteKtModule: KtModule,
+    moduleKindProvider: LLModuleKindProvider,
+    useSiteSessionFactory: (KtModule) -> LLFirSession,
+) : LLFirResolveSession(
+    moduleProvider = LLModuleProvider(useSiteKtModule),
+    moduleKindProvider = moduleKindProvider,
+    sessionProvider = LLSessionProvider(useSiteKtModule, useSiteSessionFactory),
+    scopeSessionProvider = LLDefaultScopeSessionProvider
+) {
     override fun getOrBuildFirFor(element: KtElement): FirElement? {
         val moduleComponents = getModuleComponentsForElement(element)
         return moduleComponents.elementsBuilder.getOrBuildFirFor(element, this)
@@ -90,7 +54,7 @@ internal abstract class LLFirResolvableResolveSession(
 
     protected fun getModuleComponentsForElement(element: KtElement): LLFirModuleResolveComponents {
         val module = getModule(element)
-        return getResolvableSessionFor(module).moduleComponents
+        return sessionProvider.getResolvableSession(module).moduleComponents
     }
 
     override fun resolveToFirSymbol(
@@ -100,8 +64,8 @@ internal abstract class LLFirResolvableResolveSession(
         val containingKtFile = ktDeclaration.containingKtFile
         val module = getModule(containingKtFile.originalKtFile ?: containingKtFile)
         return when (getModuleKind(module)) {
-            ModuleKind.RESOLVABLE_MODULE -> findSourceFirSymbol(ktDeclaration, module).also { resolveFirToPhase(it.fir, phase) }
-            ModuleKind.BINARY_MODULE -> findFirCompiledSymbol(ktDeclaration, module)
+            KtModuleKind.RESOLVABLE_MODULE -> findSourceFirSymbol(ktDeclaration, module).also { resolveFirToPhase(it.fir, phase) }
+            KtModuleKind.BINARY_MODULE -> findFirCompiledSymbol(ktDeclaration, module)
         }
     }
 
@@ -121,7 +85,7 @@ internal abstract class LLFirResolvableResolveSession(
     }
 
     private fun findSourceFirDeclarationByDeclaration(ktDeclaration: KtDeclaration, module: KtModule): FirBasedSymbol<*> {
-        require(getModuleKind(module) == ModuleKind.RESOLVABLE_MODULE) {
+        require(getModuleKind(module) == KtModuleKind.RESOLVABLE_MODULE) {
             "Declaration should be resolvable module, instead it had ${module::class}"
         }
 
@@ -132,7 +96,7 @@ internal abstract class LLFirResolvableResolveSession(
             }
 
         if (ktDeclaration == nonLocalDeclaration) {
-            val session = getResolvableSessionFor(module)
+            val session = sessionProvider.getResolvableSession(module)
             return nonLocalDeclaration.findSourceNonLocalFirDeclaration(
                 firFileBuilder = session.moduleComponents.firFileBuilder,
                 provider = session.firProvider,
@@ -142,7 +106,9 @@ internal abstract class LLFirResolvableResolveSession(
         return findDeclarationInSourceViaResolve(ktDeclaration)
     }
 
-    protected abstract fun getModuleKind(module: KtModule): ModuleKind
+    protected fun getModuleKind(module: KtModule): KtModuleKind {
+        return moduleKindProvider.getKind(module)
+    }
 
     private fun findDeclarationInSourceViaResolve(ktDeclaration: KtExpression): FirBasedSymbol<*> {
         val firDeclaration = when (val fir = getOrBuildFirFor(ktDeclaration)) {
@@ -165,9 +131,4 @@ internal abstract class LLFirResolvableResolveSession(
     override fun getTowerContextProvider(ktFile: KtFile): FirTowerContextProvider {
         return TowerProviderForElementForState(this)
     }
-
-    protected enum class ModuleKind {
-        RESOLVABLE_MODULE,
-        BINARY_MODULE
-    }
 }

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

@@ -5,32 +5,26 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.analysis.utils.errors.unexpectedElementError
-import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
 
 internal class LLFirScriptResolveSession(
     useSiteKtModule: KtModule,
     useSiteSessionFactory: (KtModule) -> LLFirSession,
-) : LLFirResolvableResolveSession(useSiteKtModule, useSiteSessionFactory) {
-    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
-        val moduleComponents = getModuleComponentsForElement(element)
-        return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
-    }
+) : LLFirResolvableResolveSession(
+    useSiteKtModule = useSiteKtModule,
+    moduleKindProvider = LLScriptModuleKindProvider(useSiteKtModule),
+    useSiteSessionFactory = useSiteSessionFactory
+) {
+    override val diagnosticProvider = LLSourceDiagnosticProvider(moduleProvider, sessionProvider)
+}
 
-    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> {
-        val moduleComponents = getModuleComponentsForElement(ktFile)
-        return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
-    }
-
-    override fun getModuleKind(module: KtModule): ModuleKind {
+private class LLScriptModuleKindProvider(private val useSiteModule: KtModule) : LLModuleKindProvider {
+    override fun getKind(module: KtModule): KtModuleKind {
         return when (module) {
-            useSiteKtModule, is KtSourceModule -> ModuleKind.RESOLVABLE_MODULE
-            is KtBuiltinsModule, is KtLibraryModule -> ModuleKind.BINARY_MODULE
+            useSiteModule, is KtSourceModule -> KtModuleKind.RESOLVABLE_MODULE
+            is KtBuiltinsModule, is KtLibraryModule -> KtModuleKind.BINARY_MODULE
             else -> unexpectedElementError("module", module)
         }
     }

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

@@ -5,38 +5,31 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.KtBuiltinsModule
 import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.utils.errors.unexpectedElementError
-import org.jetbrains.kotlin.diagnostics.KtPsiDiagnostic
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
 
 internal class LLFirSourceResolveSession(
     useSiteKtModule: KtModule,
     useSiteSessionFactory: (KtModule) -> LLFirSession,
-) : LLFirResolvableResolveSession(useSiteKtModule, useSiteSessionFactory) {
-    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
-        val moduleComponents = getModuleComponentsForElement(element)
-        return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
-    }
+) : LLFirResolvableResolveSession(
+    useSiteKtModule = useSiteKtModule,
+    moduleKindProvider = LLSourceModuleKindProvider,
+    useSiteSessionFactory = useSiteSessionFactory
+) {
+    override val diagnosticProvider = LLSourceDiagnosticProvider(moduleProvider, sessionProvider)
+}
 
-    override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> {
-        val moduleComponents = getModuleComponentsForElement(ktFile)
-        return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
-    }
-
-    override fun getModuleKind(module: KtModule): ModuleKind {
+private object LLSourceModuleKindProvider : LLModuleKindProvider {
+    override fun getKind(module: KtModule): KtModuleKind {
         return when (module) {
-            is KtSourceModule -> ModuleKind.RESOLVABLE_MODULE
+            is KtSourceModule -> KtModuleKind.RESOLVABLE_MODULE
             is KtBuiltinsModule,
-            is KtLibraryModule -> ModuleKind.BINARY_MODULE
+            is KtLibraryModule -> KtModuleKind.BINARY_MODULE
             else -> unexpectedElementError("module", module)
         }
     }
-}
-
+}

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

@@ -0,0 +1,37 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.state
+
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
+
+interface LLModuleKindProvider {
+    /**
+     * Returns [KtModuleKind.BINARY_MODULE] if the [module] is treated as a binary for the current session,
+     * and [KtModuleKind.RESOLVABLE_MODULE] otherwise.
+     *
+     * In some cases, modules of the same type might be treated differently by the session, and have a different [KtModuleKind].
+     * For instance, for a resolvable library session, only the target library is considered resolvable, and its dependencies are binary.
+     */
+    fun getKind(module: KtModule): KtModuleKind
+}
+
+/**
+ * Specifies the way declarations are loaded and handled in the module.
+ */
+enum class KtModuleKind {
+    /**
+     * Declarations in resolvable modules might be in an unresolved state.
+     * Call [lazyResolveToPhase] on them before usage if needed.
+     */
+    RESOLVABLE_MODULE,
+
+    /**
+     * Declarations in binary modules always come fully resolved, as they are loaded from JAR/klib, where the complete type information
+     * is available.
+     */
+    BINARY_MODULE
+}

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

@@ -0,0 +1,21 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.state
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
+
+class LLModuleProvider(val useSiteModule: KtModule) {
+    /**
+     * Returns a [KtModule] for a given [element] in context of the current session.
+     *
+     * See [ProjectStructureProvider] for more information on contextual modules.
+     */
+    fun getModule(element: PsiElement): KtModule {
+        return ProjectStructureProvider.getModule(useSiteModule.project, element, useSiteModule)
+    }
+}

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

@@ -0,0 +1,19 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.state
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+
+interface LLScopeSessionProvider {
+    fun getScopeSession(session: LLFirSession): ScopeSession
+}
+
+internal object LLDefaultScopeSessionProvider : LLScopeSessionProvider {
+    override fun getScopeSession(session: LLFirSession): ScopeSession {
+        return session.getScopeSession()
+    }
+}

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

@@ -0,0 +1,55 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.state
+
+import com.intellij.psi.util.CachedValueProvider
+import com.intellij.psi.util.CachedValuesManager
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.fir.FirSession
+
+class LLSessionProvider(
+    val useSiteModule: KtModule,
+    private val useSiteSessionFactory: (KtModule) -> LLFirSession
+) {
+    private val useSiteSessionCached = CachedValuesManager.getManager(useSiteModule.project).createCachedValue {
+        val session = useSiteSessionFactory(useSiteModule)
+        CachedValueProvider.Result.create(session, session.createValidityTracker())
+    }
+
+    val useSiteSession: LLFirSession
+        get() = useSiteSessionCached.value
+
+    /**
+     * Returns a [FirSession] for the [module].
+     * For a binary module, the resulting session will be a binary (non-resolvable) one.
+     */
+    fun getSession(module: KtModule): LLFirSession {
+        return getSession(module, preferBinary = true)
+    }
+
+    /**
+     * Returns an analyzable [FirSession] for the module.
+     * For a binary module, the resulting session will still be a resolvable one.
+     *
+     * Note: prefer using [getSession] unless you need to perform resolution actively.
+     * Resolvable sessions for libraries are much less performant.
+     */
+    fun getResolvableSession(module: KtModule): LLFirResolvableModuleSession {
+        return getSession(module, preferBinary = false) as LLFirResolvableModuleSession
+    }
+
+    private fun getSession(module: KtModule, preferBinary: Boolean): LLFirSession {
+        if (module == useSiteSession.ktModule) {
+            return useSiteSession
+        }
+
+        val cache = LLFirSessionCache.getInstance(module.project)
+        return cache.getSession(module, preferBinary)
+    }
+}