[KT-58817] Add script-related extensions for FIR session

The following two registrars were missing:
- FirScriptingCompilerExtensionIdeRegistrar
- FirScriptingSamWithReceiverExtensionRegistrar

`FirScriptingCompilerExtensionIdeRegistrar` was introduced as a copy of
`FirScriptingCompilerExtensionRegistrar` adapted for usage from IDE.

`FirRegisteredPluginAnnotations` and `FirJvmTypeMapper` are mandatory at
`FirExtensionService.registerExtensions` and reported as missing at
runtime.

analysis/low-level-api-fir/build.gradle.kts

@@ -27,6 +27,8 @@ dependencies {
     implementation(project(":analysis:analysis-api-providers"))
     implementation(project(":analysis:analysis-api"))
     implementation(project(":analysis:analysis-internal-utils"))
+    implementation(project(":kotlin-scripting-compiler"))
+    implementation(project(":kotlin-scripting-common"))
 
     // We cannot use the latest version `3.1.5` because it doesn't support Java 8.
     implementation("com.github.ben-manes.caffeine:caffeine:2.9.3")

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

@@ -0,0 +1,35 @@
+/*
+ * 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 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.definitions.ScriptDefinition
+import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionsSource
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+
+/**
+ * The class is a copy of [org.jetbrains.kotlin.scripting.compiler.plugin.FirScriptingCompilerExtensionRegistrar] adapted for the usage
+ * outside the compiler context. Difference lies in its constructor parameters, namely in
+ * [org.jetbrains.kotlin.config.CompilerConfiguration] replaced with a pair of explicit [scriptDefinitionSources] and [scriptDefinitions].
+ */
+internal class FirScriptingCompilerExtensionIdeRegistrar(
+    private val hostConfiguration: ScriptingHostConfiguration,
+    private val scriptDefinitionSources: List<ScriptDefinitionsSource>,
+    private val scriptDefinitions: List<ScriptDefinition>
+) : FirExtensionRegistrar() {
+
+    override fun ExtensionRegistrarContext.configurePlugin() {
+        val definitionSources = scriptDefinitionSources
+        val definitions = scriptDefinitions
+        if (definitionSources.isNotEmpty() || definitions.isNotEmpty()) {
+            +FirScriptDefinitionProviderService.getFactory(definitions, definitionSources)
+        }
+
+        +FirScriptConfiguratorExtensionImpl.getFactory(hostConfiguration)
+    }
+}

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

@@ -44,8 +44,12 @@ import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
+import org.jetbrains.kotlin.scripting.compiler.plugin.FirScriptingSamWithReceiverExtensionRegistrar
+import org.jetbrains.kotlin.scripting.definitions.findScriptDefinition
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.partitionIsInstance
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
 
 @OptIn(PrivateSessionConstructor::class, SessionConfiguration::class)
 internal abstract class LLFirAbstractSessionFactory(protected val project: Project) {
@@ -117,7 +121,23 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
 
             register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)
             register(DEPENDENCIES_SYMBOL_PROVIDER_QUALIFIED_KEY, dependencyProvider)
-            register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.Empty)
+            register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotationsImpl(this))
+            register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
+
+            FirSessionConfigurator(this).apply {
+                val hostConfiguration = ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration) {}
+                val scriptDefinition = module.file.findScriptDefinition()
+                    ?: error("Cannot load script definition for ${module.file.virtualFilePath}")
+
+                val extensionRegistrar = FirScriptingCompilerExtensionIdeRegistrar(
+                    hostConfiguration,
+                    scriptDefinitionSources = emptyList(),
+                    scriptDefinitions = listOf(scriptDefinition)
+                )
+
+                registerExtensions(extensionRegistrar.configure())
+                registerExtensions(FirScriptingSamWithReceiverExtensionRegistrar().configure())
+            }.configure()
 
             LLFirSessionConfigurator.configure(this)
         }

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

@@ -10,15 +10,11 @@ import com.intellij.psi.util.CachedValue
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
 import com.intellij.util.containers.CollectionFactory
-import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.*
-import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.fir.FirModuleDataImpl
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor
-import org.jetbrains.kotlin.fir.extensions.*
-import org.jetbrains.kotlin.fir.resolve.providers.*
 import org.jetbrains.kotlin.fir.session.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.JsPlatform

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/FirScriptingSamWithReceiverExtensionRegistrar.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.scripting.compiler.plugin
+
+import org.jetbrains.kotlin.builtins.functions.FunctionTypeKind
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.builder.scriptConfigurators
+import org.jetbrains.kotlin.fir.containingClassLookupTag
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.fir.extensions.extensionService
+import org.jetbrains.kotlin.fir.resolve.FirSamConversionTransformerExtension
+import org.jetbrains.kotlin.fir.resolve.createFunctionType
+import org.jetbrains.kotlin.fir.resolve.toFirRegularClassSymbol
+import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.functionTypeService
+import org.jetbrains.kotlin.scripting.compiler.plugin.services.FirScriptConfiguratorExtensionImpl
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
+
+
+class FirScriptingSamWithReceiverExtensionRegistrar() : FirExtensionRegistrar() {
+    override fun ExtensionRegistrarContext.configurePlugin() {
+        +FirScriptingSamWithReceiverExtensionRegistrar::FirScriptSamWithReceiverConventionTransformer
+    }
+
+    class FirScriptSamWithReceiverConventionTransformer(
+        session: FirSession
+    ) : FirSamConversionTransformerExtension(session) {
+
+        val knownAnnotations: Set<String> by lazy {
+            session.extensionService.scriptConfigurators.flatMapTo(mutableSetOf()) {
+                (it as? FirScriptConfiguratorExtensionImpl)?.knownAnnotationsForSamWithReceiver ?: emptySet()
+            }
+        }
+
+        override fun getCustomFunctionTypeForSamConversion(function: FirSimpleFunction): ConeLookupTagBasedType? {
+            val containingClassSymbol = function.containingClassLookupTag()?.toFirRegularClassSymbol(session) ?: return null
+            return runIf(containingClassSymbol.resolvedAnnotationClassIds.any { it.asSingleFqName().asString() in knownAnnotations }) {
+                val parameterTypes = function.valueParameters.map { it.returnTypeRef.coneType }
+                if (parameterTypes.isEmpty()) return null
+                val kind = session.functionTypeService.extractSingleSpecialKindForFunction(function.symbol) ?: FunctionTypeKind.Function
+                createFunctionType(
+                    kind,
+                    parameters = parameterTypes.subList(1, parameterTypes.size),
+                    receiverType = parameterTypes[0],
+                    rawReturnType = function.returnTypeRef.coneType
+                )
+            }
+        }
+    }
+}

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

@@ -8,16 +8,15 @@
 package org.jetbrains.kotlin.scripting.compiler.plugin.impl
 
 import com.intellij.openapi.Disposable
-import org.jetbrains.kotlin.builtins.functions.FunctionTypeKind
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.arguments.K2JVMCompilerArguments
 import org.jetbrains.kotlin.cli.common.arguments.parseCommandLineArguments
 import org.jetbrains.kotlin.cli.common.arguments.validateArguments
+import org.jetbrains.kotlin.cli.common.checkPluginsArguments
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.cli.common.reportArgumentParseProblems
 import org.jetbrains.kotlin.cli.common.setupCommonArguments
-import org.jetbrains.kotlin.cli.common.checkPluginsArguments
 import org.jetbrains.kotlin.cli.jvm.*
 import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
@@ -38,18 +37,6 @@ import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.extensions.AnnotationBasedExtension
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.builder.scriptConfigurators
-import org.jetbrains.kotlin.fir.containingClassLookupTag
-import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
-import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
-import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.resolve.FirSamConversionTransformerExtension
-import org.jetbrains.kotlin.fir.resolve.createFunctionType
-import org.jetbrains.kotlin.fir.resolve.toFirRegularClassSymbol
-import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
-import org.jetbrains.kotlin.fir.types.coneType
-import org.jetbrains.kotlin.fir.types.functionTypeService
 import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.platform.jvm.isJvm
 import org.jetbrains.kotlin.psi.KtFile
@@ -58,12 +45,10 @@ import org.jetbrains.kotlin.resolve.sam.SamWithReceiverResolver
 import org.jetbrains.kotlin.scripting.compiler.plugin.ScriptingCompilerConfigurationComponentRegistrar
 import org.jetbrains.kotlin.scripting.compiler.plugin.dependencies.ScriptsCompilationDependencies
 import org.jetbrains.kotlin.scripting.compiler.plugin.dependencies.collectScriptsCompilationDependencies
-import org.jetbrains.kotlin.scripting.compiler.plugin.services.FirScriptConfiguratorExtensionImpl
 import org.jetbrains.kotlin.scripting.configuration.ScriptingConfigurationKeys
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
 import org.jetbrains.kotlin.scripting.definitions.annotationsForSamWithReceivers
-import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import kotlin.script.experimental.api.ScriptCompilationConfiguration
 import kotlin.script.experimental.api.compilerOptions
 import kotlin.script.experimental.api.dependencies
@@ -139,39 +124,6 @@ internal class ScriptingSamWithReceiverComponentContributor(val annotations: Lis
     }
 }
 
-internal class FirScriptingSamWithReceiverExtensionRegistrar() : FirExtensionRegistrar() {
-    override fun ExtensionRegistrarContext.configurePlugin() {
-        +::FirScriptSamWithReceiverConventionTransformer
-    }
-
-    class FirScriptSamWithReceiverConventionTransformer(
-        session: FirSession
-    ) : FirSamConversionTransformerExtension(session) {
-
-        val knownAnnotations: Set<String> by lazy {
-            session.extensionService.scriptConfigurators.flatMapTo(mutableSetOf()) {
-                (it as? FirScriptConfiguratorExtensionImpl)?.knownAnnotationsForSamWithReceiver ?: emptySet()
-            }
-        }
-
-        override fun getCustomFunctionTypeForSamConversion(function: FirSimpleFunction): ConeLookupTagBasedType? {
-            val containingClassSymbol = function.containingClassLookupTag()?.toFirRegularClassSymbol(session) ?: return null
-            return runIf(containingClassSymbol.resolvedAnnotationClassIds.any { it.asSingleFqName().asString() in knownAnnotations }) {
-                val parameterTypes = function.valueParameters.map { it.returnTypeRef.coneType }
-                if (parameterTypes.isEmpty()) return null
-                val kind = session.functionTypeService.extractSingleSpecialKindForFunction(function.symbol) ?: FunctionTypeKind.Function
-                createFunctionType(
-                    kind,
-                    parameters = parameterTypes.subList(1, parameterTypes.size),
-                    receiverType = parameterTypes[0],
-                    rawReturnType = function.returnTypeRef.coneType
-                )
-            }
-        }
-    }
-
-}
-
 internal fun SharedScriptCompilationContext.applyConfigure(): SharedScriptCompilationContext = apply {
     val samWithReceiverAnnotations = baseScriptCompilationConfiguration[ScriptCompilationConfiguration.annotationsForSamWithReceivers]
     if (samWithReceiverAnnotations?.isEmpty() == false) {

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

@@ -30,7 +30,6 @@ import org.jetbrains.kotlin.scripting.compiler.plugin.definitions.CliScriptRepor
 import org.jetbrains.kotlin.scripting.compiler.plugin.extensions.JvmStandardReplFactoryExtension
 import org.jetbrains.kotlin.scripting.compiler.plugin.extensions.ScriptingCollectAdditionalSourcesExtension
 import org.jetbrains.kotlin.scripting.compiler.plugin.extensions.ScriptingProcessSourcesBeforeCompilingExtension
-import org.jetbrains.kotlin.scripting.compiler.plugin.impl.FirScriptingSamWithReceiverExtensionRegistrar
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionProvider
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
 import org.jetbrains.kotlin.scripting.extensions.ScriptExtraImportsProviderExtension