JVM_IR: optimize SAM conversion of lambda literals

Instead of materializing an invokable object and storing it in an instance of the implementation, create a specialized class that calls the lambda directly, resulting in one less nesting layer and one less class.
2019-03-26 15:17:42 +01:00
parent b74586f84e
commit 39c8939241
1 changed files with 120 additions and 45 deletions
@@ -41,7 +41,8 @@ class SingleAbstractMethodLowering(val context: CommonBackendContext) : ClassLow
    override fun lower(irClass: IrClass) {
        // TODO use `IrType`s; unfortunately, they are currently not comparable.
        //      (Can't use class symbols here, we want type parameters as well.)
-        val implementations = mutableMapOf<Pair<KotlinType, KotlinType>, IrClass>()
+        val cachedImplementations = mutableMapOf<Pair<KotlinType, KotlinType>, IrClass>()
+        val localImplementations = mutableListOf<IrClass>()
        irClass.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
            override fun visitCall(expression: IrCall): IrExpression {
                if (expression.symbol.owner.descriptor !is SamConstructorDescriptor)
@@ -55,49 +56,51 @@ class SingleAbstractMethodLowering(val context: CommonBackendContext) : ClassLow
                if (expression.operator != IrTypeOperator.SAM_CONVERSION)
                    return super.visitTypeOperator(expression)
                // Implicit conversion to an interface:
-                //    val a: (Runnable) -> Unit = { it.run() } // pretend this is a Java method
-                //    val b: () -> Unit = { print("Hello, World!") }
-                //    a(b)
+                //    // FILE: 1.java
+                //    class A { public static void f(Runnable r) { r.run() } }
+                //    // FILE: 2.kt
+                //    A.f { print("Hello, World!") }
                return createInstance(expression.typeOperand, expression.argument.transform(this, null))
            }

            // Replace a SAM-conversion of a callable object with an instance of an appropriate implementation.
            private fun createInstance(superType: IrType, invokable: IrExpression): IrExpression {
-                val implementation = implementations.getOrPut(superType.toKotlinType() to invokable.type.toKotlinType()) {
-                    createClass(superType, invokable.type)
-                }
                // Running in the class context, so empty scope means this is a field/anonymous initializer.
                val scopeOwnerSymbol = currentScope?.scope?.scopeOwnerSymbol ?: irClass.symbol
                return context.createIrBuilder(scopeOwnerSymbol).irBlock(invokable, null, superType) {
-                    // TODO if invokable is a callable reference, avoid materializing it?
-                    if (superType.isNullable() && invokable.type.isNullable()) {
-                        val invokableVariable = scope.createTemporaryVariable(invokable)
-                        val instance = irCall(implementation.constructors.single()).apply { putValueArgument(0, irGet(invokableVariable)) }
-                        +invokableVariable
-                        +irIfNull(superType, irGet(invokableVariable), irNull(), instance)
+                    // Not an exhaustive check, but enough to cover the most common cases.
+                    if (invokable is IrFunctionReference) {
+                        // Runnable(::function)
+                        createFunctionProxyInstance(superType, invokable)
+                    } else if (invokable is IrBlock && invokable.statements.last() is IrFunctionReference) {
+                        // Runnable { lambda }
+                        for (statement in invokable.statements.dropLast(1))
+                            +statement
+                        createFunctionProxyInstance(superType, invokable.statements.last() as IrFunctionReference)
                    } else {
-                        +irCall(implementation.constructors.single()).apply { putValueArgument(0, invokable) }
+                        // Fall back to materializing an invokable object.
+                        createObjectProxyInstance(superType, invokable)
                    }
                }
            }

-            // Construct a class that wraps an invokable object into an implementation of an interface:
-            //     class sam$n(private val invokable: F) : Interface() { override fun method(...) = invokable(...) }
-            private fun createClass(superType: IrType, invokableType: IrType): IrClass {
-                val invokableClass = invokableType.classifierOrFail.owner as IrClass
-                val invoke = invokableClass.functions.single { it.name == OperatorNameConventions.INVOKE }
-
+            // Construct a class that implements the specified SAM interface and contains a bunch of fields
+            // in which the arguments to the constructor are stored:
+            //     class sam$n(private val arg0: T0, ...) { override fun method(...) = <whatever buildOverride returns> }
+            private fun implement(
+                superType: IrType,
+                parameters: List<IrType>,
+                buildOverride: IrSimpleFunction.(List<IrField>) -> IrBody
+            ): IrClass {
                val superClass = superType.classifierOrFail.owner as IrClass
-                val superConstructor = if (superClass.kind != ClassKind.INTERFACE)
-                    superClass.constructors.single { it.valueParameters.size == 0 }
-                else
-                    context.irBuiltIns.anyClass.owner.constructors.single { it.valueParameters.size == 0 }
-                val superMethod = superClass.functions.single { it.modality == Modality.ABSTRACT }
+                // The language documentation prohibits casting lambdas to classes, but if it was allowed,
+                // the `irDelegatingConstructorCall` in the constructor below would need to be modified.
+                assert(superClass.kind == ClassKind.INTERFACE) { "SAM conversion to an abstract class not allowed" }

                val superClassName = superClass.fqNameSafe.pathSegments().joinToString("_") { it.toString() }
                val subclass = buildClass {
                    // TODO this is not the name some tests (e.g. kt11519) expect
-                    name = Name.identifier("sam\$${superClassName}\$${implementations.size}")
+                    name = Name.identifier("sam\$$superClassName\$${cachedImplementations.size + localImplementations.size}")
                    origin = JvmLoweredDeclarationOrigin.GENERATED_SAM_IMPLEMENTATION
                }.apply {
                    createImplicitParameterDeclarationWithWrappedDescriptor()
@@ -106,31 +109,35 @@ class SingleAbstractMethodLowering(val context: CommonBackendContext) : ClassLow
                    superTypes += superType
                }

-                val invokableField = subclass.addField {
-                    name = Name.identifier("invokable")
-                    type = invokableType
-                    origin = subclass.origin
-                    visibility = Visibilities.PRIVATE
+                val fields = parameters.mapIndexed { i, fieldType ->
+                    subclass.addField {
+                        name = Name.identifier("arg$i")
+                        type = fieldType
+                        origin = subclass.origin
+                        visibility = Visibilities.PRIVATE
+                    }
                }

                subclass.addConstructor {
-                    returnType = subclass.defaultType
                    origin = subclass.origin
                    isPrimary = true
                }.apply {
-                    addValueParameter {
-                        name = invokableField.name
-                        type = invokableType
-                        index = 0
-                        origin = subclass.origin
+                    for (field in fields) {
+                        addValueParameter {
+                            name = field.name
+                            type = field.type
+                            origin = subclass.origin
+                        }
                    }
                    body = context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
-                        +irDelegatingConstructorCall(superConstructor)
-                        +irSetField(irGet(subclass.thisReceiver!!), invokableField, irGet(valueParameters[0]))
-                        +IrInstanceInitializerCallImpl(startOffset, endOffset, superClass.symbol, context.irBuiltIns.unitType)
+                        +irDelegatingConstructorCall(context.irBuiltIns.anyClass.owner.constructors.single())
+                        for ((field, parameter) in fields.zip(valueParameters))
+                            +irSetField(irGet(subclass.thisReceiver!!), field, irGet(parameter))
+                        +IrInstanceInitializerCallImpl(startOffset, endOffset, subclass.symbol, context.irBuiltIns.unitType)
                    }
                }

+                val superMethod = superClass.functions.single { it.modality == Modality.ABSTRACT }
                subclass.addFunction {
                    name = superMethod.name
                    returnType = superMethod.returnType
@@ -140,17 +147,85 @@ class SingleAbstractMethodLowering(val context: CommonBackendContext) : ClassLow
                    overriddenSymbols.add(superMethod.symbol)
                    dispatchReceiverParameter = subclass.thisReceiver!!.copyTo(this)
                    superMethod.valueParameters.mapTo(valueParameters) { it.copyTo(this) }
-                    body = context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
-                        +irReturn(irCall(invoke.symbol).apply {
-                            dispatchReceiver = irGetField(irGet(subclass.thisReceiver!!), invokableField)
+                    body = buildOverride(fields)
+                }
+
+                return subclass
+            }
+
+            // Construct a class that wraps an invokable object into an implementation of an interface:
+            //     class sam$n(private val invokable: F) : Interface { override fun method(...) = invokable(...) }
+            private fun createObjectProxy(superType: IrType, invokableType: IrType): IrClass =
+                implement(superType, listOf(invokableType)) { fields ->
+                    context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
+                        val invokableClass = invokableType.classifierOrFail.owner as IrClass
+                        +irReturn(irCall(invokableClass.functions.single { it.name == OperatorNameConventions.INVOKE }).apply {
+                            dispatchReceiver = irGetField(irGet(dispatchReceiverParameter!!), fields[0])
                            valueParameters.forEachIndexed { i, parameter -> putValueArgument(i, irGet(parameter)) }
                        })
                    }
                }

-                return subclass
+            private fun IrBlockBuilder.createObjectProxyInstance(superType: IrType, invokable: IrExpression) {
+                val implementation = cachedImplementations.getOrPut(superType.toKotlinType() to invokable.type.toKotlinType()) {
+                    createObjectProxy(superType, invokable.type)
+                }
+                if (superType.isNullable() && invokable.type.isNullable()) {
+                    val invokableVariable = irTemporary(invokable)
+                    val instance = irCall(implementation.constructors.single()).apply { putValueArgument(0, irGet(invokableVariable)) }
+                    +irIfNull(superType, irGet(invokableVariable), irNull(), instance)
+                } else {
+                    +irCall(implementation.constructors.single()).apply { putValueArgument(0, invokable) }
+                }
+            }
+
+            // This lowering is located after LocalDeclarationsLowering. On the one hand, this allows us to cache
+            // object-wrapping classes without checking if the types are local declarations or not. On the other hand,
+            // this means we need to manually create fields for all bound arguments when wrapping callable references.
+            private val IrFunctionReference.arguments: List<IrExpression?>
+                get() = listOf(dispatchReceiver).filter { symbol.owner.dispatchReceiverParameter != null } +
+                        listOf(extensionReceiver).filter { symbol.owner.extensionReceiverParameter != null } +
+                        (0 until valueArgumentsCount).map(::getValueArgument)
+
+            // Construct a class that forwards method calls to an existing function/method:
+            //     class sam$n(private val receiver: R) : Interface { override fun method(...) = receiver.target(...) }
+            //
+            // Unlike the above variant, this avoids materializing an invokable KFunction representing
+            // the target in CallableReferenceLowering, thus producing one less class. This is actually very
+            // common, as `Interface { something }` is a local function + a SAM-conversion of a reference
+            // to it into an implementation.
+            private fun createFunctionProxy(superType: IrType, reference: IrFunctionReference): IrClass =
+                implement(superType, reference.arguments.mapNotNull { it?.type }) { fields ->
+                    context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
+                        +irReturn(irCall(reference.symbol).apply {
+                            var boundOffset = 0
+                            var offset = 0
+                            val arguments = reference.arguments.map {
+                                if (it != null)
+                                    irGetField(irGet(dispatchReceiverParameter!!), fields[boundOffset++])
+                                else
+                                    irGet(valueParameters[offset++])
+                            }.toMutableList()
+                            if (reference.symbol.owner.dispatchReceiverParameter != null) {
+                                dispatchReceiver = arguments.removeAt(0)
+                            }
+                            if (reference.symbol.owner.extensionReceiverParameter != null) {
+                                extensionReceiver = arguments.removeAt(0)
+                            }
+                            arguments.forEachIndexed(::putValueArgument)
+                        })
+                    }
+                }
+
+            private fun IrBlockBuilder.createFunctionProxyInstance(superType: IrType, reference: IrFunctionReference) {
+                val implementation = createFunctionProxy(superType, reference)
+                localImplementations += implementation
+                +irCall(implementation.constructors.single()).apply {
+                    reference.arguments.filterNotNull().forEachIndexed(::putValueArgument)
+                }
            }
        })
-        irClass.declarations.addAll(implementations.values)
+        irClass.declarations += cachedImplementations.values
+        irClass.declarations += localImplementations
    }
 }