[LL API] Provide basic support for scripts in K2 (KTIJ-21108)

2023-02-27 20:18:05 +09:00
parent ca7d1d0194
commit ae40ddcf03
13 changed files with 490 additions and 94 deletions
@@ -0,0 +1,45 @@
+/*
+ * 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.project.structure.builder
+
+import org.jetbrains.kotlin.analysis.project.structure.KtScriptModule
+import org.jetbrains.kotlin.analysis.project.structure.impl.KtScriptModuleImpl
+import org.jetbrains.kotlin.config.ApiVersion
+import org.jetbrains.kotlin.config.LanguageVersion
+import org.jetbrains.kotlin.config.LanguageVersionSettings
+import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl
+import org.jetbrains.kotlin.psi.KtFile
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.InvocationKind
+import kotlin.contracts.contract
+
+@KtModuleBuilderDsl
+public class KtScriptModuleBuilder : KtModuleBuilder() {
+    public lateinit var file: KtFile
+
+    public var languageVersionSettings: LanguageVersionSettings =
+        LanguageVersionSettingsImpl(LanguageVersion.LATEST_STABLE, ApiVersion.LATEST)
+
+    override fun build(): KtScriptModule {
+        return KtScriptModuleImpl(
+            directRegularDependencies,
+            directDependsOnDependencies,
+            directFriendDependencies,
+            platform,
+            project,
+            file,
+            languageVersionSettings
+        )
+    }
+}
+
+@OptIn(ExperimentalContracts::class)
+public inline fun buildKtScriptModule(init: KtScriptModuleBuilder.() -> Unit): KtScriptModule {
+    contract {
+        callsInPlace(init, InvocationKind.EXACTLY_ONCE)
+    }
+    return KtScriptModuleBuilder().apply(init).build()
+}
@@ -17,10 +17,7 @@ import com.intellij.psi.search.ProjectScope
 import com.intellij.util.io.URLUtil
 import org.jetbrains.kotlin.analysis.api.impl.base.util.LibraryUtils
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
-import org.jetbrains.kotlin.analysis.project.structure.builder.buildKtLibraryModule
-import org.jetbrains.kotlin.analysis.project.structure.builder.buildKtSdkModule
-import org.jetbrains.kotlin.analysis.project.structure.builder.buildKtSourceModule
-import org.jetbrains.kotlin.analysis.project.structure.builder.buildProjectStructureProvider
+import org.jetbrains.kotlin.analysis.project.structure.builder.*
 import org.jetbrains.kotlin.analyzer.common.CommonPlatformAnalyzerServices
 import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
 import org.jetbrains.kotlin.cli.jvm.config.javaSourceRoots
@@ -154,51 +151,66 @@ internal fun buildKtModuleProviderByCompilerConfiguration(
    project: Project,
    ktFiles: List<KtFile>,
 ): ProjectStructureProvider = buildProjectStructureProvider {
+    val (scriptFiles, ordinaryFiles) = ktFiles.partition { it.isScript() }
    val platform = JvmPlatforms.defaultJvmPlatform
-    addModule(
-        buildKtSourceModule {
-            val moduleName = compilerConfig.get(CommonConfigurationKeys.MODULE_NAME) ?: "<no module name provided>"

-            val libraryRoots = compilerConfig.jvmModularRoots + compilerConfig.jvmClasspathRoots
+    fun KtModuleBuilder.addModuleDependencies(moduleName: String) {
+        val libraryRoots = compilerConfig.jvmModularRoots + compilerConfig.jvmClasspathRoots
+        addRegularDependency(
+            buildKtLibraryModule {
+                contentScope = ProjectScope.getLibrariesScope(project)
+                this.platform = platform
+                this.project = project
+                binaryRoots = libraryRoots.map { it.toPath() }
+                libraryName = "Library for $moduleName"
+            }
+        )
+        compilerConfig.get(JVMConfigurationKeys.JDK_HOME)?.let { jdkHome ->
+            val vfm = VirtualFileManager.getInstance()
+            val jdkHomePath = jdkHome.toPath()
+            val jdkHomeVirtualFile = vfm.findFileByNioPath(jdkHomePath)
+            val binaryRoots = LibraryUtils.findClassesFromJdkHome(jdkHomePath).map {
+                Paths.get(URLUtil.extractPath(it))
+            }
            addRegularDependency(
-                buildKtLibraryModule {
-                    contentScope = ProjectScope.getLibrariesScope(project)
+                buildKtSdkModule {
+                    contentScope = GlobalSearchScope.fileScope(project, jdkHomeVirtualFile)
                    this.platform = platform
                    this.project = project
-                    binaryRoots = libraryRoots.map { it.toPath() }
-                    libraryName = "Library for $moduleName"
+                    this.binaryRoots = binaryRoots
+                    sdkName = "JDK for $moduleName"
                }
            )
-            compilerConfig.get(JVMConfigurationKeys.JDK_HOME)?.let { jdkHome ->
-                val vfm = VirtualFileManager.getInstance()
-                val jdkHomePath = jdkHome.toPath()
-                val jdkHomeVirtualFile = vfm.findFileByNioPath(jdkHomePath)
-                val binaryRoots = LibraryUtils.findClassesFromJdkHome(jdkHomePath).map {
-                    Paths.get(URLUtil.extractPath(it))
-                }
-                addRegularDependency(
-                    buildKtSdkModule {
-                        contentScope = GlobalSearchScope.fileScope(project, jdkHomeVirtualFile)
-                        this.platform = platform
-                        this.project = project
-                        this.binaryRoots = binaryRoots
-                        sdkName = "JDK for $moduleName"
-                    }
-                )
-            }
-
-            contentScope = TopDownAnalyzerFacadeForJVM.newModuleSearchScope(project, ktFiles)
-            this.platform = platform
-            this.project = project
-            this.moduleName = moduleName
-            addSourceRoots(
-                getPsiFilesFromPaths(
-                    project,
-                    getSourceFilePaths(compilerConfig, includeDirectoryRoot = true)
-                )
-            )
        }
-    )
+    }
+
+    for (scriptFile in scriptFiles) {
+        buildKtScriptModule {
+            this.project = project
+            this.platform = platform
+            this.file = scriptFile
+
+            addModuleDependencies("Script " + scriptFile.name)
+        }.apply(::addModule)
+    }
+
+    buildKtSourceModule {
+        this.project = project
+        this.platform = platform
+        this.moduleName = compilerConfig.get(CommonConfigurationKeys.MODULE_NAME) ?: "<no module name provided>"
+
+        addModuleDependencies(moduleName)
+
+        contentScope = TopDownAnalyzerFacadeForJVM.newModuleSearchScope(project, ordinaryFiles)
+        addSourceRoots(
+            getPsiFilesFromPaths(
+                project,
+                getSourceFilePaths(compilerConfig, includeDirectoryRoot = true)
+            )
+        )
+    }.apply(::addModule)
+
+
    this.platform = platform
    this.project = project
 }
@@ -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.project.structure.impl
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.KtScriptModule
+import org.jetbrains.kotlin.analysis.project.structure.computeTransitiveDependsOnDependencies
+import org.jetbrains.kotlin.config.LanguageVersionSettings
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
+
+internal class KtScriptModuleImpl(
+    override val directRegularDependencies: List<KtModule> = emptyList(),
+    override val directDependsOnDependencies: List<KtModule> = emptyList(),
+    override val directFriendDependencies: List<KtModule> = emptyList(),
+    override val platform: TargetPlatform = JvmPlatforms.defaultJvmPlatform,
+    override val project: Project,
+    override val file: KtFile,
+    override val languageVersionSettings: LanguageVersionSettings
+) : KtScriptModule, KtModuleWithPlatform {
+    override val transitiveDependsOnDependencies: List<KtModule> by lazy {
+        computeTransitiveDependsOnDependencies(directDependsOnDependencies)
+    }
+
+    override val analyzerServices: PlatformDependentAnalyzerServices
+        get() = super.analyzerServices
+
+    override val contentScope: GlobalSearchScope
+        get() = GlobalSearchScope.fileScope(file)
+}
@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionProv
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirLibraryOrLibrarySourceResolvableResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirNotUnderContentRootResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirScriptResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirSourceResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
 import org.jetbrains.kotlin.analysis.project.structure.*
@@ -69,6 +70,15 @@ internal class LLFirResolveSessionService(project: Project) {
                    )
                }

+                is KtScriptModule -> {
+                    LLFirScriptResolveSession(
+                        useSiteSession.moduleComponents.globalResolveComponents,
+                        sessionProviderStorage.project,
+                        useSiteKtModule,
+                        sessionProvider
+                    )
+                }
+
                is KtNotUnderContentRootModule -> {
                    LLFirNotUnderContentRootResolveSession(
                        useSiteSession.moduleComponents.globalResolveComponents,
@@ -100,6 +100,10 @@ private fun collectDesignationPath(target: FirElementWithResolvePhase): List<Fir
            return if (target.diagnostic == ConeDestructuringDeclarationsOnTopLevel) emptyList() else null
        }

+        is FirScript -> {
+            return emptyList()
+        }
+
        else -> {
            return null
        }
@@ -239,15 +243,15 @@ fun FirElementWithResolvePhase.tryCollectDesignation(): FirDesignation? =

 fun FirElementWithResolvePhase.tryCollectDesignationWithFile(): FirDesignationWithFile? {
    return when (this) {
+        is FirScript, is FirFileAnnotationsContainer -> {
+            val firFile = getContainingFile() ?: return null
+            FirDesignationWithFile(path = emptyList(), this, firFile)
+        }
        is FirDeclaration -> {
            val path = collectDesignationPath(this) ?: return null
            val firFile = getContainingFile() ?: return null
            FirDesignationWithFile(path, this, firFile)
        }
-        is FirFileAnnotationsContainer -> {
-            val firFile = getContainingFile() ?: return null
-            FirDesignationWithFile(path = emptyList(), this, firFile)
-        }
        else -> unexpectedElementError<FirElementWithResolvePhase>(this)
    }
 }
@@ -60,6 +60,7 @@ internal object PersistenceContextCollector {
            is FirClassLikeDeclaration -> declaration.symbol.classId.isLocal
            is FirCallableDeclaration -> declaration.symbol.callableId.isLocal
            is FirDanglingModifierList -> declaration.containingClass()?.classId?.isLocal == true
+            is FirScript -> false
            else -> error("Unsupported declaration ${declaration.renderWithType()}")
        }
        require(!isLocal) {
@@ -42,6 +42,11 @@ internal class LLFirModuleWithDependenciesSymbolProvider(
        dependencyProvider.getTopLevelCallableSymbolsTo(destination, packageFqName, name)
    }

+    @FirSymbolProviderInternals
+    fun getTopLevelCallableSymbolsWithoutDependencies(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
+        return buildList { getTopLevelCallableSymbolsToWithoutDependencies(this, packageFqName, name) }
+    }
+
    @FirSymbolProviderInternals
    fun getTopLevelCallableSymbolsToWithoutDependencies(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
        providers.forEach { it.getTopLevelCallableSymbolsTo(destination, packageFqName, name) }
@@ -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.sessions
+
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
+import org.jetbrains.kotlin.analysis.project.structure.KtScriptModule
+import org.jetbrains.kotlin.fir.BuiltinTypes
+import org.jetbrains.kotlin.fir.PrivateSessionConstructor
+
+internal class LLFirScriptSession @PrivateSessionConstructor constructor(
+    override val ktModule: KtScriptModule,
+    override val project: Project,
+    override val moduleComponents: LLFirModuleResolveComponents,
+    builtinTypes: BuiltinTypes
+) : LLFirResolvableModuleSession(builtinTypes)
@@ -58,17 +58,7 @@ internal object LLFirSessionFactory {

        val platform = module.platform
        val builtinsSession = LLFirBuiltinsSessionFactory.getInstance(project).getBuiltinsSession(platform)
-
-        val languageVersionSettings = object : LanguageVersionSettings by module.languageVersionSettings {
-            override fun getFeatureSupport(feature: LanguageFeature): LanguageFeature.State =
-                if (feature == LanguageFeature.EnableDfaWarningsInK2) LanguageFeature.State.ENABLED
-                else module.languageVersionSettings.getFeatureSupport(feature)
-
-            override fun supportsFeature(feature: LanguageFeature): Boolean =
-                getFeatureSupport(feature).let {
-                    it == LanguageFeature.State.ENABLED || it == LanguageFeature.State.ENABLED_WITH_WARNING
-                }
-        }
+        val languageVersionSettings = wrapLanguageVersionSettings(module.languageVersionSettings)

        val scopeProvider = FirKotlinScopeProvider(::wrapScopeWithJvmMapped)

@@ -115,35 +105,15 @@ internal object LLFirSessionFactory {
                register(FirSwitchableExtensionDeclarationsSymbolProvider::class, it)
            }

-            fun getOrCreateSessionForDependency(dependency: KtModule): LLFirSession? = when (dependency) {
-                is KtBuiltinsModule -> null //  build in is already added
-                is KtBinaryModule -> LLFirLibrarySessionFactory.getInstance(project).getLibrarySession(dependency, sessionsCache)
-                is KtSourceModule -> {
-                    createSourcesSession(
-                        project,
-                        dependency,
-                        globalResolveComponents,
-                        sessionInvalidator,
-                        sessionsCache,
-                        librariesSessionFactory = librariesSessionFactory,
-                        configureSession = configureSession,
-                    )
-                }
-                is KtNotUnderContentRootModule -> error("Module $module cannot depend on ${dependency::class}: $dependency")
-                is KtLibrarySourceModule -> error("Module $module cannot depend on ${dependency::class}: $dependency")
-            }
-
            val dependencyProvider = LLFirDependentModuleProvidersBySessions(this) {
-                module.directRegularDependencies.mapNotNullTo(this, ::getOrCreateSessionForDependency)
-
-                // The dependency provider needs to have access to all direct and indirect `dependsOn` dependencies, as `dependsOn`
-                // dependencies are transitive.
-                val directRegularDependenciesSet = module.directRegularDependencies.toSet()
-                module.transitiveDependsOnDependencies.forEach { dependency ->
-                    if (dependency !in directRegularDependenciesSet) {
-                        getOrCreateSessionForDependency(dependency)?.let(::add)
-                    }
-                }
+                processSourceDependencies(
+                    module,
+                    sessionsCache,
+                    globalResolveComponents,
+                    sessionInvalidator,
+                    librariesSessionFactory,
+                    configureSession
+                )

                add(builtinsSession)
            }
@@ -233,15 +203,21 @@ internal object LLFirSessionFactory {
            register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)

            val dependencyProvider = LLFirDependentModuleProvidersByProviders(this) {
-                // <all libraries scope> - <current library scope>
-                val librariesSearchScope =
-                    ProjectScope.getLibrariesScope(project).intersectWith(GlobalSearchScope.notScope(libraryModule.contentScope))
                add(builtinSession.symbolProvider)
-                addAll(
-                    LLFirLibraryProviderFactory.createLibraryProvidersForAllProjectLibraries(
-                        session, moduleData, scopeProvider, project, builtinTypes, librariesSearchScope
+
+                // Script dependencies are self-contained and should not depend on other libraries
+                if (module !is KtScriptDependencyModule) {
+                    // Add all libraries excluding the current one
+                    val librariesSearchScope = ProjectScope.getLibrariesScope(project)
+                        .intersectWith(GlobalSearchScope.notScope(libraryModule.contentScope))
+
+                    val restLibrariesProvider = LLFirLibraryProviderFactory.createLibraryProvidersForAllProjectLibraries(
+                        session, moduleData, scopeProvider,
+                        project, builtinTypes, librariesSearchScope
                    )
-                )
+
+                    addAll(restLibrariesProvider)
+                }
            }

            val javaSymbolProvider = createJavaSymbolProvider(this, moduleData, project, contentScope)
@@ -260,6 +236,91 @@ internal object LLFirSessionFactory {

            register(DEPENDENCIES_SYMBOL_PROVIDER_QUALIFIED_KEY, dependencyProvider)
            register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
+            register(LLFirFirClassByPsiClassProvider::class, LLFirFirClassByPsiClassProvider(this))
+
+            configureSession?.invoke(this)
+        }
+    }
+
+    fun createScriptSession(
+        project: Project,
+        module: KtScriptModule,
+        sessionInvalidator: LLFirSessionInvalidator,
+        sessionsCache: MutableMap<KtModule, LLFirSession>,
+        librariesSessionFactory: LLFirLibrarySessionFactory,
+        configureSession: (LLFirSession.() -> Unit)? = null
+    ): LLFirScriptSession {
+        sessionsCache[module]?.let { return it as LLFirScriptSession }
+        checkCanceled()
+
+        val platform = module.platform
+        val builtinsSession = LLFirBuiltinsSessionFactory.getInstance(project).getBuiltinsSession(platform)
+        val languageVersionSettings = wrapLanguageVersionSettings(module.languageVersionSettings)
+        val scopeProvider = FirKotlinScopeProvider(::wrapScopeWithJvmMapped)
+        val globalResolveComponents = LLFirGlobalResolveComponents(project)
+
+        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider, sessionInvalidator)
+        val contentScope = module.contentScope
+
+        val session = LLFirScriptSession(module, project, components, builtinsSession.builtinTypes)
+        sessionsCache[module] = session
+        components.session = session
+
+        return session.apply session@{
+            val moduleData = LLFirModuleData(module).apply { bindSession(this@session) }
+            registerModuleData(moduleData)
+            register(FirKotlinScopeProvider::class, scopeProvider)
+
+            registerIdeComponents(project)
+            registerCommonComponents(languageVersionSettings)
+            registerCommonComponentsAfterExtensionsAreConfigured()
+            registerCommonJavaComponents(JavaModuleResolver.getInstance(project))
+            registerResolveComponents()
+            registerJavaSpecificResolveComponents()
+
+            val provider = LLFirProvider(
+                this,
+                components,
+                FileBasedKotlinDeclarationProvider(module.file),
+                project.createPackageProvider(contentScope),
+                canContainKotlinPackage = true,
+            )
+
+            register(FirProvider::class, provider)
+            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver())
+
+            val dependencyProvider = LLFirDependentModuleProvidersBySessions(this) {
+                processSourceDependencies(
+                    module,
+                    sessionsCache,
+                    globalResolveComponents,
+                    sessionInvalidator,
+                    librariesSessionFactory,
+                    configureSession
+                )
+
+                add(builtinsSession)
+            }
+
+            val javaSymbolProvider = createJavaSymbolProvider(this, moduleData, project, contentScope)
+            register(JavaSymbolProvider::class, javaSymbolProvider)
+
+            register(
+                FirSymbolProvider::class,
+                LLFirModuleWithDependenciesSymbolProvider(
+                    this,
+                    dependencyProvider,
+                    providers = listOfNotNull(
+                        javaSymbolProvider,
+                        provider.symbolProvider,
+                    )
+                )
+            )
+
+            register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)
+            register(DEPENDENCIES_SYMBOL_PROVIDER_QUALIFIED_KEY, dependencyProvider)
+            register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
+            register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.Empty)

            configureSession?.invoke(this)
        }
@@ -334,6 +395,76 @@ internal object LLFirSessionFactory {
            configureSession?.invoke(this)
        }
    }
+
+    private fun wrapLanguageVersionSettings(original: LanguageVersionSettings): LanguageVersionSettings {
+        return object : LanguageVersionSettings by original {
+            override fun getFeatureSupport(feature: LanguageFeature): LanguageFeature.State {
+                return when (feature) {
+                    LanguageFeature.EnableDfaWarningsInK2 -> LanguageFeature.State.ENABLED
+                    else -> original.getFeatureSupport(feature)
+                }
+            }
+
+            override fun supportsFeature(feature: LanguageFeature): Boolean {
+                return when (getFeatureSupport(feature)) {
+                    LanguageFeature.State.ENABLED, LanguageFeature.State.ENABLED_WITH_WARNING -> true
+                    else -> false
+                }
+            }
+        }
+    }
+
+    private fun MutableList<LLFirSession>.processSourceDependencies(
+        module: KtModule,
+        sessionsCache: MutableMap<KtModule, LLFirSession>,
+        globalResolveComponents: LLFirGlobalResolveComponents,
+        sessionInvalidator: LLFirSessionInvalidator,
+        librariesSessionFactory: LLFirLibrarySessionFactory,
+        configureSession: (LLFirSession.() -> Unit)?
+    ) {
+        val project = module.project
+
+        fun getOrCreateSessionForDependency(dependency: KtModule): LLFirSession? = when (dependency) {
+            is KtBuiltinsModule -> {
+                // Built-ins are already added
+                null
+            }
+
+            is KtBinaryModule -> {
+                LLFirLibrarySessionFactory.getInstance(project).getLibrarySession(dependency, sessionsCache)
+            }
+
+            is KtSourceModule -> {
+                createSourcesSession(
+                    project,
+                    dependency,
+                    globalResolveComponents,
+                    sessionInvalidator,
+                    sessionsCache,
+                    librariesSessionFactory = librariesSessionFactory,
+                    configureSession = configureSession,
+                )
+            }
+
+            is KtScriptModule,
+            is KtScriptDependencyModule,
+            is KtNotUnderContentRootModule,
+            is KtLibrarySourceModule -> {
+                error("Module $module cannot depend on ${dependency::class}: $dependency")
+            }
+        }
+
+        module.directRegularDependencies.mapNotNullTo(this@processSourceDependencies, ::getOrCreateSessionForDependency)
+
+        // The dependency provider needs to have access to all direct and indirect `dependsOn` dependencies, as `dependsOn`
+        // dependencies are transitive.
+        val directRegularDependenciesSet = module.directRegularDependencies.toSet()
+        module.transitiveDependsOnDependencies.forEach { dependency ->
+            if (dependency !in directRegularDependenciesSet) {
+                getOrCreateSessionForDependency(dependency)?.let(this::add)
+            }
+        }
+    }
 }

@Deprecated(
@@ -40,6 +40,10 @@ class LLFirSessionProviderStorage(val project: Project) {
            createSessionProviderForLibraryOrLibrarySource(useSiteKtModule, configureSession)
        }

+        is KtScriptModule -> {
+            createSessionProviderForScriptSession(useSiteKtModule, configureSession)
+        }
+
        is KtNotUnderContentRootModule -> {
            createSessionProviderForNotUnderContentRootSession(useSiteKtModule, configureSession)
        }
@@ -89,6 +93,25 @@ class LLFirSessionProviderStorage(val project: Project) {
        return LLFirSessionProvider(project, session, KtModuleToSessionMappingByWeakValueMapImpl(sessions))
    }

+    private fun createSessionProviderForScriptSession(
+        useSiteKtModule: KtScriptModule,
+        configureSession: (LLFirSession.() -> Unit)?
+    ): LLFirSessionProvider {
+        val (sessions, session) = sourceAsUseSiteSessionCache.withMappings(project) { mappings ->
+            val sessions = mutableMapOf<KtModule, LLFirSession>().apply { putAll(mappings) }
+            val session = LLFirSessionFactory.createScriptSession(
+                project,
+                useSiteKtModule,
+                sourceAsUseSiteSessionCache.sessionInvalidator,
+                sessions,
+                librariesSessionFactory,
+                configureSession = configureSession,
+            )
+            sessions to session
+        }
+        return LLFirSessionProvider(project, session, KtModuleToSessionMappingByWeakValueMapImpl(sessions))
+    }
+
    private fun createSessionProviderForNotUnderContentRootSession(
        useSiteKtModule: KtNotUnderContentRootModule,
        configureSession: (LLFirSession.() -> Unit)?
@@ -189,6 +212,8 @@ private class FirSessionWithModificationTracker(
        val outOfBlockTracker = when (ktModule) {
            is KtSourceModule -> trackerFactory.createModuleWithoutDependenciesOutOfBlockModificationTracker(ktModule)
            is KtNotUnderContentRootModule -> ModificationTracker { ktModule.file?.modificationStamp ?: 0 }
+            is KtScriptModule -> ModificationTracker { ktModule.file.modificationStamp }
+            is KtScriptDependencyModule -> ModificationTracker { ktModule.file?.modificationStamp ?: 0 }
            else -> null
        }
        modificationTracker = CompositeModificationTracker.create(
@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveCompone
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
+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.LLFirSessionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirDeclarationForCompiledElementSearcher
@@ -49,6 +50,9 @@ internal abstract class LLFirResolvableResolveSession(
    override fun getSessionFor(module: KtModule): LLFirSession =
        sessionProvider.getSession(module)

+    protected open fun getResolvableSessionFor(module: KtModule): LLFirResolvableModuleSession =
+        sessionProvider.getResolvableSession(module)
+
    override fun getScopeSessionFor(firSession: FirSession): ScopeSession {
        requireIsInstance<LLFirSession>(firSession)
        return firSession.getScopeSession()
@@ -66,7 +70,7 @@ internal abstract class LLFirResolvableResolveSession(

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

    override fun resolveToFirSymbol(
@@ -86,7 +90,7 @@ internal abstract class LLFirResolvableResolveSession(
        }

        val ktModule = ktDeclaration.getKtModule(project)
-        val firSession = sessionProvider.getSession(ktModule)
+        val firSession = getSessionFor(ktModule)
        val searcher = FirDeclarationForCompiledElementSearcher(firSession.symbolProvider)
        val firDeclaration = searcher.findNonLocalDeclaration(ktDeclaration)
        return firDeclaration.symbol
@@ -107,7 +111,7 @@ internal abstract class LLFirResolvableResolveSession(
            }

        if (ktDeclaration == nonLocalNamedDeclaration) {
-            val session = sessionProvider.getResolvableSession(module)
+            val session = getResolvableSessionFor(module)
            return nonLocalNamedDeclaration.findSourceNonLocalFirDeclaration(
                firFileBuilder = session.moduleComponents.firFileBuilder,
                provider = session.firProvider,
@@ -0,0 +1,74 @@
+/*
+ * 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.openapi.components.ServiceManager
+import com.intellij.openapi.project.Project
+import com.intellij.util.containers.CollectionFactory
+import com.intellij.util.containers.FactoryMap
+import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirGlobalResolveComponents
+import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionService
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.*
+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(
+    override val globalComponents: LLFirGlobalResolveComponents,
+    override val project: Project,
+    override val useSiteKtModule: KtModule,
+    sessionProvider: LLFirSessionProvider
+) : LLFirResolvableResolveSession(sessionProvider) {
+    private val dependencyCache: Map<KtModule, LLFirSession> =
+        FactoryMap.createMap(::createDependencySession, CollectionFactory::createConcurrentSoftValueMap)
+
+    private fun createDependencySession(module: KtModule): LLFirSession {
+        when (module) {
+            is KtScriptDependencyModule -> {
+                val resolveSessionService = ServiceManager.getService(project, LLFirResolveSessionService::class.java)
+                return resolveSessionService.getFirResolveSession(module).useSiteFirSession
+            }
+            else -> error("Unsupported script dependency session type: ${module.javaClass.name}")
+        }
+    }
+
+    override fun getSessionFor(module: KtModule): LLFirSession {
+        if (module is KtScriptDependencyModule) {
+            return dependencyCache.getValue(module)
+        }
+
+        return super.getSessionFor(module)
+    }
+
+    override fun getResolvableSessionFor(module: KtModule): LLFirResolvableModuleSession {
+        if (module is KtScriptDependencyModule && module is KtLibrarySourceModule) {
+            return dependencyCache[module] as LLFirResolvableModuleSession
+        }
+
+        return super.getResolvableSessionFor(module)
+    }
+
+    override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
+        val moduleComponents = getModuleComponentsForElement(element)
+        return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
+    }
+
+    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 {
+        return when (module) {
+            useSiteKtModule, is KtSourceModule -> ModuleKind.RESOLVABLE_MODULE
+            is KtBuiltinsModule, is KtLibraryModule -> ModuleKind.BINARY_MODULE
+            else -> unexpectedElementError("module", module)
+        }
+    }
+}
@@ -10,6 +10,7 @@ import com.intellij.psi.PsiFile
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
 import java.nio.file.Path

@@ -177,6 +178,34 @@ public class KtBuiltinsModule(
    override fun hashCode(): Int = platform.hashCode()
 }

+/**
+ * A module for a Kotlin script file.
+ */
+public interface KtScriptModule : KtModule {
+    /**
+     * A script PSI.
+     */
+    public val file: KtFile
+
+    /**
+     * A set of Kotlin settings, like API version, supported features and flags.
+     */
+    public val languageVersionSettings: LanguageVersionSettings
+
+    override val moduleDescription: String
+        get() = "Script " + file.name
+}
+
+/**
+ * A module for Kotlin script dependencies.
+ * Must be either a [KtLibraryModule] or [KtLibrarySourceModule].
+ */
+public interface KtScriptDependencyModule : KtModule {
+    /**
+     * A `VirtualFile` that backs the dependent script PSI, or `null` if the module is for project-level dependencies.
+     */
+    public val file: KtFile?
+}

 /**
 * A set of sources which live outside the project content root. E.g, testdata files or source files of some other project.