[JS IR] Change logic of InteropCallableReferenceLowering a bit

Now it first transforms Lambda classes and after fixes the call sites

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/InteropCallableReferenceLowering.kt

@@ -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.types.IrType
-import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.ir.types.classifierOrNull
-import org.jetbrains.kotlin.ir.util.patchDeclarationParents
+import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.util.render
 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
+    override fun lower(irFile: IrFile) {
-    private val transformedLambdas = mutableMapOf<IrConstructorSymbol, IrSimpleFunctionSymbol>()
+        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()
+        error("Unreachable")
-
-        implicitDeclarationFile = container.file // TODO
-        irBody.transformChildrenVoid(CallableReferenceLowerTransformer())
-        implicitDeclarationFile.declarations.addAll(newDeclarations)
     }
 
-    inner class CallableReferenceLowerTransformer : IrElementTransformerVoid() {
+    private inner class CallableReferenceClassTransformer(private val ctorToFactoryMap: MutableMap<IrConstructorSymbol, IrSimpleFunctionSymbol>) : IrElementTransformerVoid() {
-        override fun visitConstructorCall(expression: IrConstructorCall): IrExpression {
+        override fun visitFile(declaration: IrFile): IrFile {
-            expression.transformChildrenVoid(this)
+            declaration.transformChildrenVoid()
-            if (expression.origin === JsStatementOrigins.CALLABLE_REFERENCE_CREATE) {
+            declaration.transformDeclarationsFlat { it.transformCallableReference() }
-                return transformToJavaScriptFunction(expression)
+            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
+            return declaration
-        }
-    }
-
-    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
         }
 
-        val newCall = expression.run {
+        private fun IrDeclaration.asCallableReference(): IrClass? {
-            IrCallImpl(startOffset, endOffset, type, factory, typeArgumentsCount, valueArgumentsCount, origin)
+            if (origin == CallableReferenceLowering.Companion.FUNCTION_REFERENCE_IMPL || origin == CallableReferenceLowering.Companion.LAMBDA_IMPL)
+                return this as? IrClass
+            return null
         }
 
-        newCall.dispatchReceiver = expression.dispatchReceiver
+        private fun IrDeclaration.transformCallableReference(): List<IrDeclaration>? {
-        newCall.extensionReceiver = expression.extensionReceiver
+            return asCallableReference()?.let {
-
+                replaceWithFactory(it)
-        for (i in 0 until expression.typeArgumentsCount) {
+            }
-            newCall.putTypeArgument(i, expression.getTypeArgument(i))
         }
 
-        for (i in 0 until expression.valueArgumentsCount) {
+        private fun replaceWithFactory(lambdaClass: IrClass): List<IrDeclaration> {
-            newCall.putValueArgument(i, expression.getValueArgument(i))
+            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 {
+    private fun buildFactoryBody(
-        val constructor = expression.symbol.owner
+        factoryFunction: IrSimpleFunction,
-        val lambdaClass = constructor.parentAsClass
+        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 {
+            val lambdaType = lambdaClass.defaultType
-                initializer = expression.run {
+            val instanceVal = JsIrBuilder.buildVar(lambdaType, factoryFunction, "i").apply {
-                    val newCtorCall = IrConstructorCallImpl(
+                val newCtorCall = IrConstructorCallImpl(
-                        startOffset,
+                    lambdaClass.startOffset,
-                        endOffset,
+                    lambdaClass.endOffset,
-                        type,
+                    lambdaType,
-                        symbol,
+                    constructor.symbol,
-                        typeArgumentsCount,
+                    lambdaClass.typeParameters.size,
-                        constructorTypeArgumentsCount,
+                    constructor.typeParameters.size,
-                        valueArgumentsCount,
+                    constructor.valueParameters.size
-                        origin
+                )
-                    )
-                    // TODO: forward type arguments
-                    assert(expression.dispatchReceiver == null)
-                    assert(expression.extensionReceiver == null)
 
-                    for ((i, vp) in factoryFunction.valueParameters.withIndex()) {
+                for ((i, vp) in factoryFunction.valueParameters.withIndex()) {
-                        newCtorCall.putValueArgument(i, IrGetValueImpl(startOffset, endOffset, vp.type, vp.symbol))
+                    newCtorCall.putValueArgument(i, IrGetValueImpl(startOffset, endOffset, vp.type, vp.symbol))
-                    }
-
-                    newCtorCall
                 }
+
+                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)
+            newDeclarations.addAll(lambdaInnerClasses)
-            lambdaInnerClasses.forEach { it.parent = classContainer }
         }
 
-        val functionExpression =
+        val lambdaType = lambdaClass.superTypes.single { it.classifierOrNull === superClass.symbol }
-            expression.run { IrFunctionExpressionImpl(startOffset, endOffset, type, lambdaDeclaration, expression.origin!!) }
+        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 {
+    private fun buildFactoryFunction(
-
+        lambdaClass: IrClass,
-        val constructor = expression.symbol.owner
+        ctorToFactoryMap: MutableMap<IrConstructorSymbol, IrSimpleFunctionSymbol>
-        val lambdaClass = constructor.parentAsClass
+    ): 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
+            startOffset = lambdaClass.startOffset
-            endOffset = expression.endOffset
+            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
     }