Fix handling of lambdas in top-level destructuring declarations

2021-03-08 23:46:37 +01:00
parent 22f1814911
commit 7d07010785
5 changed files with 105 additions and 24 deletions
@@ -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 {
-            transform(scriptTransformer, null)
-        }
+        val lambdaPatcher = ScriptFixLambdasTransformer(irScript, irScriptClass, context)
+
+        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) {
-                valueParameter.type = typeRemapper.remapType(valueParameter.type)
-                if (valueParameter.varargElementType != null) {
-                    valueParameter.varargElementType = typeRemapper.remapType(valueParameter.varargElementType!!)
-                }
-                irConstructor.valueParameters = irConstructor.valueParameters + valueParameter
+            fun addConstructorParameter(valueParameter: IrValueParameter, createCorrespondingProperty: Boolean): IrValueParameter {
+                val newValueParameter = valueParameter.patchForClass() as IrValueParameter
+                irConstructor.valueParameters = irConstructor.valueParameters + newValueParameter
                if (createCorrespondingProperty) {
                    irScriptClass.addSimplePropertyFrom(
-                        valueParameter,
+                        newValueParameter,
                        IrExpressionBodyImpl(
                            IrGetValueImpl(
-                                valueParameter.startOffset, valueParameter.endOffset,
-                                valueParameter.type,
-                                valueParameter.symbol,
+                                newValueParameter.startOffset, newValueParameter.endOffset,
+                                newValueParameter.type,
+                                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() {
@@ -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,
@@ -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)
    }
 }
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR

 data class Pair(val first: Int, val second: Int)

@@ -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