K2 Scripting: add support for imported scripts

2023-06-19 10:13:03 +02:00
parent d24fc3b581
commit 96bde033e1
21 changed files with 306 additions and 38 deletions
@@ -24,6 +24,7 @@ internal class LLFirBuiltinsAndCloneableSessionProvider(override val symbolProvi
    override fun getFirClassifierContainerFileIfAny(fqName: ClassId): FirFile? = null
    override fun getFirCallableContainerFile(symbol: FirCallableSymbol<*>): FirFile? = null
    override fun getFirScriptContainerFile(symbol: FirScriptSymbol): FirFile? = null
+    override fun getFirScriptByFilePath(path: String): FirScriptSymbol? = null
    override fun getFirFilesByPackage(fqName: FqName): List<FirFile> = emptyList()

    override fun getClassNamesInPackage(fqName: FqName): Set<Name> = shouldNotBeCalled()
@@ -29,6 +29,7 @@ internal class LLFirLibrarySessionProvider(
    override fun getFirClassifierContainerFileIfAny(fqName: ClassId): FirFile? = null
    override fun getFirCallableContainerFile(symbol: FirCallableSymbol<*>): FirFile? = null
    override fun getFirScriptContainerFile(symbol: FirScriptSymbol): FirFile? = null
+    override fun getFirScriptByFilePath(path: String): FirScriptSymbol? = null
    override fun getFirFilesByPackage(fqName: FqName): List<FirFile> = emptyList()

    override fun getClassNamesInPackage(fqName: FqName): Set<Name> = shouldNotBeCalled()
@@ -89,6 +89,9 @@ internal class LLFirProvider(
        return moduleComponents.cache.getContainerFirFile(symbol.fir)
    }

+    // TODO: implement
+    override fun getFirScriptByFilePath(path: String): FirScriptSymbol? = null
+
    override fun getFirFilesByPackage(fqName: FqName): List<FirFile> = error("Should not be called in FIR IDE")

    override fun getClassNamesInPackage(fqName: FqName): Set<Name> = providerHelper.getTopLevelClassNamesInPackage(fqName)
@@ -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.backend.Fir2IrScriptConfiguratorExtension
 import org.jetbrains.kotlin.fir.builder.FirScriptConfiguratorExtension
 import org.jetbrains.kotlin.fir.resolve.FirSamConversionTransformerExtension
 import kotlin.reflect.KClass
@@ -31,6 +32,7 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
            FirSamConversionTransformerExtension::class,
            FirAssignExpressionAltererExtension::class,
            FirScriptConfiguratorExtension::class,
+            Fir2IrScriptConfiguratorExtension::class,
            FirFunctionTypeKindExtension::class,
            FirDeclarationsForMetadataProviderExtension::class,
        )
@@ -96,6 +98,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
            registerExtension(FirScriptConfiguratorExtension::class, this)
        }

+        @JvmName("plusFir2IrScriptConfiguratorExtension")
+        operator fun (Fir2IrScriptConfiguratorExtension.Factory).unaryPlus() {
+            registerExtension(Fir2IrScriptConfiguratorExtension::class, this)
+        }
+
        @JvmName("plusFunctionTypeKindExtension")
        operator fun (FirFunctionTypeKindExtension.Factory).unaryPlus() {
            registerExtension(FirFunctionTypeKindExtension::class, this)
@@ -158,6 +165,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
            FirScriptConfiguratorExtension.Factory { this.invoke(it) }.unaryPlus()
        }

+        @JvmName("plusFir2IrScriptConfiguratorExtension")
+        operator fun ((FirSession) -> Fir2IrScriptConfiguratorExtension).unaryPlus() {
+            Fir2IrScriptConfiguratorExtension.Factory { this.invoke(it) }.unaryPlus()
+        }
+
        @JvmName("plusFunctionTypeKindExtension")
        operator fun ((FirSession) -> FirFunctionTypeKindExtension).unaryPlus() {
            FirFunctionTypeKindExtension.Factory { this.invoke(it) }.unaryPlus()
@@ -0,0 +1,35 @@
+/*
+ * 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.backend
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirScript
+import org.jetbrains.kotlin.fir.extensions.FirExtension
+import org.jetbrains.kotlin.fir.extensions.FirExtensionPointName
+import org.jetbrains.kotlin.fir.extensions.FirExtensionService
+import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
+import org.jetbrains.kotlin.ir.declarations.IrScript
+import org.jetbrains.kotlin.ir.symbols.IrScriptSymbol
+import kotlin.reflect.KClass
+
+abstract class Fir2IrScriptConfiguratorExtension(
+    session: FirSession,
+) : FirExtension(session) {
+    companion object {
+        val NAME = FirExtensionPointName("Fir2IrScriptConversion")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension> = Fir2IrScriptConfiguratorExtension::class
+
+    fun interface Factory : FirExtension.Factory<Fir2IrScriptConfiguratorExtension>
+
+    abstract fun IrScript.configure(script: FirScript, getIrScriptByFirSymbol: (FirScriptSymbol) -> IrScriptSymbol?)
+}
+
+val FirExtensionService.fir2IrScriptConfigurators: List<Fir2IrScriptConfiguratorExtension> by FirExtensionService.registeredExtensions()
@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirContractCallBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.expressions.impl.FirUnitExpression
+import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.references.*
 import org.jetbrains.kotlin.fir.resolve.isIteratorNext
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -304,6 +305,11 @@ class Fir2IrVisitor(
                    irScript.statements.add(irStatement!!)
                }
            }
+            for (configurator in session.extensionService.fir2IrScriptConfigurators) {
+                with(configurator) {
+                    irScript.configure(script) { declarationStorage.getCachedIrScript(it.fir)?.symbol }
+                }
+            }
            declarationStorage.leaveScope(irScript)
        }
    }
@@ -8,9 +8,7 @@ package org.jetbrains.kotlin.fir.resolve.providers
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
@@ -42,6 +40,8 @@ abstract class FirProvider : FirSessionComponent {

    abstract fun getFirScriptContainerFile(symbol: FirScriptSymbol): FirFile?

+    abstract fun getFirScriptByFilePath(path: String): FirScriptSymbol?
+
    abstract fun getFirFilesByPackage(fqName: FqName): List<FirFile>

    abstract fun getClassNamesInPackage(fqName: FqName): Set<Name>
@@ -26,6 +26,7 @@ class FirLibrarySessionProvider(
    override fun getFirClassifierContainerFileIfAny(fqName: ClassId): FirFile? = null
    override fun getFirCallableContainerFile(symbol: FirCallableSymbol<*>): FirFile? = null
    override fun getFirScriptContainerFile(symbol: FirScriptSymbol): FirFile? = null
+    override fun getFirScriptByFilePath(path: String): FirScriptSymbol? = null
    override fun getFirFilesByPackage(fqName: FqName): List<FirFile> = emptyList()

    override fun getClassNamesInPackage(fqName: FqName): Set<Name> = shouldNotBeCalled()
@@ -39,6 +39,10 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
        return state.scriptContainerMap[symbol]
    }

+    override fun getFirScriptByFilePath(path: String): FirScriptSymbol? {
+        return state.scriptByFilePathMap[path]
+    }
+
    override fun getFirClassifierContainerFile(fqName: ClassId): FirFile {
        return state.classifierContainerFileMap[fqName] ?: error("Couldn't find container for $fqName")
    }
@@ -184,6 +188,7 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
        override fun visitScript(script: FirScript, data: FirRecorderData) {
            val symbol = script.symbol
            data.state.scriptContainerMap[symbol] = data.file
+            data.file.sourceFile?.path?.let { data.state.scriptByFilePathMap[it] = symbol }
            script.acceptChildren(this, data)
        }
    }
@@ -202,6 +207,7 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
        val constructorMap = mutableMapOf<CallableId, List<FirConstructorSymbol>>()
        val callableContainerMap = mutableMapOf<FirCallableSymbol<*>, FirFile>()
        val scriptContainerMap = mutableMapOf<FirScriptSymbol, FirFile>()
+        val scriptByFilePathMap = mutableMapOf<String, FirScriptSymbol>()

        fun setFrom(other: State) {
            fileMap.clear()
@@ -213,6 +219,7 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
            constructorMap.clear()
            callableContainerMap.clear()
            scriptContainerMap.clear()
+            scriptByFilePathMap.clear()

            fileMap.putAll(other.fileMap)
            allSubPackages.addAll(other.allSubPackages)
@@ -223,6 +230,7 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
            constructorMap.putAll(other.constructorMap)
            callableContainerMap.putAll(other.callableContainerMap)
            scriptContainerMap.putAll(other.scriptContainerMap)
+            scriptByFilePathMap.putAll(other.scriptByFilePathMap)
            classesInPackage.putAll(other.classesInPackage)
            classifierInPackage.putAll(other.classifierInPackage)
        }
@@ -45,6 +45,7 @@ import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmBackendErrors
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
+import org.jetbrains.kotlin.utils.topologicalSort

 internal val scriptsToClassesPhase = makeCustomPhase<JvmBackendContext, IrModuleFragment>(
    name = "ScriptsToClasses",
@@ -58,7 +59,8 @@ internal val scriptsToClassesPhase = makeCustomPhase<JvmBackendContext, IrModule
 private class ScriptsToClassesLowering(val context: JvmBackendContext, val innerClassesSupport: JvmInnerClassesSupport) {

    fun lower(module: IrModuleFragment) {
-        val scriptsToClasses = mutableMapOf<IrScript, IrClass>()
+        val scripts = mutableListOf<IrScript>()
+        val scriptDependencies = mutableMapOf<IrScript, List<IrScript>>()

        for (irFile in module.files) {
            val iterator = irFile.declarations.listIterator()
@@ -66,7 +68,10 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
                val declaration = iterator.next()
                if (declaration is IrScript) {
                    val scriptClass = prepareScriptClass(irFile, declaration)
-                    scriptsToClasses[declaration] = scriptClass
+                    scripts.add(declaration)
+                    declaration.importedScripts.takeUnless { it.isNullOrEmpty() }?.let {
+                        scriptDependencies[declaration] = it.map { it.owner }
+                    }
                    iterator.set(scriptClass)
                }
            }
@@ -74,10 +79,11 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner

        val symbolRemapper = ScriptsToClassesSymbolRemapper()

-        for ((irScript, irScriptClass) in scriptsToClasses) {
-            finalizeScriptClass(irScriptClass, irScript, symbolRemapper)
+        val orderedScripts = topologicalSort(scripts) { scriptDependencies[this] ?: emptyList() }.reversed()
+        for (irScript in orderedScripts) {
+            finalizeScriptClass(irScript, symbolRemapper)
            // TODO fix parents in script classes
-            irScriptClass.patchDeclarationParents(irScript.parent)
+            irScript.targetClass!!.owner.patchDeclarationParents(irScript.parent)
        }
    }

@@ -111,6 +117,9 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
            it.addIfNotNull(irScript.thisReceiver?.type)
        }
        irScript.earlierScripts?.forEach { scriptsReceivers.addIfNotNull(it.owner.thisReceiver?.type) }
+        irScript.importedScripts?.forEach {
+            scriptsReceivers.add(it.owner.targetClass!!.owner.thisReceiver!!.type)
+        }
        irScript.implicitReceiversParameters.forEach {
            scriptsReceivers.add(it.type)
            scriptsReceivers.add(typeRemapper.remapType(it.type))
@@ -165,7 +174,8 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
        return capturingClasses
    }

-    private fun finalizeScriptClass(irScriptClass: IrClass, irScript: IrScript, symbolRemapper: ScriptsToClassesSymbolRemapper) {
+    private fun finalizeScriptClass(irScript: IrScript, symbolRemapper: ScriptsToClassesSymbolRemapper) {
+        val irScriptClass = irScript.targetClass!!.owner
        val typeRemapper = SimpleTypeRemapper(symbolRemapper)
        val capturingClasses = collectCapturingClasses(irScript, typeRemapper)

@@ -188,22 +198,44 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
        }

        val implicitReceiversFieldsWithParameters = arrayListOf<Pair<IrField, IrValueParameter>>().apply {
-            irScript.implicitReceiversParameters.forEach { param ->
+
+            fun createField(name: Name, type: IrType): IrField {
                val field = irScriptClass.factory.createField(
                    startOffset = UNDEFINED_OFFSET,
                    endOffset = UNDEFINED_OFFSET,
                    origin = IrDeclarationOrigin.SCRIPT_IMPLICIT_RECEIVER,
-                    name = Name.identifier("\$\$implicitReceiver_${param.type.classFqName?.shortName()?.asString()!!}"),
+                    name = name,
                    visibility = DescriptorVisibilities.PRIVATE,
                    symbol = IrFieldSymbolImpl(),
-                    type = typeRemapper.remapType(param.type),
+                    type = typeRemapper.remapType(type),
                    isFinal = true,
                    isStatic = false,
                    isExternal = false
                )
                field.parent = irScriptClass
                irScriptClass.declarations.add(field)
-                add(field to param)
+                return field
+            }
+
+            irScript.importedScripts?.forEach {
+                val importedScriptClass = it.owner.targetClass!!.owner
+                val type = importedScriptClass.defaultType
+                val name = Name.identifier("\$\$importedScript_${type.classFqName?.shortName()?.asString()!!}")
+                val param = irScriptClass.factory.createValueParameter(
+                    UNDEFINED_OFFSET, UNDEFINED_OFFSET, IrDeclarationOrigin.SCRIPT_IMPLICIT_RECEIVER, name, type, isAssignable = false,
+                    IrValueParameterSymbolImpl(), UNDEFINED_PARAMETER_INDEX, varargElementType = null,
+                    isCrossinline = false, isNoinline = false, isHidden = false,
+                )
+                param.parent = irScriptClass
+                add(createField(name, type) to param)
+            }
+            irScript.implicitReceiversParameters.forEach { param ->
+                add(
+                    createField(
+                        Name.identifier("\$\$implicitReceiver_${param.type.classFqName?.shortName()?.asString()!!}"),
+                        param.type
+                    ) to param
+                )
            }
        }

@@ -221,19 +253,30 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner

        irScriptClass.thisReceiver = scriptTransformer.scriptClassReceiver

-        val defaultContext = ScriptToClassTransformerContext(
+        val defaultContextForDeclarations = ScriptToClassTransformerContext(
            valueParameterForScriptThis = irScriptClass.thisReceiver?.symbol,
            fieldForScriptThis = null,
            valueParameterForFieldReceiver = null,
            isInScriptConstructor = false
        )

-        fun <E : IrElement> E.patchForClass(): IrElement =
+        val defaultContextForTopLevelStatements = ScriptToClassTransformerContext(
+            valueParameterForScriptThis = irScriptClass.thisReceiver?.symbol,
+            fieldForScriptThis = null,
+            valueParameterForFieldReceiver = null,
+            isInScriptConstructor = true
+        )
+
+        fun <E : IrElement> E.patchDeclarationForClass(): IrElement =
            transform(
                scriptTransformer,
-                (this as? IrDeclaration)?.let { defaultContext.copy( topLevelDeclaration = it) } ?: defaultContext
+                (this as? IrDeclaration)?.let { defaultContextForDeclarations.copy(topLevelDeclaration = it) } ?: defaultContextForDeclarations
            ).transform(lambdaPatcher, ScriptFixLambdasTransformerContext())

+        fun <E : IrElement> E.patchTopLevelStatementForClass(): IrElement =
+            transform(scriptTransformer, defaultContextForTopLevelStatements)
+                .transform(lambdaPatcher, ScriptFixLambdasTransformerContext())
+
        val explicitParametersWithFields = irScript.explicitCallParameters.map { parameter ->
            val field = irScriptClass.addField {
                startOffset = parameter.startOffset
@@ -247,7 +290,7 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
            parameter to field
        }

-        (irScript.constructor?.patchForClass() as? IrConstructor
+        (irScript.constructor?.patchDeclarationForClass() as? IrConstructor
            ?: createConstructor(irScriptClass, irScript, implicitReceiversFieldsWithParameters)).also { constructor ->
            val explicitParamsStartIndex = if (irScript.earlierScriptsParameter == null) 0 else 1
            val explicitParameters = constructor.valueParameters.subList(
@@ -286,7 +329,7 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
                    +irSetField(
                        irGet(irScriptClass.thisReceiver!!),
                        field,
-                        irGet(correspondingParameter.patchForClass() as IrValueParameter)
+                        irGet(correspondingParameter.patchDeclarationForClass() as IrValueParameter)
                    )
                }
                +IrInstanceInitializerCallImpl(
@@ -303,11 +346,11 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
        irScript.statements.forEach { scriptStatement ->
            when (scriptStatement) {
                is IrVariable -> {
-                    val copy = scriptStatement.patchForClass() as IrVariable
+                    val copy = scriptStatement.patchDeclarationForClass() as IrVariable
                    irScriptClass.addSimplePropertyFrom(copy)
                }
                is IrDeclaration -> {
-                    val copy = scriptStatement.patchForClass() as IrDeclaration
+                    val copy = scriptStatement.patchDeclarationForClass() as IrDeclaration
                    irScriptClass.declarations.add(copy)
                    // temporary way to avoid name clashes
                    // TODO: remove as soon as main generation become an explicit configuration option
@@ -316,7 +359,7 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
                    }
                }
                else -> {
-                    val transformedStatement = scriptStatement.patchForClass() as IrStatement
+                    val transformedStatement = scriptStatement.patchTopLevelStatementForClass() as IrStatement
                    irScriptClass.addAnonymousInitializer().also { irInitializer ->
                        irInitializer.body =
                            context.createIrBuilder(irInitializer.symbol).irBlockBody {
@@ -707,10 +750,20 @@ private class ScriptToClassTransformer(
            expression.putTypeArgument(i, expression.getTypeArgument(i)?.remapType())
        }
        if (expression.dispatchReceiver == null && (expression.symbol.owner as? IrDeclaration)?.needsScriptReceiver() == true) {
+            val memberAccessTargetReceiverType = when (val callee = expression.symbol.owner) {
+                is IrFunction -> callee.dispatchReceiverParameter?.type
+                is IrProperty -> callee.getter?.dispatchReceiverParameter?.type
+                else -> null
+            }
            expression.dispatchReceiver =
-                getAccessCallForScriptInstance(
-                    data, expression.startOffset, expression.endOffset, expression.origin, originalReceiverParameter = null
-                )
+                if (memberAccessTargetReceiverType != null && memberAccessTargetReceiverType != scriptClassReceiver.type)
+                    getAccessCallForImplicitReceiver(
+                        data, expression, memberAccessTargetReceiverType, expression.origin, originalReceiverParameter = null
+                    )
+                else
+                    getAccessCallForScriptInstance(
+                        data, expression.startOffset, expression.endOffset, expression.origin, originalReceiverParameter = null
+                    )
        }
        return super.visitMemberAccess(expression, data) as IrExpression
    }
@@ -909,7 +962,17 @@ private class ScriptToClassTransformer(
    }

    private fun IrDeclaration.needsScriptReceiver() =
-        (this as? IrFunction)?.dispatchReceiverParameter?.origin == IrDeclarationOrigin.SCRIPT_THIS_RECEIVER
+        when (this) {
+            is IrFunction -> this.dispatchReceiverParameter
+            is IrProperty -> {
+                this.getter?.takeIf {
+                    // without this exception, the PropertyReferenceLowering generates clinit with attempt to use script as receiver
+                    // TODO: find whether it is a valid exception and maybe how to make it more obvious
+                    it.origin != IrDeclarationOrigin.DELEGATED_PROPERTY_ACCESSOR
+                }?.dispatchReceiverParameter
+            }
+            else -> null
+        }?.origin == IrDeclarationOrigin.SCRIPT_THIS_RECEIVER
 }

 private class ScriptFixLambdasTransformer(val irScriptClass: IrClass) : IrElementTransformer<ScriptFixLambdasTransformerContext> {
@@ -994,4 +1057,4 @@ private inline fun IrClass.addAnonymousInitializer(builder: IrFunctionBuilder.()
 private val IrScript.needsReceiverProcessing: Boolean
    // in K2 we need to add dispatch receiver to the top-level declarations, and in all cases receivers should be replaced
    // for all kinds of implicit receivers
-    get() = origin == SCRIPT_K2_ORIGIN || earlierScripts?.isNotEmpty() == true || implicitReceiversParameters.isNotEmpty()
+    get() = origin == SCRIPT_K2_ORIGIN || importedScripts?.isNotEmpty() == true || earlierScripts?.isNotEmpty() == true || implicitReceiversParameters.isNotEmpty()
@@ -43,6 +43,8 @@ abstract class IrScript : IrDeclarationBase(), IrDeclarationWithName, IrDeclarat

    abstract var earlierScriptsParameter: IrValueParameter?

+    abstract var importedScripts: List<IrScriptSymbol>?
+
    abstract var earlierScripts: List<IrScriptSymbol>?

    abstract var targetClass: IrClassSymbol?
@@ -53,6 +53,7 @@ class IrScriptImpl(
    override lateinit var providedPropertiesParameters: List<IrValueParameter>
    override var resultProperty: IrPropertySymbol? = null
    override var earlierScriptsParameter: IrValueParameter? = null
+    override var importedScripts: List<IrScriptSymbol>? = null
    override var earlierScripts: List<IrScriptSymbol>? = null
    override var targetClass: IrClassSymbol? = null
    override var constructor: IrConstructor? = null
@@ -131,6 +131,7 @@ open class DeepCopyIrTreeWithSymbols(
        ).also { scriptCopy ->
            scriptCopy.thisReceiver = declaration.thisReceiver?.transform()
            declaration.statements.mapTo(scriptCopy.statements) { it.transform() }
+            scriptCopy.importedScripts = declaration.importedScripts
            scriptCopy.earlierScripts = declaration.earlierScripts
            scriptCopy.earlierScriptsParameter = declaration.earlierScriptsParameter
            scriptCopy.explicitCallParameters = declaration.explicitCallParameters.memoryOptimizedMap { it.transform() }
@@ -548,6 +548,7 @@ object IrTree : AbstractTreeBuilder() {
        +listField("providedPropertiesParameters", valueParameter, mutability = Var, isChild = true)
        +field("resultProperty", propertySymbolType, nullable = true)
        +field("earlierScriptsParameter", valueParameter, nullable = true, isChild = true)
+        +listField("importedScripts", scriptSymbolType, mutability = Var, nullable = true)
        +listField("earlierScripts", scriptSymbolType, mutability = Var, nullable = true)
        +field("targetClass", classSymbolType, nullable = true)
        +field("constructor", constructor, nullable = true) // K1
@@ -42,7 +42,7 @@ class CachingTest : TestCase() {
    }

    @Test
-    fun testSimpleImportWithMemoryCache() = expectTestToFailOnK2 {
+    fun testSimpleImportWithMemoryCache() {
        val cache = SimpleMemoryScriptsCache()
        checkWithCache(
            cache, scriptWithImport, scriptWithImportExpectedOutput,
@@ -62,7 +62,7 @@ class CachingTest : TestCase() {
    }

    @Test
-    fun testSimpleImportWithFileCache() = expectTestToFailOnK2 {
+    fun testSimpleImportWithFileCache() {
        withTempDir("scriptingTestCache") { cacheDir ->
            val cache = FileBasedScriptCache(cacheDir)
            Assert.assertEquals(true, cache.baseDir.listFiles()?.isEmpty())
@@ -89,7 +89,7 @@ class CachingTest : TestCase() {
    }

    @Test
-    fun testSimpleImportWithJarCache() = expectTestToFailOnK2 {
+    fun testSimpleImportWithJarCache() {
        withTempDir("scriptingTestJarCache") { cacheDir ->
            val cache = TestCompiledScriptJarsCache(cacheDir)
            Assert.assertTrue(cache.baseDir.listFiles()!!.isEmpty())
@@ -191,7 +191,7 @@ class ScriptingHostTest : TestCase() {
    }

    @Test
-    fun testSimpleImport() = expectTestToFailOnK2 {
+    fun testSimpleImport() {
        val greeting = listOf("Hello from helloWithVal script!", "Hello from imported helloWithVal script!")
        val script = "println(\"Hello from imported \$helloScriptName script!\")"
        val compilationConfiguration = createJvmCompilationConfigurationFromTemplate<SimpleScriptTemplate> {
@@ -204,7 +204,7 @@ class ScriptingHostTest : TestCase() {
    }

    @Test
-    fun testSimpleImportWithImplicitReceiver() = expectTestToFailOnK2 {
+    fun testSimpleImportWithImplicitReceiver() {
        val greeting = listOf("Hello from helloWithVal script!", "Hello from imported helloWithVal script!")
        val script = "println(\"Hello from imported \$helloScriptName script!\")"
        val definition = createJvmScriptDefinitionFromTemplate<SimpleScriptTemplate>(
@@ -224,6 +224,27 @@ class ScriptingHostTest : TestCase() {
        Assert.assertEquals(greeting, output)
    }

+    @Test
+    fun testSimpleImportWithImplicitReceiverRef() {
+        val greeting = listOf("Hello from helloWithVal script!", "Hello from imported helloWithVal script!")
+        val script = "println(\"Hello from imported \${(::helloScriptName).get()} script!\")"
+        val definition = createJvmScriptDefinitionFromTemplate<SimpleScriptTemplate>(
+            compilation = {
+                makeSimpleConfigurationWithTestImport()
+                implicitReceivers(String::class)
+            },
+            evaluation = {
+                implicitReceivers("abc")
+            }
+        )
+        val output = captureOut {
+            BasicJvmScriptingHost().eval(
+                script.toScriptSource(), definition.compilationConfiguration, definition.evaluationConfiguration
+            ).throwOnFailure()
+        }.lines()
+        Assert.assertEquals(greeting, output)
+    }
+
    @Test
    fun testSimpleScriptWithImplicitReceiver() {
        val greeting = listOf("3")
@@ -295,14 +316,14 @@ class ScriptingHostTest : TestCase() {
    }

    @Test
-    fun testDiamondImportWithoutSharing() = expectTestToFailOnK2 {
+    fun testDiamondImportWithoutSharing() {
        val greeting = listOf("Hi from common", "Hi from middle", "Hi from common", "sharedVar == 3")
        val output = doDiamondImportTest()
        Assert.assertEquals(greeting, output)
    }

    @Test
-    fun testDiamondImportWithSharing() = expectTestToFailOnK2 {
+    fun testDiamondImportWithSharing() {
        val greeting = listOf("Hi from common", "Hi from middle", "sharedVar == 5")
        val output = doDiamondImportTest(
            ScriptEvaluationConfiguration {
@@ -7,6 +7,7 @@ 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.Fir2IrScriptConfiguratorExtensionImpl
 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
@@ -27,5 +28,6 @@ class FirScriptingCompilerExtensionRegistrar(
        }

        +FirScriptConfiguratorExtensionImpl.getFactory(hostConfiguration)
+        +Fir2IrScriptConfiguratorExtensionImpl.getFactory(hostConfiguration)
    }
 }
@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 import org.jetbrains.kotlin.fir.pipeline.*
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
@@ -39,6 +40,9 @@ import org.jetbrains.kotlin.scripting.compiler.plugin.dependencies.ScriptsCompil
 import org.jetbrains.kotlin.scripting.compiler.plugin.services.scriptDefinitionProviderService
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionProvider
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
+import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
+import org.jetbrains.kotlin.scripting.resolve.resolvedImportScripts
+import org.jetbrains.kotlin.utils.topologicalSort
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.ScriptingHostConfiguration
 import kotlin.script.experimental.jvm.JvmDependency
@@ -370,14 +374,49 @@ private fun doCompileWithK2(
        }
    ).single().session

-    session.scriptDefinitionProviderService?.run {
+    val scriptDefinitionProviderService = session.scriptDefinitionProviderService
+
+    scriptDefinitionProviderService?.run {
        definitionProvider = ScriptDefinitionProvider.getInstance(context.environment.project)
        configurationProvider = ScriptDependenciesProvider.getInstance(context.environment.project)
    }

-    val rawFir = session.buildFirFromKtFiles(sourceFiles)
+    val rawFir = session.buildFirFromKtFiles(sourceFiles) //.reversed()

-    val (scopeSession, fir) = session.runResolution(rawFir)
+    val orderedRawFir =
+        if (scriptDefinitionProviderService == null) rawFir
+        else {
+            val rawFirDeps = rawFir.associateWith { firFile ->
+                ((firFile.sourceFile as? KtPsiSourceFile)?.psiFile as? KtFile)?.let { ktFile ->
+                    val scriptCompilationConfiguration = scriptDefinitionProviderService.configurationProvider?.getScriptConfiguration(ktFile)?.configuration
+                    scriptCompilationConfiguration?.get(ScriptCompilationConfiguration.resolvedImportScripts)?.mapNotNull { depSource ->
+                        (depSource as? VirtualFileScriptSource)?.virtualFile?.let { depVFile ->
+                            rawFir.find { ((it.sourceFile as? KtPsiSourceFile)?.psiFile as? KtFile)?.virtualFile == depVFile }
+                        }
+                    }
+                }.orEmpty()
+            }
+
+            class CycleDetected(val node: FirFile) : Throwable()
+
+            try {
+                topologicalSort(
+                    rawFir, reportCycle = { throw CycleDetected(it) }
+                ) {
+                    rawFirDeps[this] ?: emptyList()
+                }.reversed()
+            } catch (e: CycleDetected) {
+                return ResultWithDiagnostics.Failure(
+                    ScriptDiagnostic(
+                        ScriptDiagnostic.unspecifiedError,
+                        "Unable to handle recursive script dependencies, cycle detected on file ${e.node.name}",
+                        sourcePath = e.node.sourceFile?.path
+                    )
+                )
+            }
+        }
+
+    val (scopeSession, fir) = session.runResolution(orderedRawFir)
    // checkers
    session.runCheckers(scopeSession, fir, diagnosticsReporter)

@@ -0,0 +1,44 @@
+/*
+ * 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.services
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.backend.Fir2IrScriptConfiguratorExtension
+import org.jetbrains.kotlin.fir.backend.Fir2IrScriptConfiguratorExtension.Factory
+import org.jetbrains.kotlin.fir.declarations.FirScript
+import org.jetbrains.kotlin.fir.resolve.providers.firProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
+import org.jetbrains.kotlin.ir.declarations.IrScript
+import org.jetbrains.kotlin.ir.symbols.IrScriptSymbol
+import org.jetbrains.kotlin.scripting.resolve.resolvedImportScripts
+import kotlin.script.experimental.api.ScriptCompilationConfiguration
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+
+class Fir2IrScriptConfiguratorExtensionImpl(
+    session: FirSession,
+    @Suppress("UNUSED_PARAMETER") hostConfiguration: ScriptingHostConfiguration
+) : Fir2IrScriptConfiguratorExtension(session) {
+    override fun IrScript.configure(script: FirScript, getIrScriptByFirSymbol: (FirScriptSymbol) -> IrScriptSymbol?) {
+        // processing only refined scripts here
+        val scriptFile = session.firProvider.getFirScriptContainerFile(script.symbol) ?: return
+        val scriptSourceFile = scriptFile.sourceFile?.toSourceCode() ?: return
+        val compilationConfiguration = session.getScriptCompilationConfiguration(scriptSourceFile, getDefault = { null }) ?: return
+
+        // assuming that if the script is compiled, the import files should be all resolved already
+        val importedScripts = compilationConfiguration[ScriptCompilationConfiguration.resolvedImportScripts]?.takeIf { it.isNotEmpty() } ?: return
+        val importedScriptSymbols = importedScripts.mapNotNull {
+            session.firProvider.getFirScriptByFilePath(it.locationId!!)  // TODO: all !! should be converted to diagnostics
+        }
+
+        this.importedScripts = importedScriptSymbols.map { getIrScriptByFirSymbol(it)!! }.takeIf { it.isNotEmpty() }
+    }
+
+    companion object {
+        fun getFactory(hostConfiguration: ScriptingHostConfiguration): Factory {
+            return Factory { session -> Fir2IrScriptConfiguratorExtensionImpl(session, hostConfiguration) }
+        }
+    }
+}
@@ -0,0 +1,29 @@
+/*
+ * 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.services
+
+import org.jetbrains.kotlin.*
+import org.jetbrains.kotlin.fir.FirSession
+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.ScriptCompilationConfiguration
+import kotlin.script.experimental.api.SourceCode
+import kotlin.script.experimental.api.valueOrNull
+import kotlin.script.experimental.host.FileScriptSource
+import kotlin.script.experimental.host.StringScriptSource
+
+fun FirSession.getScriptCompilationConfiguration(
+    sourceCode: SourceCode?,
+    getDefault: FirScriptDefinitionProviderService.() -> ScriptCompilationConfiguration? = { definitionProvider?.getDefaultDefinition()?.compilationConfiguration }
+) =
+    scriptDefinitionProviderService?.let { providerService ->
+        sourceCode?.let { script ->
+            val ktFile = (script as? KtFileScriptSource)?.ktFile ?: error("only PSI scripts are supported at the moment")
+            providerService.configurationProvider?.getScriptConfigurationResult(ktFile)?.valueOrNull()?.configuration
+                ?: providerService.getDefault()
+        } ?: providerService.getDefault()
+    }
@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.script.loadScriptingPlugin
 import org.jetbrains.kotlin.scripting.compiler.plugin.TestDisposable
-import org.jetbrains.kotlin.scripting.compiler.plugin.expectTestToFailOnK2
 import org.jetbrains.kotlin.scripting.compiler.plugin.impl.ScriptJvmCompilerFromEnvironment
 import org.jetbrains.kotlin.scripting.compiler.plugin.updateWithBaseCompilerArguments
 import org.jetbrains.kotlin.scripting.configuration.ScriptingConfigurationKeys
@@ -24,7 +23,6 @@ import org.jetbrains.kotlin.test.TestJdkKind
 import org.junit.Assert
 import java.io.File
 import java.nio.file.Files
-import kotlin.io.path.*
 import kotlin.script.experimental.annotations.KotlinScript
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.ScriptingHostConfiguration
@@ -38,7 +36,7 @@ private const val testDataPath = "plugins/scripting/scripting-compiler/testData/
 class CompileTimeFibonacciTest : TestCase() {
    private val testRootDisposable: Disposable = TestDisposable()

-    fun testFibonacciWithSupportedNumbersImplementsTheCorrectConstants() = expectTestToFailOnK2 {
+    fun testFibonacciWithSupportedNumbersImplementsTheCorrectConstants() {
        val outputLines = runScript("supported.fib.kts")
            .valueOr { failure ->
                val message = failure.reports.joinToString("\n") { it.message }