Support bound callable reference inlining in IR

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

@@ -8,13 +8,15 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 import org.jetbrains.kotlin.backend.common.ir.isInlineParameter
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.codegen.*
+import org.jetbrains.kotlin.codegen.AsmUtil.BOUND_REFERENCE_RECEIVER
 import org.jetbrains.kotlin.codegen.inline.*
 import org.jetbrains.kotlin.codegen.state.GenerationState
-import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
+import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.ir.util.getArguments
 import org.jetbrains.kotlin.utils.keysToMap
 import org.jetbrains.org.objectweb.asm.Type
@@ -53,7 +55,15 @@ class IrInlineCodegen(
         if (irValueParameter?.isInlineParameter() == true && isInlineIrExpression(argumentExpression)) {
             val irReference: IrFunctionReference =
                 (argumentExpression as IrBlock).statements.filterIsInstance<IrFunctionReference>().single()
-            rememberClosure(irReference, parameterType, irValueParameter) as IrExpressionLambdaImpl
+            val boundReceiver = argumentExpression.statements.filterIsInstance<IrVariable>().singleOrNull()
+            val lambdaInfo =
+                rememberClosure(irReference, parameterType, irValueParameter, boundReceiver) as IrExpressionLambdaImpl
+
+            if (boundReceiver != null) {
+                activeLambda = lambdaInfo
+                putCapturedValueOnStack(boundReceiver.initializer!!, lambdaInfo.capturedParamsInDesc.single(), 0)
+                activeLambda = null
+            }
         } else {
             putValueOnStack(argumentExpression, parameterType, irValueParameter?.index ?: -1, blockInfo)
         }
@@ -103,11 +113,15 @@ class IrInlineCodegen(
         }
     }
 
-    private fun rememberClosure(irReference: IrFunctionReference, type: Type, parameter: IrValueParameter): LambdaInfo {
-        //assert(InlineUtil.isInlinableParameterExpression(ktLambda)) { "Couldn't find inline expression in ${expression.text}" }
-        val expression = irReference.symbol.owner
+    private fun rememberClosure(
+        irReference: IrFunctionReference,
+        type: Type,
+        parameter: IrValueParameter,
+        boundReceiver: IrVariable?
+    ): LambdaInfo {
+        val referencedFunction = irReference.symbol.owner
         return IrExpressionLambdaImpl(
-            irReference, expression, typeMapper, parameter.isCrossinline, false/*TODO*/,
+            irReference, referencedFunction, codegen.typeMapper, parameter.isCrossinline, boundReceiver != null,
             parameter.type.isExtensionFunctionType
         ).also { lambda ->
             val closureInfo = invocationParamBuilder.addNextValueParameter(type, true, null, parameter.index)
@@ -120,27 +134,36 @@ class IrInlineCodegen(
 class IrExpressionLambdaImpl(
     val reference: IrFunctionReference,
     val function: IrFunction,
-    typeMapper: KotlinTypeMapper,
+    private val typeMapper: IrTypeMapper,
     isCrossInline: Boolean,
     override val isBoundCallableReference: Boolean,
     override val isExtensionLambda: Boolean
-) : ExpressionLambda(typeMapper, isCrossInline), IrExpressionLambda {
+) : ExpressionLambda(isCrossInline), IrExpressionLambda {
 
     override fun isReturnFromMe(labelName: String): Boolean {
         return false //always false
     }
 
-    override val lambdaClassType: Type = Type.getObjectType("test123")
+    companion object {
+        private var counter: Int = 123//TODO: pass proper type
+    }
+
+    override val lambdaClassType: Type = Type.getObjectType("test${counter++}")
 
     override val capturedVars: List<CapturedParamDesc> =
         arrayListOf<CapturedParamDesc>().apply {
             reference.getArguments().forEachIndexed { _, (_, ir) ->
-                val getValue = ir as? IrGetValue ?: error("Unrecognized expression: $ir")
-                add(capturedParamDesc(getValue.descriptor.name.asString(), typeMapper.mapType(getValue.descriptor.type)))
+                add(
+                    when (ir) {
+                        is IrGetValue -> capturedParamDesc(ir.descriptor.name.asString(), typeMapper.mapType(ir.type))
+                        is IrConst<*> -> capturedParamDesc(BOUND_REFERENCE_RECEIVER, typeMapper.mapType(ir.type))
+                        else -> error("Unrecognized expression: ${ir.dump()}")
+                    }
+                )
             }
         }
 
-    private val loweredMethod = typeMapper.mapAsmMethod(function.descriptor)
+    private val loweredMethod = typeMapper.mapAsmMethod(function)
 
     val capturedParamsInDesc: List<Type> =
         loweredMethod.argumentTypes.drop(if (isExtensionLambda) 1 else 0).take(capturedVars.size)
@@ -164,8 +187,11 @@ class IrExpressionLambdaImpl(
 fun isInlineIrExpression(argumentExpression: IrExpression) =
     when (argumentExpression) {
         is IrBlock -> (argumentExpression.origin == IrStatementOrigin.LAMBDA || argumentExpression.origin == IrStatementOrigin.ANONYMOUS_FUNCTION)
-        //TODO: support bound CR
-        is IrCallableReference -> argumentExpression.dispatchReceiver == null && argumentExpression.extensionReceiver == null
+        is IrCallableReference -> true.also {
+            assert((0 until argumentExpression.valueArgumentsCount).count { argumentExpression.getValueArgument(it) != null } == 0) {
+                "Expecting 0 value arguments for bounded callable reference: ${argumentExpression.dump()}"
+            }
+        }
         else -> false
     }
 

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

@@ -8,6 +8,7 @@ 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.ScopeWithIr
+import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom
 import org.jetbrains.kotlin.backend.common.ir.copyValueParametersToStatic
 import org.jetbrains.kotlin.backend.common.ir.isInlineParameter
@@ -18,6 +19,7 @@ import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.codegen.isInlineFunctionCall
 import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrExpression
+import org.jetbrains.kotlin.codegen.AsmUtil.BOUND_REFERENCE_RECEIVER
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
@@ -25,6 +27,7 @@ import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.IrConstructor
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
 import org.jetbrains.kotlin.ir.util.defaultType
@@ -81,8 +84,13 @@ internal class InlineCallableReferenceToLambdaPhase(val context: JvmBackendConte
                 //..else use field itself
                 val irBuilder =
                     context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol, expression.startOffset, expression.endOffset)
-
+                val boundReceiver = expression.dispatchReceiver ?: expression.extensionReceiver
                 return irBuilder.irBlock(expression, IrStatementOrigin.LAMBDA) {
+                    lateinit var variableForBoundReceiver: IrVariable
+                    if (boundReceiver != null) {
+                        variableForBoundReceiver = createTmpVariable(boundReceiver, BOUND_REFERENCE_RECEIVER)
+                    }
+
                     val newLambda = buildFun {
                         setSourceRange(expression)
                         origin = JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
@@ -93,8 +101,11 @@ internal class InlineCallableReferenceToLambdaPhase(val context: JvmBackendConte
                     }.apply {
 
                         val receiver =
-                            if (field.isStatic) null
-                            else addValueParameter("receiver", field.parentAsClass.defaultType)
+                            when {
+                                field.isStatic -> null
+                                boundReceiver != null -> variableForBoundReceiver
+                                else -> addValueParameter("receiver", field.parentAsClass.defaultType)
+                            }
 
                         val lambdaBodyBuilder = this@InlineCallableReferenceToLambdaPhase.context.createIrBuilder(this.symbol)
                         body = lambdaBodyBuilder.irBlockBody(startOffset, endOffset) {
@@ -130,7 +141,13 @@ internal class InlineCallableReferenceToLambdaPhase(val context: JvmBackendConte
         val irBuilder =
             context.createIrBuilder(scope.scope.scopeOwnerSymbol, expression.startOffset, expression.endOffset)
 
+        val boundReceiver = expression.dispatchReceiver ?: expression.extensionReceiver
         return irBuilder.irBlock(expression, IrStatementOrigin.LAMBDA) {
+            lateinit var variableForBoundReceiver: IrVariable
+            if (boundReceiver != null) {
+                variableForBoundReceiver = createTmpVariable(boundReceiver, BOUND_REFERENCE_RECEIVER)
+            }
+
             val newLambda = buildFun {
                 setSourceRange(expression)
                 origin = JvmLoweredDeclarationOrigin.FUNCTION_REFERENCE_IMPL
@@ -143,23 +160,42 @@ internal class InlineCallableReferenceToLambdaPhase(val context: JvmBackendConte
                     copyTypeParametersFrom(referencedFunction.parentAsClass)
                 }
                 copyTypeParametersFrom(referencedFunction)
-                copyValueParametersToStatic(referencedFunction, origin)
+                if (boundReceiver == null) {
+                    copyValueParametersToStatic(referencedFunction, origin)
+                } else {
+                    for (oldValueParameter in referencedFunction.valueParameters) {
+                        valueParameters.add(
+                            oldValueParameter.copyTo(
+                                this,
+                                origin = origin,
+                                index = oldValueParameter.index
+                            )
+                        )
+                    }
+                }
                 val lambdaBodyBuilder = this@InlineCallableReferenceToLambdaPhase.context.createIrBuilder(this.symbol)
                 body = lambdaBodyBuilder.irBlockBody(startOffset, endOffset) {
                     var shift = 0
-                    +irReturn(
-                        irCall(referencedFunction.symbol).also { call ->
+                    val irCall =
+                        if (expression is IrPropertyReference)
+                            irGet(referencedFunction.returnType, null, referencedFunction.symbol)
+                        else irCall(referencedFunction.symbol)
 
+                    +irReturn(
+                        irCall.also { call ->
                             for (it in this@apply.typeParameters) {
                                 call.putTypeArgument(it.index, expression.getTypeArgument(it.index))
                             }
 
                             referencedFunction.dispatchReceiverParameter?.let {
-                                call.dispatchReceiver = irGet(valueParameters[shift++])
+                                call.dispatchReceiver =
+                                    irGet(if (expression.dispatchReceiver != null) variableForBoundReceiver else valueParameters[shift++])
                             }
                             referencedFunction.extensionReceiverParameter?.let {
-                                call.extensionReceiver = irGet(valueParameters[shift++])
+                                call.extensionReceiver =
+                                    irGet(if (expression.extensionReceiver != null) variableForBoundReceiver else valueParameters[shift++])
                             }
+
                             for (it in referencedFunction.valueParameters.indices) {
                                 call.putValueArgument(it, irGet(valueParameters[shift++]))
                             }