K2 Scripting: fix capturing in scripts created from Fir

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -607,6 +607,17 @@ class Fir2IrVisitor(
                     }
                 }
             }
+        } else if (boundSymbol is FirScriptSymbol && calleeReference.contextReceiverNumber >= 0) {
+            val firScript = boundSymbol.fir
+            val irScript = declarationStorage.getCachedIrScript(firScript) ?: error("IrScript for ${firScript.name} not found")
+            val receiverParameter = irScript.implicitReceiversParameters.find { it.index == calleeReference.contextReceiverNumber }
+            if (receiverParameter != null) {
+                return thisReceiverExpression.convertWithOffsets { startOffset, endOffset ->
+                    IrGetValueImpl(startOffset, endOffset, receiverParameter.type, receiverParameter.symbol)
+                }
+            } else {
+                error("Expecting implicit receiver") // TODO: check if any valid situations possible here
+            }
         } else if (boundSymbol is FirCallableSymbol) {
             val irFunction = when (boundSymbol) {
                 is FirFunctionSymbol -> declarationStorage.getIrFunctionSymbol(boundSymbol).owner

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

@@ -283,6 +283,7 @@ class ClosureAnnotator(irElement: IrElement, declaration: IrDeclaration) {
 
         override fun visitFunctionAccess(expression: IrFunctionAccessExpression, data: ClosureBuilder?) {
             super.visitFunctionAccess(expression, data)
+            processScriptCapturing(expression.dispatchReceiver, expression.symbol.owner, data)
             processMemberAccess(expression.symbol.owner, data)
         }
 
@@ -310,6 +311,15 @@ class ClosureAnnotator(irElement: IrElement, declaration: IrDeclaration) {
             typeParameterContainerScopeBuilder?.seeType(expression.type)
         }
 
+        private fun processScriptCapturing(receiverExpression: IrExpression?, declaration: IrDeclaration, data: ClosureBuilder?) {
+            if (receiverExpression == null) {
+                val parent = declaration.parent
+                if (parent is IrScript && parent.thisReceiver != null) {
+                    data?.seeVariable(parent.thisReceiver!!.symbol)
+                }
+            }
+        }
+
         private fun processMemberAccess(declaration: IrDeclaration, parentClosure: ClosureBuilder?) {
             if (declaration.isLocal) {
                 if (declaration is IrSimpleFunction && declaration.visibility != DescriptorVisibilities.LOCAL) {

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

@@ -109,6 +109,16 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
     private fun collectCapturingClasses(irScript: IrScript, typeRemapper: SimpleTypeRemapper): Set<IrClassImpl> {
         val annotator = ClosureAnnotator(irScript, irScript)
         val capturingClasses = mutableSetOf<IrClassImpl>()
+
+        val scriptsReceivers = mutableSetOf<IrType>().also {
+            it.addIfNotNull(irScript.thisReceiver?.type)
+        }
+        irScript.earlierScripts?.forEach { scriptsReceivers.addIfNotNull(it.owner.thisReceiver?.type) }
+        irScript.implicitReceiversParameters.forEach {
+            scriptsReceivers.add(it.type)
+            scriptsReceivers.add(typeRemapper.remapType(it.type))
+        }
+
         val collector = object : IrElementVisitorVoid {
             override fun visitElement(element: IrElement) {
                 element.acceptChildrenVoid(this)
@@ -117,15 +127,6 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext, val inner
             override fun visitClass(declaration: IrClass) {
                 if (declaration is IrClassImpl && !declaration.isInner) {
                     val closure = annotator.getClassClosure(declaration)
-                    val scriptsReceivers = mutableSetOf<IrType>().also {
-                        it.addIfNotNull(irScript.thisReceiver?.type)
-                    }
-                    irScript.earlierScripts?.forEach { scriptsReceivers.addIfNotNull(it.owner.thisReceiver?.type) }
-                    irScript.implicitReceiversParameters.forEach {
-                        scriptsReceivers.add(it.type)
-                        scriptsReceivers.add(typeRemapper.remapType(it.type))
-                    }
-
                     if (closure.capturedValues.any { it.owner.type in scriptsReceivers }) {
                         fun reportError(factory: KtDiagnosticFactory1<String>, name: Name? = null) {
                             context.ktDiagnosticReporter.at(declaration).report(factory, (name ?: declaration.name).asString())
@@ -854,7 +855,7 @@ private class ScriptToClassTransformer(
             }
         } else if (irScript.needsReceiverProcessing) {
             val getValueParameter = expression.symbol.owner as? IrValueParameter
-            if (getValueParameter != null && getValueParameter.name == SpecialNames.THIS) {
+            if (getValueParameter != null && isValidNameForReceiver(getValueParameter.name)) {
                 val newExpression = getDispatchReceiverExpression(
                     data, expression, getValueParameter.type, expression.origin, getValueParameter
                 )
@@ -866,6 +867,9 @@ private class ScriptToClassTransformer(
         return super.visitGetValue(expression, data)
     }
 
+    private fun isValidNameForReceiver(name: Name) =
+        name == SpecialNames.THIS || irScript.implicitReceiversParameters.any { it.name == name }
+
     private fun IrDeclaration.isCurrentScriptTopLevelDeclaration(data: ScriptToClassTransformerContext): Boolean {
         if (data.topLevelDeclaration == null || (parent != irScript && parent != irScriptClass)) return false
         val declarationToCompare = if (this is IrFunction) this.propertyIfAccessor else this

compiler/testData/codegen/script/scriptInstanceCapturing/classCapturesExtensionIndirect.kts

@@ -1,6 +1,5 @@
 // IGNORE_BACKEND: JS, JS_IR, JS_IR_ES6, NATIVE, WASM
 // IGNORE_BACKEND: JVM
-// IGNORE_BACKEND_K2: JVM_IR
 
 // expected: rv: kotlin.Unit
 

compiler/testData/codegen/script/scriptInstanceCapturing/classCapturesExtensionIndirect2x.kts

@@ -1,6 +1,5 @@
 // IGNORE_BACKEND: JS, JS_IR, JS_IR_ES6, NATIVE, WASM
 // IGNORE_BACKEND: JVM
-// IGNORE_BACKEND_K2: JVM_IR
 
 // expected: rv: kotlin.Unit