[JS IR] Change logic of InteropCallableReferenceLowering a bit

Now it first transforms Lambda classes and after fixes the call sites
2021-09-21 13:08:06 +07:00
parent 693a5740c1
commit 76e08356d8
1 changed files with 122 additions and 87 deletions
@@ -21,65 +21,86 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.*
-import org.jetbrains.kotlin.ir.util.file
-import org.jetbrains.kotlin.ir.util.parentAsClass
-import org.jetbrains.kotlin.ir.util.patchDeclarationParents
-import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.classifierOrNull
+import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addToStdlib.cast

 class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLoweringPass {
-    private val newDeclarations = mutableListOf<IrDeclaration>()
-    private lateinit var implicitDeclarationFile: IrFile
-    private val transformedLambdas = mutableMapOf<IrConstructorSymbol, IrSimpleFunctionSymbol>()
+
+    override fun lower(irFile: IrFile) {
+        val ctorToFactoryMap = mutableMapOf<IrConstructorSymbol, IrSimpleFunctionSymbol>()
+        irFile.transform(CallableReferenceClassTransformer(ctorToFactoryMap), null)
+        irFile.transformChildrenVoid(object : IrElementTransformerVoid() {
+            override fun visitConstructorCall(expression: IrConstructorCall): IrExpression {
+                expression.transformChildrenVoid()
+                if (expression.origin != JsStatementOrigins.CALLABLE_REFERENCE_CREATE) return expression
+                return ctorToFactoryMap[expression.symbol]?.let { factory ->
+                    val newCall = expression.run {
+                        IrCallImpl(startOffset, endOffset, type, factory, typeArgumentsCount, valueArgumentsCount, origin)
+                    }
+
+                    newCall.dispatchReceiver = expression.dispatchReceiver
+                    newCall.extensionReceiver = expression.extensionReceiver
+
+                    for (i in 0 until expression.typeArgumentsCount) {
+                        newCall.putTypeArgument(i, expression.getTypeArgument(i))
+                    }
+
+                    for (i in 0 until expression.valueArgumentsCount) {
+                        newCall.putValueArgument(i, expression.getValueArgument(i))
+                    }
+
+                    newCall
+                } ?: expression
+            }
+        })
+    }

    override fun lower(irBody: IrBody, container: IrDeclaration) {
-        newDeclarations.clear()
-
-        implicitDeclarationFile = container.file // TODO
-        irBody.transformChildrenVoid(CallableReferenceLowerTransformer())
-        implicitDeclarationFile.declarations.addAll(newDeclarations)
+        error("Unreachable")
    }

-    inner class CallableReferenceLowerTransformer : IrElementTransformerVoid() {
-        override fun visitConstructorCall(expression: IrConstructorCall): IrExpression {
-            expression.transformChildrenVoid(this)
-            if (expression.origin === JsStatementOrigins.CALLABLE_REFERENCE_CREATE) {
-                return transformToJavaScriptFunction(expression)
+    private inner class CallableReferenceClassTransformer(private val ctorToFactoryMap: MutableMap<IrConstructorSymbol, IrSimpleFunctionSymbol>) : IrElementTransformerVoid() {
+        override fun visitFile(declaration: IrFile): IrFile {
+            declaration.transformChildrenVoid()
+            declaration.transformDeclarationsFlat { it.transformCallableReference() }
+            return declaration
+        }
+
+        override fun visitClass(declaration: IrClass): IrStatement {
+            declaration.transformChildrenVoid()
+            declaration.transformDeclarationsFlat { it.transformCallableReference() }
+            return declaration
+        }
+
+        override fun visitScript(declaration: IrScript): IrStatement {
+            declaration.transformChildrenVoid()
+            declaration.statements.transformFlat { s ->
+                if (s is IrDeclaration) s.transformCallableReference()
+                else null
            }
-            return expression
-        }
-    }
-
-    private fun transformToJavaScriptFunction(expression: IrConstructorCall): IrExpression {
-        val irConstructor = expression.symbol
-
-        // There could be more than one lambda instantiation so don't create redundant copies
-        // For testcase take a look into `boxInline/suspend/twiceRegeneratedAnonymousObject.kt`
-        val factory = transformedLambdas.getOrPut(irConstructor) {
-            buildFactoryFunction(expression).also { f ->
-                newDeclarations += f
-            }.symbol
+            return declaration
        }

-        val newCall = expression.run {
-            IrCallImpl(startOffset, endOffset, type, factory, typeArgumentsCount, valueArgumentsCount, origin)
+        private fun IrDeclaration.asCallableReference(): IrClass? {
+            if (origin == CallableReferenceLowering.Companion.FUNCTION_REFERENCE_IMPL || origin == CallableReferenceLowering.Companion.LAMBDA_IMPL)
+                return this as? IrClass
+            return null
        }

-        newCall.dispatchReceiver = expression.dispatchReceiver
-        newCall.extensionReceiver = expression.extensionReceiver
-
-        for (i in 0 until expression.typeArgumentsCount) {
-            newCall.putTypeArgument(i, expression.getTypeArgument(i))
+        private fun IrDeclaration.transformCallableReference(): List<IrDeclaration>? {
+            return asCallableReference()?.let {
+                replaceWithFactory(it)
+            }
        }

-        for (i in 0 until expression.valueArgumentsCount) {
-            newCall.putValueArgument(i, expression.getValueArgument(i))
+        private fun replaceWithFactory(lambdaClass: IrClass): List<IrDeclaration> {
+            return buildFactoryFunction(lambdaClass, ctorToFactoryMap).onEach { it.parent = lambdaClass.parent }
        }
-
-        return newCall
    }

    private fun inlineLambdaBody(
@@ -91,6 +112,8 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
        val body = invokeFun.body ?: error("invoke() method has to have a body")

        fun IrExpression.getValue(d: IrValueSymbol): IrExpression = IrGetValueImpl(startOffset, endOffset, d)
+        fun IrExpression.getCastedValue(d: IrValueSymbol, toType: IrType): IrExpression =
+            IrTypeOperatorCallImpl(startOffset, endOffset, toType, IrTypeOperator.IMPLICIT_CAST, toType, getValue(d))

        // TODO: remap type parameters???
        body.transformChildrenVoid(object : IrElementTransformerVoid() {
@@ -103,7 +126,8 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
            override fun visitGetValue(expression: IrGetValue): IrExpression {
                expression.transformChildrenVoid()
                val parameter = invokeMapping[expression.symbol] ?: return expression
-                return expression.getValue(parameter)
+                val parameterType = invokeFun.valueParameters[parameter.owner.index].type
+                return expression.getCastedValue(parameter, parameterType)
            }

            override fun visitReturn(expression: IrReturn): IrExpression {
@@ -176,18 +200,22 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
        return returnStmt.value
    }

-    private fun buildFactoryBody(factoryFunction: IrSimpleFunction, expression: IrConstructorCall): IrBlockBody {
-        val constructor = expression.symbol.owner
-        val lambdaClass = constructor.parentAsClass
+    private fun buildFactoryBody(
+        factoryFunction: IrSimpleFunction,
+        lambdaClass: IrClass,
+        newDeclarations: MutableList<IrDeclaration>
+    ): IrBlockBody {
        val invokeFun = lambdaClass.declarations.filterIsInstance<IrSimpleFunction>().single { it.name.asString() == "invoke" }
        val superInvokeFun = invokeFun.overriddenSymbols.single { it.owner.isSuspend == invokeFun.isSuspend }.owner
        val lambdaName = Name.identifier("${lambdaClass.name.asString()}\$lambda")

+        val superClass = superInvokeFun.parentAsClass
+        val anyNType = context.irBuiltIns.anyNType
        val lambdaDeclaration = context.irFactory.buildFun {
            startOffset = invokeFun.startOffset
            endOffset = invokeFun.endOffset
            // Since box/unbox is done on declaration side in case of suspend function use the specified type
-            returnType = if (invokeFun.isSuspend) invokeFun.returnType else superInvokeFun.returnType
+            returnType = if (invokeFun.isSuspend) invokeFun.returnType else anyNType
            visibility = DescriptorVisibilities.LOCAL
            name = lambdaName
            isSuspend = invokeFun.isSuspend
@@ -195,41 +223,40 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo

        lambdaDeclaration.parent = factoryFunction

-        lambdaDeclaration.valueParameters = superInvokeFun.valueParameters.map { it.copyTo(lambdaDeclaration) }
+        lambdaDeclaration.valueParameters = superInvokeFun.valueParameters.mapIndexed { id, vp ->
+            vp.copyTo(lambdaDeclaration, type = anyNType, name = invokeFun.valueParameters[id].name)
+        }

        val statements = ArrayList<IrStatement>(4)
        val isSuspendLambda = invokeFun.overriddenSymbols.any { it.owner.isSuspend }
+        val constructor = lambdaClass.declarations.firstNotNullOf { it as? IrConstructor }

        if (isSuspendLambda) {
            // Due to suspend lambda is a class itself it's not easy to inline it correctly and moreover I see no reason to do so
-            val instanceVal = JsIrBuilder.buildVar(expression.type, factoryFunction, "i").apply {
-                initializer = expression.run {
-                    val newCtorCall = IrConstructorCallImpl(
-                        startOffset,
-                        endOffset,
-                        type,
-                        symbol,
-                        typeArgumentsCount,
-                        constructorTypeArgumentsCount,
-                        valueArgumentsCount,
-                        origin
-                    )
-                    // TODO: forward type arguments
-                    assert(expression.dispatchReceiver == null)
-                    assert(expression.extensionReceiver == null)
+            val lambdaType = lambdaClass.defaultType
+            val instanceVal = JsIrBuilder.buildVar(lambdaType, factoryFunction, "i").apply {
+                val newCtorCall = IrConstructorCallImpl(
+                    lambdaClass.startOffset,
+                    lambdaClass.endOffset,
+                    lambdaType,
+                    constructor.symbol,
+                    lambdaClass.typeParameters.size,
+                    constructor.typeParameters.size,
+                    constructor.valueParameters.size
+                )

-                    for ((i, vp) in factoryFunction.valueParameters.withIndex()) {
-                        newCtorCall.putValueArgument(i, IrGetValueImpl(startOffset, endOffset, vp.type, vp.symbol))
-                    }
-
-                    newCtorCall
+                for ((i, vp) in factoryFunction.valueParameters.withIndex()) {
+                    newCtorCall.putValueArgument(i, IrGetValueImpl(startOffset, endOffset, vp.type, vp.symbol))
                }
+
+                initializer = newCtorCall
            }

            statements.add(instanceVal)

            lambdaDeclaration.body = buildLambdaBody(instanceVal, lambdaDeclaration, invokeFun)

+            newDeclarations.add(lambdaClass)
        } else {
            val fieldToParameterMapping = capturedFieldsToParametersMap(constructor, factoryFunction)
            val oldToNewInvokeParametersMapping = mutableMapOf<IrValueParameterSymbol, IrValueParameterSymbol>()
@@ -239,18 +266,17 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
            lambdaDeclaration.body =
                inlineLambdaBody(lambdaDeclaration, invokeFun, oldToNewInvokeParametersMapping, fieldToParameterMapping)

-            val classContainer = lambdaClass.parent as IrDeclarationContainer
-
            // lambdas could contain another lambdas and local classes in so let do not lose them
            val lambdaInnerClasses =
                lambdaClass.declarations.filter { it is IrClass || (it is IrSimpleFunction && it.dispatchReceiverParameter == null) }
-            classContainer.declarations.remove(lambdaClass)
-            classContainer.declarations.addAll(lambdaInnerClasses)
-            lambdaInnerClasses.forEach { it.parent = classContainer }
+
+            newDeclarations.addAll(lambdaInnerClasses)
        }

-        val functionExpression =
-            expression.run { IrFunctionExpressionImpl(startOffset, endOffset, type, lambdaDeclaration, expression.origin!!) }
+        val lambdaType = lambdaClass.superTypes.single { it.classifierOrNull === superClass.symbol }
+        val functionExpression = lambdaClass.run {
+            IrFunctionExpressionImpl(startOffset, endOffset, lambdaType, lambdaDeclaration, JsStatementOrigins.CALLABLE_REFERENCE_CREATE)
+        }

        val nameGetter = context.mapping.reflectedNameAccessor[lambdaClass]

@@ -276,32 +302,38 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
                )
            }

-            statements.add(JsIrBuilder.buildReturn(factoryFunction.symbol, JsIrBuilder.buildGetValue(tmpVar.symbol), context.irBuiltIns.nothingType))
+            statements.add(
+                JsIrBuilder.buildReturn(
+                    factoryFunction.symbol,
+                    JsIrBuilder.buildGetValue(tmpVar.symbol),
+                    context.irBuiltIns.nothingType
+                )
+            )
        } else {
            statements.add(JsIrBuilder.buildReturn(factoryFunction.symbol, functionExpression, context.irBuiltIns.nothingType))
        }

-        return context.irFactory.createBlockBody(expression.startOffset, expression.endOffset, statements)
+        return context.irFactory.createBlockBody(lambdaClass.startOffset, lambdaClass.endOffset, statements)
    }

-    private fun buildFactoryFunction(expression: IrConstructorCall): IrSimpleFunction {
-
-        val constructor = expression.symbol.owner
-        val lambdaClass = constructor.parentAsClass
+    private fun buildFactoryFunction(
+        lambdaClass: IrClass,
+        ctorToFactoryMap: MutableMap<IrConstructorSymbol, IrSimpleFunctionSymbol>
+    ): List<IrDeclaration> {
+        val newDeclarations = mutableListOf<IrDeclaration>()
+        val constructor = lambdaClass.declarations.single { it is IrConstructor } as IrConstructor

        val factoryName = Name.identifier("${lambdaClass.name.asString()}\$factory")

        val factoryDeclaration = context.irFactory.buildFun {
-            startOffset = expression.startOffset
-            endOffset = expression.endOffset
+            startOffset = lambdaClass.startOffset
+            endOffset = lambdaClass.endOffset
            visibility = lambdaClass.visibility
-            returnType = expression.type
+            returnType = lambdaClass.defaultType
            name = factoryName
            origin = JsStatementOrigins.FACTORY_ORIGIN
        }

-        factoryDeclaration.parent = implicitDeclarationFile
-
        factoryDeclaration.valueParameters = constructor.valueParameters.map { it.copyTo(factoryDeclaration) }
        factoryDeclaration.typeParameters = constructor.typeParameters.map {
            it.copyToWithoutSuperTypes(factoryDeclaration).also { tp ->
@@ -310,9 +342,12 @@ class InteropCallableReferenceLowering(val context: JsIrBackendContext) : BodyLo
            }
        }

-        factoryDeclaration.body = buildFactoryBody(factoryDeclaration, expression)
+        factoryDeclaration.body = buildFactoryBody(factoryDeclaration, lambdaClass, newDeclarations)

-        return factoryDeclaration
+        newDeclarations.add(factoryDeclaration)
+        ctorToFactoryMap[constructor.symbol] = factoryDeclaration.symbol
+
+        return newDeclarations
    }