K2 Scripting: compiler extension for script configuration

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
+import org.jetbrains.kotlin.fir.builder.FirScriptConfiguratorExtension
 import org.jetbrains.kotlin.fir.resolve.FirSamConversionTransformerExtension
 import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.reflect.KClass
@@ -30,6 +31,7 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirExtensionSessionComponent::class,
             FirSamConversionTransformerExtension::class,
             FirAssignExpressionAltererExtension::class,
+            FirScriptConfiguratorExtension::class,
         )
     }
 
@@ -83,6 +85,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             registerExtension(FirAssignExpressionAltererExtension::class, this)
         }
 
+        @JvmName("plusScriptConfiguratorExtension")
+        operator fun (FirScriptConfiguratorExtension.Factory).unaryPlus() {
+            registerExtension(FirScriptConfiguratorExtension::class, this)
+        }
+
         // ------------------ reference methods ------------------
 
         @JvmName("plusStatusTransformerExtension")
@@ -130,6 +137,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirAssignExpressionAltererExtension.Factory { this.invoke(it) }.unaryPlus()
         }
 
+        @JvmName("plusScriptConfiguratorExtension")
+        operator fun ((FirSession) -> FirScriptConfiguratorExtension).unaryPlus() {
+            FirScriptConfiguratorExtension.Factory { this.invoke(it) }.unaryPlus()
+        }
+
         // ------------------ utilities ------------------
 
         @JvmName("bindLeft")

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.fir.diagnostics.*
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
+import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.references.builder.*
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -1076,6 +1077,7 @@ open class RawFirBuilder(
                         }
                     }
                 }
+                baseSession.extensionService.scriptConfigurators.forEach { with(it) { configure(containingFile) } }
             }
         }
 

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/FirScriptConfiguratorExtension.kt

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.builder
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.builder.FirFileBuilder
+import org.jetbrains.kotlin.fir.declarations.builder.FirScriptBuilder
+import org.jetbrains.kotlin.fir.extensions.FirExtension
+import org.jetbrains.kotlin.fir.extensions.FirExtensionPointName
+import org.jetbrains.kotlin.fir.extensions.FirExtensionService
+import kotlin.reflect.KClass
+
+abstract class FirScriptConfiguratorExtension(
+    session: FirSession,
+) : FirExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("ScriptConfigurator")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension> = FirScriptConfiguratorExtension::class
+
+    fun interface Factory : FirExtension.Factory<FirScriptConfiguratorExtension>
+
+    abstract fun FirScriptBuilder.configure(fileBuilder: FirFileBuilder)
+}
+
+val FirExtensionService.scriptConfigurators: List<FirScriptConfiguratorExtension> by FirExtensionService.registeredExtensions()

plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/definitions/ScriptDependenciesProvider.kt

@@ -17,6 +17,10 @@ import kotlin.script.experimental.api.ScriptCompilationConfiguration
 import kotlin.script.experimental.api.valueOrNull
 import kotlin.script.experimental.dependencies.ScriptDependencies
 
+// Note: misleading name, it is now general configurations provider, not only for dependencies
+// but we are affraid to touch it so far, since the impact should be assessed first
+// TODO: consider deprecating completely and swith to a new interface in the K2
+// TODO: support SourceCode (or KtSourceFile) as a key
 open class ScriptDependenciesProvider constructor(
     protected val project: Project
 ) {

plugins/scripting/scripting-compiler/build.gradle.kts

@@ -11,6 +11,7 @@ dependencies {
     compileOnly(project(":compiler:psi"))
     compileOnly(project(":compiler:plugin-api"))
     compileOnly(project(":compiler:fir:entrypoint"))
+    compileOnly(project(":compiler:fir:raw-fir:raw-fir.common"))
     compileOnly(project(":compiler:cli"))
     compileOnly(project(":core:descriptors.runtime"))
     compileOnly(project(":compiler:ir.tree"))

plugins/scripting/scripting-compiler/resources/META-INF/services/org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar

@@ -0,0 +1 @@
+org.jetbrains.kotlin.scripting.compiler.plugin.ScriptingK2CompilerPluginRegistrar

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/FirScriptingCompilerExtensionRegistrar.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.scripting.compiler.plugin
+
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.scripting.compiler.plugin.services.FirScriptConfiguratorExtensionImpl
+import org.jetbrains.kotlin.scripting.compiler.plugin.services.FirScriptDefinitionProviderService
+import org.jetbrains.kotlin.scripting.configuration.ScriptingConfigurationKeys
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+
+class FirScriptingCompilerExtensionRegistrar(
+    private val hostConfiguration: ScriptingHostConfiguration, private val compilerConfiguration: CompilerConfiguration
+) : FirExtensionRegistrar() {
+
+    override fun ExtensionRegistrarContext.configurePlugin() {
+        if (compilerConfiguration.getBoolean(ScriptingConfigurationKeys.DISABLE_SCRIPTING_PLUGIN_OPTION)) return
+
+        configureScriptDefinitions(compilerConfiguration, hostConfiguration, this::class.java.classLoader)
+        val definitionSources = compilerConfiguration.getList(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS_SOURCES)
+        val definitions = compilerConfiguration.getList(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS)
+        if (definitionSources.isNotEmpty() || definitions.isNotEmpty()) {
+            +FirScriptDefinitionProviderService.getFactory(definitions, definitionSources)
+        }
+
+        +FirScriptConfiguratorExtensionImpl.getFactory(hostConfiguration)
+    }
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptingCompilerConfigurationExtension.kt

@@ -33,7 +33,6 @@ class ScriptingCompilerConfigurationExtension(
 
         if (!configuration.getBoolean(ScriptingConfigurationKeys.DISABLE_SCRIPTING_PLUGIN_OPTION)) {
 
-            val messageCollector = configuration.get(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY) ?: MessageCollector.NONE
             val projectRoot = project.run { basePath ?: baseDir?.canonicalPath }?.let(::File)
             if (projectRoot != null) {
                 configuration.put(
@@ -48,43 +47,7 @@ class ScriptingCompilerConfigurationExtension(
                 }
             }
 
-            val explicitScriptDefinitions = configuration.getList(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS_CLASSES)
+            configureScriptDefinitions(configuration, hostConfiguration, this::class.java.classLoader)
-
-            if (explicitScriptDefinitions.isNotEmpty()) {
-                configureScriptDefinitions(
-                    explicitScriptDefinitions,
-                    configuration,
-                    this::class.java.classLoader,
-                    messageCollector,
-                    hostConfiguration
-                )
-            }
-            // If not disabled explicitly, we should always support at least the standard script definition
-            if (!configuration.getBoolean(JVMConfigurationKeys.DISABLE_STANDARD_SCRIPT_DEFINITION) &&
-                configuration.getList(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS).none { it.isDefault }
-            ) {
-                configuration.add(
-                    ScriptingConfigurationKeys.SCRIPT_DEFINITIONS,
-                    ScriptDefinition.getDefault(hostConfiguration)
-                )
-            }
-
-            val definitionsFromClasspath =
-                if (configuration.getBoolean(ScriptingConfigurationKeys.DISABLE_SCRIPT_DEFINITIONS_FROM_CLASSPATH_OPTION)) null
-                else
-                    ScriptDefinitionsFromClasspathDiscoverySource(
-                        configuration.jvmClasspathRoots,
-                        hostConfiguration,
-                        messageCollector.reporter
-                    )
-            val autoloadedScriptDefinitions =
-                if (configuration.getBoolean(ScriptingConfigurationKeys.DISABLE_SCRIPT_DEFINITIONS_AUTOLOADING_OPTION)) null
-                else AutoloadedScriptDefinitions(hostConfiguration, this::class.java.classLoader, messageCollector.reporter)
-
-            configuration.addAll(
-                ScriptingConfigurationKeys.SCRIPT_DEFINITIONS_SOURCES,
-                listOfNotNull(definitionsFromClasspath, autoloadedScriptDefinitions)
-            )
 
             val scriptDefinitionProvider = ScriptDefinitionProvider.getInstance(project) as? CliScriptDefinitionProvider
             if (scriptDefinitionProvider != null) {
@@ -111,3 +74,48 @@ class ScriptingCompilerConfigurationExtension(
     }
 }
 
+internal fun configureScriptDefinitions(
+    configuration: CompilerConfiguration,
+    hostConfiguration: ScriptingHostConfiguration,
+    classLoader: ClassLoader
+) {
+    val messageCollector = configuration.get(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY) ?: MessageCollector.NONE
+
+    val explicitScriptDefinitions = configuration.getList(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS_CLASSES)
+
+    if (explicitScriptDefinitions.isNotEmpty()) {
+        configureScriptDefinitions(
+            explicitScriptDefinitions,
+            configuration,
+            classLoader,
+            messageCollector,
+            hostConfiguration
+        )
+    }
+    // If not disabled explicitly, we should always support at least the standard script definition
+    if (!configuration.getBoolean(JVMConfigurationKeys.DISABLE_STANDARD_SCRIPT_DEFINITION) &&
+        configuration.getList(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS).none { it.isDefault }
+    ) {
+        configuration.add(
+            ScriptingConfigurationKeys.SCRIPT_DEFINITIONS,
+            ScriptDefinition.getDefault(hostConfiguration)
+        )
+    }
+
+    val definitionsFromClasspath =
+        if (configuration.getBoolean(ScriptingConfigurationKeys.DISABLE_SCRIPT_DEFINITIONS_FROM_CLASSPATH_OPTION)) null
+        else
+            ScriptDefinitionsFromClasspathDiscoverySource(
+                configuration.jvmClasspathRoots,
+                hostConfiguration,
+                messageCollector.reporter
+            )
+    val autoloadedScriptDefinitions =
+        if (configuration.getBoolean(ScriptingConfigurationKeys.DISABLE_SCRIPT_DEFINITIONS_AUTOLOADING_OPTION)) null
+        else AutoloadedScriptDefinitions(hostConfiguration, classLoader, messageCollector.reporter)
+
+    configuration.addAll(
+        ScriptingConfigurationKeys.SCRIPT_DEFINITIONS_SOURCES,
+        listOfNotNull(definitionsFromClasspath, autoloadedScriptDefinitions)
+    )
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/ScriptJvmCompilerImpls.kt

@@ -42,6 +42,8 @@ import org.jetbrains.kotlin.resolve.jvm.extensions.PackageFragmentProviderExtens
 import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
 import org.jetbrains.kotlin.scripting.compiler.plugin.ScriptCompilerProxy
 import org.jetbrains.kotlin.scripting.compiler.plugin.dependencies.ScriptsCompilationDependencies
+import org.jetbrains.kotlin.scripting.compiler.plugin.services.scriptDefinitionProviderService
+import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionProvider
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.ScriptingHostConfiguration
@@ -363,6 +365,11 @@ private fun doCompileWithK2(
         friendDependencies(kotlinCompilerConfiguration[JVMConfigurationKeys.FRIEND_PATHS] ?: emptyList())
     }
 
+    session.scriptDefinitionProviderService?.run {
+        definitionProvider = ScriptDefinitionProvider.getInstance(context.environment.project)
+        configurationProvider = ScriptDependenciesProvider.getInstance(context.environment.project)
+    }
+
     val rawFir = session.buildFirFromKtFiles(sourceFiles)
 
     val (scopeSession, fir) = session.runResolution(rawFir)

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/pluginRegisrar.kt

@@ -14,12 +14,14 @@ import org.jetbrains.kotlin.cli.common.extensions.ScriptEvaluationExtension
 import org.jetbrains.kotlin.cli.common.extensions.ShellExtension
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.extensions.CollectAdditionalSourcesExtension
 import org.jetbrains.kotlin.extensions.CompilerConfigurationExtension
 import org.jetbrains.kotlin.extensions.ProcessSourcesBeforeCompilingExtension
 import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
+import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
 import org.jetbrains.kotlin.resolve.extensions.ExtraImportsProviderExtension
 import org.jetbrains.kotlin.resolve.extensions.SyntheticResolveExtension
 import org.jetbrains.kotlin.scripting.compiler.plugin.definitions.CliScriptDefinitionProvider
@@ -76,6 +78,26 @@ class ScriptingCompilerConfigurationComponentRegistrar : ComponentRegistrar {
     }
 }
 
+// Scripting infrastructure still depends on project-based components, therefore we still need a separate registrar above - ScriptingCompilerConfigurationComponentRegistrar
+// TODO: refactor components and migrate the plugin to the project-independent operation
+class ScriptingK2CompilerPluginRegistrar : CompilerPluginRegistrar() {
+    companion object {
+        fun registerComponents(extensionStorage: ExtensionStorage, compilerConfiguration: CompilerConfiguration) = with(extensionStorage) {
+            val hostConfiguration = ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration) {
+                // TODO: add jdk path and other params if needed
+            }
+            FirExtensionRegistrarAdapter.registerExtension(FirScriptingCompilerExtensionRegistrar(hostConfiguration, compilerConfiguration))
+        }
+    }
+
+    override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
+        registerComponents(this, configuration)
+    }
+
+    override val supportsK2: Boolean
+        get() = true
+}
+
 private inline fun withClassloadingProblemsReporting(messageCollector: MessageCollector?, body: () -> Unit) {
     try {
         body()

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/services/FirScriptConfigurationExtensionImpl.kt

@@ -0,0 +1,148 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.scripting.compiler.plugin.services
+
+import org.jetbrains.kotlin.*
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.builder.FirScriptConfiguratorExtension
+import org.jetbrains.kotlin.fir.builder.FirScriptConfiguratorExtension.Factory
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.builder.*
+import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.declarations.primaryConstructorIfAny
+import org.jetbrains.kotlin.fir.moduleData
+import org.jetbrains.kotlin.fir.resolve.providers.dependenciesSymbolProvider
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.types.builder.buildUserTypeRef
+import org.jetbrains.kotlin.fir.types.impl.FirQualifierPartImpl
+import org.jetbrains.kotlin.fir.types.impl.FirTypeArgumentListImpl
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.scripting.resolve.KtFileScriptSource
+import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.FileScriptSource
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+import kotlin.script.experimental.host.StringScriptSource
+
+
+class FirScriptConfiguratorExtensionImpl(
+    session: FirSession,
+    // TODO: left here because it seems it will be needed soon, remove supression if used or remove the param if it is not the case
+    @Suppress("UNUSED_PARAMETER") hostConfiguration: ScriptingHostConfiguration
+) : FirScriptConfiguratorExtension(session) {
+    @OptIn(SymbolInternals::class)
+    override fun FirScriptBuilder.configure(fileBuilder: FirFileBuilder) {
+
+        // TODO: rewrite/extract decision logic for clarity
+        val compilationConfiguration = session.scriptDefinitionProviderService?.let { providerService ->
+            fileBuilder.sourceFile?.toSourceCode()?.let { script ->
+                val ktFile = (script as? KtFileScriptSource)?.ktFile ?: error("only PSI scripts are supported at the moment")
+                providerService.configurationProvider?.getScriptConfigurationResult(ktFile)?.valueOrNull()?.configuration
+                    ?: providerService.definitionProvider?.findDefinition(script)?.compilationConfiguration
+            } ?: providerService.definitionProvider?.getDefaultDefinition()?.compilationConfiguration
+        }
+
+        if (compilationConfiguration != null) {
+
+            compilationConfiguration[ScriptCompilationConfiguration.defaultImports]?.forEach { defaultImport ->
+                val trimmed = defaultImport.trim()
+                val endsWithStar = trimmed.endsWith("*")
+                val stripped = if (endsWithStar) trimmed.substring(0, trimmed.length - 1) else trimmed
+                val fqName = FqName.fromSegments(stripped.split("."))
+                fileBuilder.imports += buildImport {
+                    importedFqName = fqName
+                    isAllUnder = endsWithStar
+                }
+            }
+
+            compilationConfiguration[ScriptCompilationConfiguration.baseClass]?.let { baseClass ->
+                val baseClassFqn = FqName.fromSegments(baseClass.typeName.split("."))
+                contextReceivers.add(buildContextReceiverWithFqName(baseClassFqn))
+
+                val baseClassSymbol =
+                    session.dependenciesSymbolProvider.getClassLikeSymbolByClassId(ClassId(baseClassFqn.parent(), baseClassFqn.shortName()))
+                            as? FirRegularClassSymbol
+                if (baseClassSymbol != null) {
+                    // assuming that if base class will be unresolved, the error will be reported on the contextReceiver
+                    baseClassSymbol.fir.primaryConstructorIfAny(session)?.fir?.valueParameters?.forEach { baseCtorParameter ->
+                        parameters.add(
+                            buildProperty {
+                                moduleData = session.moduleData
+                                origin = FirDeclarationOrigin.ScriptCustomization
+                                // TODO: copy type parameters?
+                                returnTypeRef = baseCtorParameter.returnTypeRef
+                                this.name = baseCtorParameter.name
+                                this.symbol = FirPropertySymbol(this.name)
+                                status = FirDeclarationStatusImpl(Visibilities.Local, Modality.FINAL)
+                                isLocal = true
+                                isVar = false
+                            }
+                        )
+                    }
+                }
+            }
+
+            compilationConfiguration[ScriptCompilationConfiguration.implicitReceivers]?.forEach { implicitReceiver ->
+                contextReceivers.add(buildContextReceiverWithFqName(FqName.fromSegments(implicitReceiver.typeName.split("."))))
+            }
+
+            compilationConfiguration[ScriptCompilationConfiguration.providedProperties]?.forEach { propertyName, propertyType ->
+                val typeRef = buildUserTypeRef {
+                    isMarkedNullable = propertyType.isNullable
+                    propertyType.typeName.split(".").forEach {
+                        qualifier.add(FirQualifierPartImpl(null, Name.identifier(it), FirTypeArgumentListImpl(null)))
+                    }
+                }
+                parameters.add(
+                    buildProperty {
+                        moduleData = session.moduleData
+                        origin = FirDeclarationOrigin.ScriptCustomization
+                        returnTypeRef = typeRef
+                        this.name = Name.identifier(propertyName)
+                        this.symbol = FirPropertySymbol(this.name)
+                        status = FirDeclarationStatusImpl(Visibilities.Local, Modality.FINAL)
+                        isLocal = true
+                        isVar = false
+                    }
+                )
+            }
+        }
+    }
+
+    private fun buildContextReceiverWithFqName(baseClassFqn: FqName) =
+        buildContextReceiver {
+            typeRef = buildUserTypeRef {
+                isMarkedNullable = false
+                qualifier.addAll(
+                    baseClassFqn.pathSegments().map {
+                        FirQualifierPartImpl(null, it, FirTypeArgumentListImpl(null))
+                    }
+                )
+            }
+        }
+
+    companion object {
+        fun getFactory(hostConfiguration: ScriptingHostConfiguration): Factory {
+            return Factory { session -> FirScriptConfiguratorExtensionImpl(session, hostConfiguration) }
+        }
+    }
+
+}
+
+fun KtSourceFile.toSourceCode(): SourceCode? = when (this) {
+    is KtPsiSourceFile -> (psiFile as? KtFile)?.let(::KtFileScriptSource) ?: VirtualFileScriptSource(psiFile.virtualFile)
+    is KtVirtualFileSourceFile -> VirtualFileScriptSource(virtualFile)
+    is KtIoFileSourceFile -> FileScriptSource(file)
+    is KtInMemoryTextSourceFile -> StringScriptSource(text.toString(), name)
+    else -> null
+}

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/services/FirScriptDefinitionProviderService.kt

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.scripting.compiler.plugin.services
+
+import com.intellij.mock.MockProject
+import com.intellij.openapi.util.Disposer
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.extensions.FirExtensionSessionComponent
+import org.jetbrains.kotlin.fir.extensions.FirExtensionSessionComponent.Factory
+import org.jetbrains.kotlin.scripting.compiler.plugin.definitions.CliScriptDefinitionProvider
+import org.jetbrains.kotlin.scripting.compiler.plugin.definitions.CliScriptDependenciesProvider
+import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
+import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionProvider
+import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionsSource
+import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
+
+class FirScriptDefinitionProviderService(
+    session: FirSession,
+    private val makeDefaultDefinitionProvider: () -> ScriptDefinitionProvider,
+    private val makeDefaultConfigurationProvider: () -> ScriptDependenciesProvider
+) : FirExtensionSessionComponent(session) {
+
+    // TODO: get rid of project-based implementation, write and use own singleton in K2
+
+    private var _definitionProvider: ScriptDefinitionProvider? = null
+    var definitionProvider: ScriptDefinitionProvider?
+        get() = synchronized(this) {
+            if (_definitionProvider == null) _definitionProvider = makeDefaultDefinitionProvider()
+            _definitionProvider
+        }
+        set(value) { synchronized(this) { _definitionProvider = value} }
+
+    private var _configurationProvider: ScriptDependenciesProvider? = null
+    var configurationProvider: ScriptDependenciesProvider?
+        get() = synchronized(this) {
+            if (_configurationProvider == null) _configurationProvider = makeDefaultConfigurationProvider()
+            _configurationProvider
+        }
+        set(value) { synchronized(this) { _configurationProvider = value} }
+
+    companion object {
+        fun getFactory(
+            definitions: List<ScriptDefinition>,
+            definitionSources: List<ScriptDefinitionsSource>
+        ): Factory {
+            return Factory { session ->
+                FirScriptDefinitionProviderService(
+                    session,
+                    makeDefaultDefinitionProvider = {
+                        CliScriptDefinitionProvider().also {
+                            it.setScriptDefinitionsSources(definitionSources)
+                            it.setScriptDefinitions(definitions)
+                        }
+                    },
+                    makeDefaultConfigurationProvider = {
+                        // TODO: check if memory can leak in MockProject (probably not too important, since currently the providers are set externaly in important cases)
+                        CliScriptDependenciesProvider(MockProject(null, Disposer.newDisposable()))
+                    }
+                )
+            }
+        }
+    }
+}
+
+val FirSession.scriptDefinitionProviderService: FirScriptDefinitionProviderService? by FirSession.nullableSessionComponentAccessor()

plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptingCompilerPluginTest.kt

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.cli.jvm.compiler.KotlinToJVMBytecodeCompiler
 import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoot
 import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
 import org.jetbrains.kotlin.compiler.plugin.AbstractCliOption
+import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
@@ -75,6 +76,7 @@ class ScriptingCompilerPluginTest : TestCase() {
             confBody()
         }
         configuration.add(ComponentRegistrar.PLUGIN_COMPONENT_REGISTRARS, ScriptingCompilerConfigurationComponentRegistrar())
+        configuration.add(CompilerPluginRegistrar.COMPILER_PLUGIN_REGISTRARS, ScriptingK2CompilerPluginRegistrar())
 
         return KotlinCoreEnvironment.createForTests(disposable, configuration, EnvironmentConfigFiles.JVM_CONFIG_FILES)
     }