Fix handling of lambdas in top-level destructuring declarations

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

@@ -27,17 +27,19 @@ import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.symbols.impl.IrAnonymousInitializerSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
 internal val scriptsToClassesPhase = makeCustomPhase<JvmBackendContext, IrModuleFragment>(
     name = "ScriptsToClasses",
     description = "Put script declarations into classes",
     op = { context, input ->
         ScriptsToClassesLowering(context).lower(input)
-   }
+    }
 )
 
 
@@ -86,39 +88,40 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext) {
     private fun finalizeScriptClass(irScriptClass: IrClass, irScript: IrScript, symbolRemapper: ScriptsToClassesSymbolRemapper) {
         val typeRemapper = SimpleTypeRemapper(symbolRemapper)
         val scriptTransformer = ScriptToClassTransformer(irScript, irScriptClass, symbolRemapper, typeRemapper, context)
-        irScriptClass.thisReceiver = irScript.thisReceiver.run {
+        val lambdaPatcher = ScriptFixLambdasTransformer(irScript, irScriptClass, context)
-            transform(scriptTransformer, null)
+
-        }
+        fun <E: IrElement> E.patchForClass(): IrElement =
+            transform(scriptTransformer, null).transform(lambdaPatcher, ScriptToClassTransformerContext())
+
+        irScriptClass.thisReceiver = irScript.thisReceiver.transform(scriptTransformer, null)
 
         irScriptClass.addConstructor {
             isPrimary = true
         }.also { irConstructor ->
 
-            fun addConstructorParameter(valueParameter: IrValueParameter, createCorrespondingProperty: Boolean) {
+            fun addConstructorParameter(valueParameter: IrValueParameter, createCorrespondingProperty: Boolean): IrValueParameter {
-                valueParameter.type = typeRemapper.remapType(valueParameter.type)
+                val newValueParameter = valueParameter.patchForClass() as IrValueParameter
-                if (valueParameter.varargElementType != null) {
+                irConstructor.valueParameters = irConstructor.valueParameters + newValueParameter
-                    valueParameter.varargElementType = typeRemapper.remapType(valueParameter.varargElementType!!)
-                }
-                irConstructor.valueParameters = irConstructor.valueParameters + valueParameter
                 if (createCorrespondingProperty) {
                     irScriptClass.addSimplePropertyFrom(
-                        valueParameter,
+                        newValueParameter,
                         IrExpressionBodyImpl(
                             IrGetValueImpl(
-                                valueParameter.startOffset, valueParameter.endOffset,
+                                newValueParameter.startOffset, newValueParameter.endOffset,
-                                valueParameter.type,
+                                newValueParameter.type,
-                                valueParameter.symbol,
+                                newValueParameter.symbol,
                                 IrStatementOrigin.INITIALIZE_PROPERTY_FROM_PARAMETER
                             )
                         )
                     )
                 }
+                return newValueParameter
             }
 
             irScript.earlierScriptsParameter?.let { earlierScriptdParameter ->
                 addConstructorParameter(earlierScriptdParameter, false)
             }
-            irScript.explicitCallParameters.forEach { addConstructorParameter(it, false) }
+            val copiedExplicitParameters = irScript.explicitCallParameters.map { addConstructorParameter(it, false) }
             irScript.implicitReceiversParameters.forEach { addConstructorParameter(it, false) }
             irScript.providedProperties.forEach { addConstructorParameter(it.first, false) }
 
@@ -129,7 +132,7 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext) {
                     +irDelegatingConstructorCall(context.irBuiltIns.anyClass.owner.constructors.single())
                 } else {
                     +irDelegatingConstructorCall(baseClassCtor).also {
-                        irScript.explicitCallParameters.forEachIndexed { idx, valueParameter ->
+                        copiedExplicitParameters.forEachIndexed { idx, valueParameter ->
                             it.putValueArgument(
                                 idx,
                                 IrGetValueImpl(
@@ -151,9 +154,12 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext) {
         var hasMain = false
         irScript.statements.forEach { scriptStatement ->
             when (scriptStatement) {
-                is IrVariable -> irScriptClass.addSimplePropertyFrom(scriptStatement)
+                is IrVariable -> {
+                    val copy = scriptStatement.patchForClass() as IrVariable
+                    irScriptClass.addSimplePropertyFrom(copy)
+                }
                 is IrDeclaration -> {
-                    val copy = scriptStatement.transform(scriptTransformer, null) as IrDeclaration
+                    val copy = scriptStatement.patchForClass() as IrDeclaration
                     irScriptClass.declarations.add(copy)
                     // temporary way to avoid name clashes
                     // TODO: remove as soon as main generation become an explicit configuration option
@@ -162,7 +168,7 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext) {
                     }
                 }
                 else -> {
-                    val transformedStatement = scriptStatement.transformStatement(scriptTransformer)
+                    val transformedStatement = scriptStatement.patchForClass() as IrStatement
                     irScriptClass.addAnonymousInitializer().also { irInitializer ->
                         irInitializer.body =
                             context.createIrBuilder(irInitializer.symbol).irBlockBody {
@@ -302,6 +308,11 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext) {
         }
 }
 
+data class ScriptToClassTransformerContext(
+    val insideTopLevelDestructuringDeclaration: Boolean = false,
+    val valueParameterToReplaceWithScript: IrValueParameter? = null
+)
+
 private class ScriptToClassTransformer(
     val irScript: IrScript,
     val irScriptClass: IrClass,
@@ -370,7 +381,6 @@ private class ScriptToClassTransformer(
     override fun visitSimpleFunction(declaration: IrSimpleFunction): IrSimpleFunction = declaration.apply {
         transformParent()
         transformFunctionChildren()
-//        transformChildren(this@ScriptToClassTransformer, null)
     }
 
     override fun visitConstructor(declaration: IrConstructor): IrConstructor = declaration.apply {
@@ -498,10 +508,73 @@ private class ScriptToClassTransformer(
         }
         return super.visitCall(expression)
     }
-
-
 }
 
+private class ScriptFixLambdasTransformer(
+    val irScript: IrScript,
+    val irScriptClass: IrClass,
+    val context: JvmBackendContext
+) : IrElementTransformer<ScriptToClassTransformerContext> {
+
+    private fun unexpectedElement(element: IrElement): Nothing =
+        throw IllegalArgumentException("Unsupported element type: $element")
+
+    override fun visitElement(element: IrElement, data: ScriptToClassTransformerContext): IrElement = unexpectedElement(element)
+
+    override fun visitModuleFragment(declaration: IrModuleFragment, data: ScriptToClassTransformerContext): IrModuleFragment =
+        unexpectedElement(declaration)
+
+    override fun visitExternalPackageFragment(
+        declaration: IrExternalPackageFragment,
+        data: ScriptToClassTransformerContext
+    ): IrExternalPackageFragment =
+        unexpectedElement(declaration)
+
+    override fun visitFile(declaration: IrFile, data: ScriptToClassTransformerContext): IrFile = unexpectedElement(declaration)
+    override fun visitScript(declaration: IrScript, data: ScriptToClassTransformerContext): IrScript = unexpectedElement(declaration)
+
+    override fun visitGetValue(expression: IrGetValue, data: ScriptToClassTransformerContext): IrExpression {
+        if (data.valueParameterToReplaceWithScript == expression.symbol.owner) {
+            val newGetValue = IrGetValueImpl(
+                expression.startOffset, expression.endOffset,
+                expression.type,
+                irScriptClass.thisReceiver!!.symbol,
+                expression.origin
+            )
+            return super.visitGetValue(newGetValue, data)
+        } else return super.visitGetValue(expression, data)
+    }
+
+    override fun visitSimpleFunction(declaration: IrSimpleFunction, data: ScriptToClassTransformerContext): IrSimpleFunction =
+        with(declaration) {
+            if (data.insideTopLevelDestructuringDeclaration && origin == IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA) {
+                visibility = DescriptorVisibilities.LOCAL
+                val dataForChildren =
+                    if (dispatchReceiverParameter?.type == irScriptClass.defaultType) {
+                        val oldDispatchReceiver = dispatchReceiverParameter
+                        dispatchReceiverParameter = null
+                        data.copy(valueParameterToReplaceWithScript = oldDispatchReceiver)
+                    } else data
+                super.visitSimpleFunction(this, dataForChildren)
+            } else {
+                super.visitSimpleFunction(this, data)
+            }
+        } as IrSimpleFunction
+
+    override fun visitComposite(expression: IrComposite, data: ScriptToClassTransformerContext): IrComposite {
+        val dataForChildren =
+            if (expression.origin == IrStatementOrigin.DESTRUCTURING_DECLARATION &&
+                expression.statements.firstIsInstanceOrNull<IrDeclaration>()?.parent == irScriptClass
+            ) {
+                data.copy(insideTopLevelDestructuringDeclaration = true)
+            } else {
+                data
+            }
+        return super.visitComposite(expression, dataForChildren) as IrComposite
+    }
+}
+
+
 private class ScriptsToClassesSymbolRemapper(
     val scriptsToClasses: Map<IrScript, IrClass>
 ) : SymbolRemapper.Empty() {

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ScriptGenerator.kt

@@ -116,7 +116,13 @@ class ScriptGenerator(declarationGenerator: DeclarationGenerator) : DeclarationG
                     val type = providedProperty.type.toIrType()
                     val valueParameter = context.symbolTable.declareValueParameter(
                         startOffset, endOffset, IrDeclarationOrigin.SCRIPT_PROVIDED_PROPERTY, parameter, type
-                    )
+                    ) { symbol ->
+                        context.irFactory.createValueParameter(
+                            startOffset, endOffset, IrDeclarationOrigin.SCRIPT_PROVIDED_PROPERTY, symbol, descriptor.name,
+                            parametersIndex, type, null, isCrossinline = false, isNoinline = false, isHidden = false, isAssignable = false
+                        ).also { it.parent = irScript }
+                    }
+                    parametersIndex++
                     val irProperty =
                         PropertyGenerator(declarationGenerator).generateSyntheticProperty(
                             ktScript,

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/impl/IrScriptImpl.kt

@@ -75,6 +75,7 @@ class IrScriptImpl(
         explicitCallParameters.forEach { it.accept(visitor, data) }
         implicitReceiversParameters.forEach { it.accept(visitor, data) }
         providedProperties.forEach { it.first.accept(visitor, data) }
+        earlierScriptsParameter?.accept(visitor, data)
     }
 
     override fun <D> transformChildren(transformer: IrElementTransformer<D>, data: D) {
@@ -83,5 +84,6 @@ class IrScriptImpl(
         explicitCallParameters = explicitCallParameters.map { it.transform(transformer, data) }
         implicitReceiversParameters = implicitReceiversParameters.map { it.transform(transformer, data) }
         providedProperties = providedProperties.map { it.first.transform(transformer, data) to it.second }
+        earlierScriptsParameter = earlierScriptsParameter?.transform(transformer, data)
     }
 }

compiler/testData/codegen/script/kt22029.kts

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 
 data class Pair(val first: Int, val second: Int)
 

plugins/scripting/scripting-compiler/testData/compiler/kt42530.kts

@@ -2,4 +2,5 @@
 fun Int.isOdd() = (this % 2) == 1
 val list: List<Pair<Int, String>> = listOf(1 to "a", 2 to "b")
 val (odds, evens) = list.partition { (i, _) -> i.isOdd() }
+println(odds)
 odds