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