JVM_IR: Change CallableReferenceLowering to match the JVM backend

In particular, the JVM backend inlines the bodies of lambdas into the
generated lambda subclasses and stores bound receivers in the receiver
field of CallableReference. Both are required for reflection.

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/DeclarationOrigins.kt

@@ -27,6 +27,7 @@ interface JvmLoweredDeclarationOrigin : IrDeclarationOrigin {
     object DEFAULT_IMPLS_BRIDGE : IrDeclarationOriginImpl("DEFAULT_IMPLS_BRIDGE")
     object DEFAULT_IMPLS_BRIDGE_TO_SYNTHETIC : IrDeclarationOriginImpl("DEFAULT_IMPLS_BRIDGE_TO_SYNTHETIC", isSynthetic = true)
     object FIELD_FOR_OUTER_THIS : IrDeclarationOriginImpl("FIELD_FOR_OUTER_THIS")
+    object LAMBDA_IMPL : IrDeclarationOriginImpl("LAMBDA_IMPL")
     object FUNCTION_REFERENCE_IMPL : IrDeclarationOriginImpl("FUNCTION_REFERENCE_IMPL", isSynthetic = true)
     object SYNTHETIC_ACCESSOR : IrDeclarationOriginImpl("SYNTHETIC_ACCESSOR", isSynthetic = true)
     object TO_ARRAY : IrDeclarationOriginImpl("TO_ARRAY")

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmSymbols.kt

@@ -136,6 +136,8 @@ class JvmSymbols(
             addValueParameter("receiver", irBuiltIns.anyNType)
         }
 
+        klass.addField("receiver", irBuiltIns.anyNType, Visibilities.PROTECTED)
+
         generateCallableReferenceMethods(klass)
     }.symbol
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt

@@ -157,6 +157,7 @@ class IrExpressionLambdaImpl(
                     when (ir) {
                         is IrGetValue -> capturedParamDesc(ir.descriptor.name.asString(), typeMapper.mapType(ir.type))
                         is IrConst<*> -> capturedParamDesc(BOUND_REFERENCE_RECEIVER, typeMapper.mapType(ir.type))
+                        is IrGetField -> capturedParamDesc(ir.descriptor.name.asString(), typeMapper.mapType(ir.type))
                         else -> error("Unrecognized expression: ${ir.dump()}")
                     }
                 )

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

@@ -1,24 +1,12 @@
 /*
- * Copyright 2010-2017 JetBrains s.r.o.
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
- *
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
  */
 
 package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
-import org.jetbrains.kotlin.backend.common.descriptors.isFunctionOrKFunctionType
 import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
@@ -32,12 +20,14 @@ import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrExpression
 import org.jetbrains.kotlin.backend.jvm.ir.isInlineParameter
 import org.jetbrains.kotlin.builtins.functions.FunctionInvokeDescriptor
 import org.jetbrains.kotlin.codegen.PropertyReferenceCodegen
+import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.addConstructor
-import org.jetbrains.kotlin.ir.builders.declarations.addField
 import org.jetbrains.kotlin.ir.builders.declarations.addFunction
+import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
 import org.jetbrains.kotlin.ir.builders.declarations.buildClass
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
@@ -45,12 +35,17 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrClassReferenceImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrVarargImpl
-import org.jetbrains.kotlin.ir.types.*
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.types.IrSimpleType
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.IrTypeProjection
+import org.jetbrains.kotlin.ir.types.typeWith
 import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.load.java.JavaVisibilities
-import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addIfNotNull
 
 internal val callableReferencePhase = makeIrFilePhase(
     ::CallableReferenceLowering,
@@ -58,125 +53,117 @@ internal val callableReferencePhase = makeIrFilePhase(
     description = "Handle callable references"
 )
 
-//Originally was copied from K/Native
+// Originally copied from K/Native
-internal class CallableReferenceLowering(val context: JvmBackendContext) : FileLoweringPass {
+internal class CallableReferenceLowering(private val context: JvmBackendContext) : FileLoweringPass, IrElementTransformerVoidWithContext() {
-    private val inlineLambdaReferences = mutableSetOf<IrFunctionReference>()
+    // This pass ignores function references used in inline arguments to inline functions references or in SAM conversions.
+    // We also implicitly ignore all suspend function references by only dealing with subclasses of (K)Function and not
+    // (K)SuspendFunction.
+    private val ignoredFunctionReferences = mutableSetOf<IrFunctionReference>()
 
-    override fun lower(irFile: IrFile) {
+    private val IrFunctionReference.isIgnored: Boolean
-        irFile.transformChildrenVoid(object : IrElementTransformerVoidWithContext() {
+        get() = !type.isFunctionOrKFunction() || ignoredFunctionReferences.contains(this)
-            override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
-                if (expression.operator == IrTypeOperator.SAM_CONVERSION) {
-                    val invokable = expression.argument
-                    if (invokable is IrFunctionReference) {
-                        inlineLambdaReferences += invokable
-                    } else if (invokable is IrBlock && invokable.statements.last() is IrFunctionReference) {
-                        inlineLambdaReferences += invokable.statements.last() as IrFunctionReference
-                    }
-                }
-                return super.visitTypeOperator(expression)
-            }
 
-            override fun visitFunctionAccess(expression: IrFunctionAccessExpression): IrExpression {
+    override fun lower(irFile: IrFile) = irFile.transformChildrenVoid(this)
-                val callee = expression.symbol.owner
-                if (callee.isInlineFunctionCall(context)) {
-                    //TODO: more wise filtering
-                    for (valueParameter in callee.valueParameters) {
-                        if (valueParameter.isInlineParameter()) {
-                            expression.getValueArgument(valueParameter.index)?.let {
-                                if (isInlineIrExpression(it)) {
-                                    inlineLambdaReferences += (it as IrBlock).statements.filterIsInstance<IrFunctionReference>()
-                                }
-                            }
-                        }
-                    }
-                }
 
-                val argumentsCount = expression.valueArgumentsCount
+    // Change calls to big arity invoke functions to vararg calls.
-                // Change calls to FunctionN with large N to varargs calls.
+    override fun visitFunctionAccess(expression: IrFunctionAccessExpression): IrExpression {
-                val newCall = if (argumentsCount >= FunctionInvokeDescriptor.Factory.BIG_ARITY &&
+        markInlineFunctionReferences(expression)
-                    callee.parentAsClass.defaultType.isFunctionOrKFunction()
+        expression.transformChildrenVoid(this)
-                ) {
-                    val vararg = IrVarargImpl(
-                        UNDEFINED_OFFSET, UNDEFINED_OFFSET,
-                        context.ir.symbols.array.typeWith(context.irBuiltIns.anyNType),
-                        context.irBuiltIns.anyNType,
-                        (0 until argumentsCount).map { i -> expression.getValueArgument(i)!! }
-                    )
-                    val invokeFun = context.ir.symbols.functionN.owner.declarations.single {
-                        it is IrSimpleFunction && it.name.asString() == "invoke"
-                    } as IrSimpleFunction
 
-                    IrCallImpl(
+        if (expression.valueArgumentsCount < FunctionInvokeDescriptor.BIG_ARITY ||
-                        UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+            !expression.symbol.owner.parentAsClass.defaultType.isFunctionOrKFunction())
-                        expression.type,
+            return expression
-                        invokeFun.symbol, invokeFun.descriptor,
-                        1,
-                        expression.origin,
-                        (expression as? IrCall)?.superQualifier?.let { context.ir.symbols.externalSymbolTable.referenceClass(it) }
-                    ).apply {
-                        putTypeArgument(0, expression.type)
-                        dispatchReceiver = expression.dispatchReceiver
-                        extensionReceiver = expression.extensionReceiver
-                        putValueArgument(0, vararg)
-                    }
-                } else expression
 
-                //TODO: clean
+        return IrCallImpl(
-                return super.visitFunctionAccess(newCall)
+            expression.startOffset, expression.endOffset,
-            }
+            expression.type, functionNInvokeFun.symbol, functionNInvokeFun.descriptor,
-
+            1, expression.origin
-            override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+        ).apply {
-                expression.transformChildrenVoid(this)
+            putTypeArgument(0, expression.type)
-
+            dispatchReceiver = expression.dispatchReceiver
-                if (!expression.type.toKotlinType().isFunctionOrKFunctionType || inlineLambdaReferences.contains(expression)) {
+            extensionReceiver = expression.extensionReceiver
-                    // Not a subject of this lowering.
+            val vararg = IrVarargImpl(
-                    return expression
+                UNDEFINED_OFFSET, UNDEFINED_OFFSET,
-                }
+                context.ir.symbols.array.typeWith(context.irBuiltIns.anyNType),
-
+                context.irBuiltIns.anyNType,
-                val currentDeclarationParent = allScopes.map { it.irElement }.last { it is IrDeclarationParent } as IrDeclarationParent
+                (0 until expression.valueArgumentsCount).map { expression.getValueArgument(it)!! }
-                val loweredFunctionReference = FunctionReferenceBuilder(currentDeclarationParent, expression).build()
+            )
-                return context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol).irBlock(expression) {
+            putValueArgument(0, vararg)
-                    +loweredFunctionReference.functionReferenceClass
+        }
-                    +irCall(loweredFunctionReference.functionReferenceConstructor.symbol).apply {
-                        expression.getArguments().forEachIndexed { index, argument ->
-                            putValueArgument(index, argument.second)
-                        }
-                    }
-                }
-            }
-        })
     }
 
-    private val arrayGetFun by lazy {
+    private fun markInlineFunctionReferences(expression: IrFunctionAccessExpression) {
-        context.irBuiltIns.arrayClass.owner.functions.find { it.name.asString() == "get" }!!
+        val function = expression.symbol.owner
+        if (!function.isInlineFunctionCall(context))
+            return
+
+        for (parameter in function.valueParameters) {
+            if (!parameter.isInlineParameter())
+                continue
+
+            val valueArgument = expression.getValueArgument(parameter.index) ?: continue
+            if (!isInlineIrExpression(valueArgument))
+                continue
+
+            if (valueArgument is IrFunctionReference) {
+                ignoredFunctionReferences.add(valueArgument)
+            } else if (valueArgument is IrBlock) {
+                ignoredFunctionReferences.addIfNotNull(valueArgument.statements.filterIsInstance<IrFunctionReference>().singleOrNull())
+            }
+        }
     }
 
-    private val arraySizeProperty by lazy {
+    // Ignore function references handled in SAM conversion
-        context.irBuiltIns.arrayClass.owner.properties.find { it.name.toString() == "size" }!!
+    override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
+        if (expression.operator == IrTypeOperator.SAM_CONVERSION) {
+            val invokable = expression.argument
+            if (invokable is IrFunctionReference) {
+                ignoredFunctionReferences += invokable
+            } else if (invokable is IrBlock && invokable.statements.last() is IrFunctionReference) {
+                ignoredFunctionReferences += invokable.statements.last() as IrFunctionReference
+            }
+        }
+        return super.visitTypeOperator(expression)
     }
 
-    private class BuiltFunctionReference(
+    override fun visitBlock(expression: IrBlock): IrExpression {
-        val functionReferenceClass: IrClass,
+        if (!expression.origin.isLambda)
-        val functionReferenceConstructor: IrConstructor
+            return super.visitBlock(expression)
-    )
 
-    private inner class FunctionReferenceBuilder(
+        val reference = expression.statements.last() as IrFunctionReference
-        val referenceParent: IrDeclarationParent,
+        if (reference.isIgnored)
-        val irFunctionReference: IrFunctionReference
+            return super.visitBlock(expression)
-    ) {
 
-        private val isLambda = irFunctionReference.origin == IrStatementOrigin.LAMBDA
+        expression.statements.dropLast(1).forEach { it.transform(this, null) }
+        reference.transformChildrenVoid(this)
+        return FunctionReferenceBuilder(reference).build()
+    }
+
+    override fun visitFunctionReference(expression: IrFunctionReference): IrExpression {
+        expression.transformChildrenVoid(this)
+        return if (expression.isIgnored) expression else FunctionReferenceBuilder(expression).build()
+    }
+
+    private inner class FunctionReferenceBuilder(val irFunctionReference: IrFunctionReference) {
+        private val isLambda = irFunctionReference.origin.isLambda
 
         private val functionReferenceOrLambda = if (isLambda) context.ir.symbols.lambdaClass else context.ir.symbols.functionReference
 
         private val callee = irFunctionReference.symbol.owner
-        private val calleeParameters = callee.explicitParameters
+
-        private val boundCalleeParameters = irFunctionReference.getArgumentsWithIr().map { it.first }
+        // Only function references can bind a receiver and even then we can only bind either an extension or a dispatch receiver.
+        // However, when we bind a value of an inline class type as a receiver, the receiver will turn into an argument of
+        // the function in question. Yet we still need to record it as the "receiver" in CallableReference in order for reflection
+        // to work correctly.
+        private val boundReceiver: Pair<IrValueParameter, IrExpression>? = irFunctionReference.getArgumentsWithIr().singleOrNull()
 
         // The type of the reference is KFunction<in A1, ..., in An, out R>
-        private val argumentTypes = (irFunctionReference.type as IrSimpleType).arguments.dropLast(1).map { (it as IrTypeProjection).type }
+        private val parameterTypes = (irFunctionReference.type as IrSimpleType).arguments.map { (it as IrTypeProjection).type }
-        private val returnType = ((irFunctionReference.type as IrSimpleType).arguments.last() as IrTypeProjection).type
+        private val argumentTypes = parameterTypes.dropLast(1)
-        private val useVararg = (argumentTypes.size >= FunctionInvokeDescriptor.Factory.BIG_ARITY)
+        private val returnType = parameterTypes.last()
+
+        private val useVararg
+            get() = argumentTypes.size >= FunctionInvokeDescriptor.BIG_ARITY
 
         private val typeParameters = if (callee is IrConstructor)
             callee.parentAsClass.typeParameters + callee.typeParameters
@@ -189,118 +176,212 @@ internal class CallableReferenceLowering(val context: JvmBackendContext) : FileL
         private val functionReferenceClass = buildClass {
             setSourceRange(irFunctionReference)
             visibility = Visibilities.LOCAL
-            origin = JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
+            // A callable reference results in a synthetic class, while a lambda is not synthetic.
+            origin = if (isLambda) JvmLoweredDeclarationOrigin.LAMBDA_IMPL else JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
             name = Name.special("<function reference to ${callee.fqNameWhenAvailable}>")
         }.apply {
-            parent = referenceParent
+            parent = currentDeclarationParent
             superTypes += functionReferenceOrLambda.owner.defaultType
             createImplicitParameterDeclarationWithWrappedDescriptor()
             copyAttributes(irFunctionReference)
         }
 
-        private val argumentToFieldMap = boundCalleeParameters.associateWith { parameter ->
+        fun build(): IrExpression {
-            // TODO: do not store receivers to fields and get rid of this
-            val safeName = parameter.name.takeUnless(Name::isSpecial) ?: parameter.name.asString().let { name ->
-                Name.identifier("$${name.substring(1, name.length - 1)}")
-            }
-            buildField(safeName, parameter.type)
-        }
-
-        fun build(): BuiltFunctionReference {
-            val actualFunctionClass = if (useVararg)
-                context.ir.symbols.functionN
-            else
-                context.ir.symbols.getJvmFunctionClass(argumentTypes.size)
-            functionReferenceClass.superTypes +=
-                actualFunctionClass.typeWith(if (useVararg) listOf(returnType) else argumentTypes + returnType)
-
-            var suspendFunctionClass: IrClass? = null
-            val lastParameterType = (calleeParameters - boundCalleeParameters).lastOrNull()?.type
-            if (lastParameterType is IrSimpleType &&
-                lastParameterType.classOrNull?.owner?.fqNameWhenAvailable?.asString() == "kotlin.coroutines.experimental.Continuation"
-            ) {
-                // If the last parameter is Continuation<> inherit from SuspendFunction.
-                suspendFunctionClass = context.getTopLevelClass(FqName("kotlin.coroutines.SuspendFunction${argumentTypes.size - 1}")).owner
-                val continuationType = (lastParameterType.arguments.single() as IrTypeProjection).type
-                functionReferenceClass.superTypes += suspendFunctionClass.typeWith(argumentTypes + continuationType)
-            }
-
             val constructor = createConstructor()
-            createInvokeMethod(actualFunctionClass.owner.functions.find { it.name.asString() == "invoke" }!!)
+            val invokeMethod = createInvokeMethod()
+            createBridge(invokeMethod)
 
             if (!isLambda) {
-                createGetSignatureMethod(functionReferenceOrLambda.owner.functions.find { it.name.asString() == "getSignature" }!!)
+                createGetSignatureMethod(functionGetSignature)
-                createGetNameMethod(functionReferenceOrLambda.owner.functions.find { it.name.asString() == "getName" }!!)
+                createGetNameMethod(functionGetName)
-                createGetOwnerMethod(functionReferenceOrLambda.owner.functions.find { it.name.asString() == "getOwner" }!!)
+                createGetOwnerMethod(functionGetOwner)
-                if (suspendFunctionClass != null) {
-                    createInvokeMethod(suspendFunctionClass.functions.find { it.name.asString() == "invoke" }!!)
-                }
             }
 
-            return BuiltFunctionReference(functionReferenceClass, constructor)
+            return context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol).irBlock(irFunctionReference) {
+                +functionReferenceClass
+                +irCall(constructor.symbol).apply {
+                    boundReceiver?.second?.let { putValueArgument(0, it) }
+                }
+            }
         }
 
         private fun createConstructor(): IrConstructor =
             functionReferenceClass.addConstructor {
                 setSourceRange(irFunctionReference)
                 origin = JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
-                visibility = Visibilities.PUBLIC
+                visibility = if (inInlineFunctionScope) Visibilities.PUBLIC else JavaVisibilities.PACKAGE_VISIBILITY
                 returnType = functionReferenceClass.defaultType
                 isPrimary = true
             }.apply {
-                for (param in boundCalleeParameters) {
+                // Add receiver parameter for bound function references
+                boundReceiver?.first?.let { param ->
                     valueParameters += param.copyTo(
-                        this,
+                        irFunction = this,
                         index = valueParameters.size,
                         type = param.type.substitute(typeArgumentsMap)
                     )
                 }
 
-                // The syntax (object::method) only allows to bind one of them.
+                // Super constructor:
-                val hasReceiver = irFunctionReference.dispatchReceiver != null || irFunctionReference.extensionReceiver != null
+                // - For function references with bound receivers, accepts arity and receiver
-                val kFunctionRefConstructor =
+                // - For lambdas and function references without bound receivers, accepts arity
-                    functionReferenceOrLambda.owner.constructors.single { it.valueParameters.size == if (hasReceiver) 2 else 1 }
+                val kFunctionRefConstructor = functionReferenceOrLambda.owner.constructors.single {
+                    it.valueParameters.size == if (boundReceiver != null) 2 else 1
+                }
 
                 body = context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
                     +irDelegatingConstructorCall(kFunctionRefConstructor).apply {
                         putValueArgument(0, irInt(argumentTypes.size))
-                        if (hasReceiver) {
+                        if (boundReceiver != null)
-                            putValueArgument(1, irGet(valueParameters[0]))
+                            putValueArgument(1, irGet(valueParameters.first()))
-                        }
-                    }
-
-                    // Save all arguments to fields.
-                    // TODO don't write receiver again: use it from base class
-                    boundCalleeParameters.forEachIndexed { index, it ->
-                        +irSetField(
-                            irGet(functionReferenceClass.thisReceiver!!),
-                            argumentToFieldMap[it]!!,
-                            irGet(valueParameters[index])
-                        )
                     }
                     +IrInstanceInitializerCallImpl(startOffset, endOffset, functionReferenceClass.symbol, context.irBuiltIns.unitType)
                 }
             }
 
-        private fun createInvokeMethod(superFunction: IrSimpleFunction): IrSimpleFunction =
+        private fun createInvokeMethod(): IrSimpleFunction =
-            buildOverride(superFunction, callee.returnType).apply {
+            functionReferenceClass.addFunction {
-                annotations.addAll(callee.annotations)
+                name = Name.identifier("invoke")
+                returnType = callee.returnType
+                isSuspend = callee.isSuspend
+            }.apply {
+                dispatchReceiverParameter = parentAsClass.thisReceiver!!.copyTo(this)
+                if (isLambda) createLambdaInvokeMethod() else createFunctionReferenceInvokeMethod()
+            }
 
-                if (useVararg) {
+        // Inline the body of an anonymous function into the generated lambda subclass.
-                    valueParameters.add(superFunction.valueParameters[0].copyTo(this))
+        private fun IrSimpleFunction.createLambdaInvokeMethod() {
-                } else {
+            annotations += callee.annotations
-                    for ((parameter, type) in superFunction.valueParameters.zip(argumentTypes)) {
+            val valueParameterMap = callee.explicitParameters.withIndex().associate { (index, param) ->
-                        valueParameters += parameter.copyTo(this, type = type)
+                param to param.copyTo(this, index = index)
-                    }
+            }
+            valueParameters += valueParameterMap.values
+
+            body = context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
+                callee.body?.statements?.forEach { statement ->
+                    +statement.transform(object : IrElementTransformerVoid() {
+                        override fun visitGetValue(expression: IrGetValue): IrExpression {
+                            val replacement = valueParameterMap[expression.symbol.owner]
+                                ?: return super.visitGetValue(expression)
+
+                            at(expression.startOffset, expression.endOffset)
+                            return irGet(replacement)
+                        }
+
+                        override fun visitReturn(expression: IrReturn): IrExpression =
+                            if (expression.returnTargetSymbol != callee.symbol) {
+                                super.visitReturn(expression)
+                            } else {
+                                at(expression.startOffset, expression.endOffset)
+                                irReturn(expression.value.transform(this, null))
+                            }
+
+                        override fun visitDeclaration(declaration: IrDeclaration): IrStatement {
+                            if (declaration.parent == callee)
+                                declaration.parent = this@createLambdaInvokeMethod
+                            return super.visitDeclaration(declaration)
+                        }
+                    }, null)
                 }
+            }
+        }
+
+        private fun IrSimpleFunction.createFunctionReferenceInvokeMethod() {
+            for ((index, argumentType) in argumentTypes.withIndex()) {
+                addValueParameter {
+                    name = Name.identifier("p$index")
+                    type = argumentType
+                }
+            }
+
+            body = context.createIrBuilder(symbol).run {
+                var unboundIndex = 0
+                irExprBody(irCall(callee).apply {
+                    for ((typeParameter, typeArgument) in typeArgumentsMap) {
+                        putTypeArgument(typeParameter.owner.index, typeArgument)
+                    }
+
+                    for (parameter in callee.explicitParameters) {
+                        when {
+                            boundReceiver?.first == parameter ->
+                                // Bound receiver parameter
+                                irImplicitCast(
+                                    irGetField(irGet(dispatchReceiverParameter!!), functionReferenceReceiverField),
+                                    boundReceiver.second.type
+                                )
+
+                            unboundIndex >= argumentTypes.size ->
+                                // Unbound, but out of range - empty vararg or default value.
+                                // TODO For suspend functions the last argument is continuation and it is implicit:
+                                //      irCall(getContinuationSymbol, listOf(ourSymbol.descriptor.returnType!!))
+                                null
+
+                            // If a vararg parameter corresponds to exactly one KFunction argument, which is an array, that array
+                            // is forwarded as a spread. In all other cases, excess arguments are packed into a new array.
+                            //
+                            //     fun f(x: (Int, Array<String>) -> String) = x(0, arrayOf("OK", "FAIL"))
+                            //     fun g(x: (Int, String, String) -> String) = x(0, "OK", "FAIL")
+                            //     fun h(i: Int, vararg xs: String) = xs[i]
+                            //     f(::h) == g(::h)
+                            //
+                            parameter.isVararg && (unboundIndex < argumentTypes.size - 1 || argumentTypes.last() != parameter.type) ->
+                                IrVarargImpl(
+                                    startOffset, endOffset, parameter.type, parameter.varargElementType!!,
+                                    (unboundIndex until argumentTypes.size).map { irGet(valueParameters[unboundIndex++]) }
+                                )
+
+                            else ->
+                                irGet(valueParameters[unboundIndex++])
+                        }?.let { putArgument(callee, parameter, it) }
+                    }
+                })
+            }
+        }
+
+        // Build a bridge to the monomorphic invoke method. This is more elaborate than the usual BridgeLowering,
+        // since we have special handling for functions with large arity (> 22 arguments). Large arity functions
+        // are translated to functions with a single vararg argument, which checks the number of arguments and types
+        // dynamically.
+        private fun createBridge(invoke: IrSimpleFunction) {
+            // Add supertypes
+            val actualFunctionClass = if (useVararg)
+                context.ir.symbols.functionN
+            else
+                context.ir.symbols.getJvmFunctionClass(argumentTypes.size)
+
+            functionReferenceClass.superTypes += actualFunctionClass.typeWith(if (useVararg) listOf(returnType) else parameterTypes)
+            val superFunction = actualFunctionClass.owner.functions.find { it.name.asString() == "invoke" }!!
+
+            // Only add a bridge method when necessary
+            if (context.state.typeMapper.mapAsmMethod(superFunction.descriptor) ==
+                context.state.typeMapper.mapAsmMethod(invoke.descriptor)
+            ) {
+                invoke.overriddenSymbols += superFunction.symbol
+                return
+            }
+
+            // Add the invoke bridge
+            functionReferenceClass.addFunction {
+                name = superFunction.name
+                returnType = context.irBuiltIns.anyNType
+                modality = Modality.FINAL
+                visibility = Visibilities.PUBLIC
+                origin = IrDeclarationOrigin.BRIDGE
+            }.apply {
+                overriddenSymbols += superFunction.symbol
+                dispatchReceiverParameter = parentAsClass.thisReceiver!!.copyTo(this)
+                if (useVararg)
+                    valueParameters += superFunction.valueParameters[0].copyTo(this)
+                else
+                    superFunction.valueParameters.forEach { valueParameters += it.copyTo(this, type = context.irBuiltIns.anyNType) }
 
                 body = context.createIrBuilder(symbol).irBlockBody(startOffset, endOffset) {
+                    // Check the number of arguments for large arity functions
                     if (useVararg) {
-                        val varargParam = valueParameters.single()
                         +irIfThen(
                             irNotEquals(
                                 irCall(arraySizeProperty.getter!!).apply {
-                                    dispatchReceiver = irGet(varargParam)
+                                    dispatchReceiver = irGet(valueParameters.single())
                                 },
                                 irInt(argumentTypes.size)
                             ),
@@ -310,74 +391,26 @@ internal class CallableReferenceLowering(val context: JvmBackendContext) : FileL
                         )
                     }
 
-                    var unboundIndex = 0
+                    +irReturn(irCall(invoke).apply {
-                    fun consumeNextArgument() = if (useVararg) {
+                        dispatchReceiver = irGet(dispatchReceiverParameter!!)
-                        val type = argumentTypes[unboundIndex]
+
-                        irBlock(resultType = type) {
+                        for (parameter in invoke.valueParameters) {
-                            val argArray = irGet(valueParameters.single())
+                            val index = parameter.index
-                            val argIndex = irInt(unboundIndex++)
+
-                            val argValue = irTemporary(irCallOp(arrayGetFun.symbol, context.irBuiltIns.anyNType, argArray, argIndex))
+                            val argument = if (useVararg) {
-                            +irIfThen(
+                                val argArray = irGet(valueParameters.single())
-                                irNotIs(irGet(argValue), type),
+                                val argIndex = irInt(index)
-                                irCall(context.irBuiltIns.illegalArgumentExceptionSymbol).apply {
+                                irCallOp(arrayGetFun.symbol, context.irBuiltIns.anyNType, argArray, argIndex)
-                                    putValueArgument(0, irString("Wrong type, expected $type"))
+                            } else {
-                                }
+                                irGet(valueParameters[index])
-                            )
+                            }
-                            +irImplicitCast(irGet(argValue), type)
+
+                            putValueArgument(index, irImplicitCast(argument, argumentTypes[index]))
                         }
-                    } else {
+                    })
-                        irGet(valueParameters[unboundIndex++])
-                    }
-
-                    val delegation = irCall(irFunctionReference.symbol).apply {
-                        for ((typeParameter, typeArgument) in typeArgumentsMap) {
-                            putTypeArgument(typeParameter.owner.index, typeArgument)
-                        }
-
-                        for (parameter in calleeParameters) {
-                            when {
-                                argumentToFieldMap.contains(parameter) ->
-                                    // Bound parameter - read from field.
-                                    irGetField(irGet(dispatchReceiverParameter!!), argumentToFieldMap[parameter]!!)
-
-                                unboundIndex >= argumentTypes.size ->
-                                    // Unbound, but out of range - empty vararg or default value.
-                                    // TODO For suspend functions the last argument is continuation and it is implicit:
-                                    //      irCall(getContinuationSymbol, listOf(ourSymbol.descriptor.returnType!!))
-                                    null
-
-                                // If a vararg parameter corresponds to exactly one KFunction argument, which is an array, that array
-                                // is forwarded as a spread. In all other cases, excess arguments are packed into a new array.
-                                //
-                                //     fun f(x: (Int, Array<String>) -> String) = x(0, arrayOf("OK", "FAIL"))
-                                //     fun g(x: (Int, String, String) -> String) = x(0, "OK", "FAIL")
-                                //     fun h(i: Int, vararg xs: String) = xs[i]
-                                //     f(::h) == g(::h)
-                                //
-                                parameter.isVararg && (unboundIndex < argumentTypes.size - 1 || argumentTypes.last() != parameter.type) ->
-                                    IrVarargImpl(
-                                        startOffset, endOffset, parameter.type, parameter.varargElementType!!,
-                                        (unboundIndex until argumentTypes.size).map { consumeNextArgument() }
-                                    )
-
-                                else ->
-                                    consumeNextArgument()
-                            }?.let { putArgument(callee, parameter, it) }
-                        }
-                    }
-                    +irReturn(delegation)
                 }
             }
-
+        }
-        private fun buildField(fieldName: Name, fieldType: IrType): IrField =
-            functionReferenceClass.addField {
-                setSourceRange(irFunctionReference)
-                origin = JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
-                name = fieldName
-                type = fieldType
-                visibility = JavaVisibilities.PACKAGE_VISIBILITY
-                isFinal = true
-            }
 
         private fun buildOverride(superFunction: IrSimpleFunction, newReturnType: IrType = superFunction.returnType): IrSimpleFunction =
             functionReferenceClass.addFunction {
@@ -447,4 +480,38 @@ internal class CallableReferenceLowering(val context: JvmBackendContext) : FileL
             }
         }
     }
+
+    private val IrStatementOrigin?.isLambda
+        get() = this == IrStatementOrigin.LAMBDA || this == IrStatementOrigin.ANONYMOUS_FUNCTION
+
+    private val currentDeclarationParent
+        get() = allScopes.last { it.irElement is IrDeclarationParent }.irElement as IrDeclarationParent
+
+    private val inInlineFunctionScope: Boolean
+        get() = allScopes.any { scope -> (scope.irElement as? IrFunction)?.isInline ?: false }
+
+    private fun IrClassSymbol.functionByName(name: String) = owner.functions.single { it.name.asString() == name }
+
+    private val arraySizeProperty by lazy {
+        context.irBuiltIns.arrayClass.owner.properties.single { it.name.toString() == "size" }
+    }
+
+    private val arrayGetFun by lazy {
+        context.irBuiltIns.arrayClass.functionByName("get")
+    }
+
+    private val functionReferenceReceiverField =
+        context.ir.symbols.functionReference.owner.declarations.single { it is IrField && it.name.toString() == "receiver" } as IrField
+
+    private val functionGetSignature =
+        context.ir.symbols.functionReference.functionByName("getSignature")
+
+    private val functionGetName =
+        context.ir.symbols.functionReference.functionByName("getName")
+
+    private val functionGetOwner =
+        context.ir.symbols.functionReference.functionByName("getOwner")
+
+    private val functionNInvokeFun =
+        context.ir.symbols.functionN.functionByName("invoke")
 }