K2 Scripting: fix capturing from the imported scripts

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt

@@ -314,7 +314,9 @@ class ClosureAnnotator(irElement: IrElement, declaration: IrDeclaration) {
         private fun processScriptCapturing(receiverExpression: IrExpression?, declaration: IrDeclaration, data: ClosureBuilder?) {
             if (receiverExpression == null) {
                 val parent = declaration.parent
-                if (parent is IrScript && parent.thisReceiver != null) {
+                if (parent is IrScript) {
+                    data?.seeVariable(parent.thisReceiver!!.symbol)
+                } else if (parent is IrClass && parent.origin == IrDeclarationOrigin.SCRIPT_CLASS) {
                     data?.seeVariable(parent.thisReceiver!!.symbol)
                 }
             }

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/ScriptLowering.kt

@@ -175,6 +175,14 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
     }
 
     private fun finalizeScriptClass(irScript: IrScript, symbolRemapper: ScriptsToClassesSymbolRemapper) {
+
+        if (irScript.thisReceiver == null) {
+            // This is a placeholder that is transformed to a proper receiver for script class down below, but it is needed for
+            // collecting captured script instances (see ClosureAnnotator.ClosureCollectorVisitor.processScriptCapturing)
+            val type = IrSimpleTypeImpl(irScript.symbol, false, emptyList(), emptyList())
+            irScript.thisReceiver = irScript.createThisReceiverParameter(context, IrDeclarationOrigin.INSTANCE_RECEIVER, type)
+        }
+
         val irScriptClass = irScript.targetClass!!.owner
         val typeRemapper = SimpleTypeRemapper(symbolRemapper)
         val capturingClasses = collectCapturingClasses(irScript, typeRemapper)
@@ -558,16 +566,10 @@ private class ScriptToClassTransformer(
         }
     }
 
-    val scriptThisType = IrSimpleTypeImpl(irScriptClass.symbol, false, emptyList(), emptyList())
     val scriptClassReceiver =
-        irScript.thisReceiver?.let {
+        irScript.thisReceiver!!.let {
-            it.type = scriptThisType
+            it.type = IrSimpleTypeImpl(irScriptClass.symbol, false, emptyList(), emptyList())
             it.transform(this, ScriptToClassTransformerContext(null, null, null, false))
-        } ?: run {
-            context.symbolTable.enterScope(irScriptClass)
-            val newReceiver = irScriptClass.createThisReceiverParameter(IrDeclarationOrigin.INSTANCE_RECEIVER, scriptThisType)
-            context.symbolTable.leaveScope(irScriptClass)
-            newReceiver
         }
 
     private fun IrDeclaration.transformParent() {
@@ -592,7 +594,7 @@ private class ScriptToClassTransformer(
             typeRemapper.withinScope(this) {
                 val newDispatchReceiverParameter = dispatchReceiverParameter?.transform(data) ?: run {
                     if (this.isCurrentScriptTopLevelDeclaration(data)) {
-                        createThisReceiverParameter(IrDeclarationOrigin.SCRIPT_THIS_RECEIVER, scriptClassReceiver.type)
+                        createThisReceiverParameter(context, IrDeclarationOrigin.SCRIPT_THIS_RECEIVER, scriptClassReceiver.type)
                     } else null
                 }
                 val isInScriptConstructor = this@transformFunctionChildren is IrConstructor && (parent == irScript || parent == irScriptClass)
@@ -621,24 +623,6 @@ private class ScriptToClassTransformer(
             }
         }
 
-    private fun IrDeclarationParent.createThisReceiverParameter(origin: IrDeclarationOrigin, type: IrType): IrValueParameter =
-        context.symbolTable.irFactory.createValueParameter(
-            startOffset = startOffset,
-            endOffset = endOffset,
-            origin = origin,
-            name = SpecialNames.THIS,
-            type = type,
-            isAssignable = false,
-            symbol = IrValueParameterSymbolImpl(),
-            index = UNDEFINED_PARAMETER_INDEX,
-            varargElementType = null,
-            isCrossinline = false,
-            isNoinline = false,
-            isHidden = false,
-        ).also {
-            it.parent = this
-        }
-
     private fun IrTypeParameter.remapSuperTypes(): IrTypeParameter = apply {
         superTypes = superTypes.map { it.remapType() }
     }
@@ -691,7 +675,7 @@ private class ScriptToClassTransformer(
     override fun visitConstructor(declaration: IrConstructor, data: ScriptToClassTransformerContext): IrConstructor = declaration.apply {
         if (declaration in capturingClassesConstructors) {
             declaration.dispatchReceiverParameter =
-                declaration.createThisReceiverParameter(IrDeclarationOrigin.INSTANCE_RECEIVER, scriptClassReceiver.type)
+                declaration.createThisReceiverParameter(context, IrDeclarationOrigin.INSTANCE_RECEIVER, scriptClassReceiver.type)
         }
         transformParent()
         transformFunctionChildren(data)
@@ -1061,4 +1045,26 @@ 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 || importedScripts?.isNotEmpty() == true || earlierScripts?.isNotEmpty() == true || implicitReceiversParameters.isNotEmpty()
+    get() = origin == SCRIPT_K2_ORIGIN || importedScripts?.isNotEmpty() == true || earlierScripts?.isNotEmpty() == true || implicitReceiversParameters.isNotEmpty()
+
+private fun IrDeclarationParent.createThisReceiverParameter(
+    context: JvmBackendContext,
+    origin: IrDeclarationOrigin,
+    type: IrType
+): IrValueParameter =
+    context.symbolTable.irFactory.createValueParameter(
+        startOffset = startOffset,
+        endOffset = endOffset,
+        origin = origin,
+        name = SpecialNames.THIS,
+        type = type,
+        isAssignable = false,
+        symbol = IrValueParameterSymbolImpl(),
+        index = UNDEFINED_PARAMETER_INDEX,
+        varargElementType = null,
+        isCrossinline = false,
+        isNoinline = false,
+        isHidden = false,
+    ).also {
+        it.parent = this
+    }

libraries/tools/kotlin-main-kts-test/test/org/jetbrains/kotlin/mainKts/test/mainKtsTest.kt

@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.mainKts.MainKtsScript
 import org.jetbrains.kotlin.mainKts.SCRIPT_FILE_LOCATION_DEFAULT_VARIABLE_NAME
 import org.jetbrains.kotlin.mainKts.impl.Directories
 import org.jetbrains.kotlin.scripting.compiler.plugin.assertTrue
-import org.jetbrains.kotlin.scripting.compiler.plugin.expectTestToFailOnK2
 import org.junit.Assert
 import org.junit.Assert.assertEquals
 import org.junit.Ignore
@@ -137,7 +136,7 @@ class MainKtsTest {
     }
 
     @Test
-    fun testImportWithCapture() = expectTestToFailOnK2 {
+    fun testImportWithCapture() {
 
         val out = captureOut {
             val res = evalFile(File("$TEST_DATA_ROOT/import-with-capture-test.main.kts"))

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

@@ -56,7 +56,7 @@ class FirScriptConfiguratorExtensionImpl(
         check(configuration != null) { "Configuration for ${sourceFile.asString()} wasn't found" }
 
         // TODO: rewrite/extract decision logic for clarity
-        configuration[ScriptCompilationConfiguration.baseClass]?.let { baseClass ->
+        configuration.getNoDefault(ScriptCompilationConfiguration.baseClass)?.let { baseClass ->
             val baseClassFqn = FqName.fromSegments(baseClass.typeName.split("."))
             contextReceivers.add(buildContextReceiverWithFqName(baseClassFqn, Name.special(SCRIPT_SPECIAL_NAME_STRING)))