JVM_IR: when convering references to lambdas, bind the receiver

instead of storing it into a captured temporary.

This makes the resulting lambdas equivalend to adapted function
references produced by PSI2IR or FIR2IR.

 #KT-42492 Fixed

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

@@ -77,15 +77,15 @@ class IrInlineCodegen(
     ) {
         val isInlineParameter = irValueParameter.isInlineParameter()
         if (isInlineParameter && isInlineIrExpression(argumentExpression)) {
-            val irReference: IrFunctionReference =
+            val irReference = (argumentExpression as IrBlock).statements.filterIsInstance<IrFunctionReference>().single()
-                (argumentExpression as IrBlock).statements.filterIsInstance<IrFunctionReference>().single()
+            val lambdaInfo = IrExpressionLambdaImpl(codegen, irReference, irValueParameter)
-            val boundReceiver = argumentExpression.statements.filterIsInstance<IrVariable>().singleOrNull()
+            val closureInfo = invocationParamBuilder.addNextValueParameter(parameterType, true, null, irValueParameter.index)
-            val lambdaInfo =
+            closureInfo.functionalArgument = lambdaInfo
-                rememberClosure(irReference, parameterType, irValueParameter, boundReceiver) as IrExpressionLambdaImpl
+            expressionMap[closureInfo.index] = lambdaInfo
-
+            val boundReceiver = irReference.extensionReceiver
             if (boundReceiver != null) {
                 activeLambda = lambdaInfo
-                putCapturedValueOnStack(boundReceiver.initializer!!, lambdaInfo.capturedParamsInDesc.single(), 0)
+                putCapturedValueOnStack(boundReceiver, lambdaInfo.capturedParamsInDesc.single(), 0)
                 activeLambda = null
             }
         } else {
@@ -153,23 +153,6 @@ class IrInlineCodegen(
         )
     }
 
-    private fun rememberClosure(
-        irReference: IrFunctionReference,
-        type: Type,
-        parameter: IrValueParameter,
-        boundReceiver: IrVariable?
-    ): LambdaInfo {
-        val referencedFunction = irReference.symbol.owner
-        return IrExpressionLambdaImpl(
-            irReference, referencedFunction, codegen.typeMapper, codegen.methodSignatureMapper, codegen.context, parameter.isCrossinline,
-            boundReceiver != null, parameter.type.isExtensionFunctionType
-        ).also { lambda ->
-            val closureInfo = invocationParamBuilder.addNextValueParameter(type, true, null, parameter.index)
-            closureInfo.functionalArgument = lambda
-            expressionMap[closureInfo.index] = lambda
-        }
-    }
-
     override fun extractDefaultLambdas(node: MethodNode): List<DefaultLambda> {
         if (maskStartIndex == -1) return listOf()
         return expandMaskConditionsAndUpdateVariableNodes(
@@ -181,15 +164,17 @@ class IrInlineCodegen(
 }
 
 class IrExpressionLambdaImpl(
+    codegen: ExpressionCodegen,
     val reference: IrFunctionReference,
-    val function: IrFunction,
+    irValueParameter: IrValueParameter
-    private val typeMapper: IrTypeMapper,
+) : ExpressionLambda(irValueParameter.isCrossinline), IrExpressionLambda {
-    methodSignatureMapper: MethodSignatureMapper,
+    override val isExtensionLambda: Boolean = irValueParameter.type.isExtensionFunctionType
-    context: JvmBackendContext,
+
-    isCrossInline: Boolean,
+    val function: IrFunction
-    override val isBoundCallableReference: Boolean,
+        get() = reference.symbol.owner
-    override val isExtensionLambda: Boolean
+
-) : ExpressionLambda(isCrossInline), IrExpressionLambda {
+    override val isBoundCallableReference: Boolean
+        get() = reference.extensionReceiver != null
 
     override val isSuspend: Boolean = function.isSuspend
 
@@ -201,17 +186,17 @@ class IrExpressionLambdaImpl(
     // arguments apart from any other scope's. So long as it's unique, any value is fine.
     // This particular string slightly aids in debugging internal compiler errors as it at least
     // points towards the function containing the lambda.
-    override val lambdaClassType: Type =
+    override val lambdaClassType: Type = codegen.context.getLocalClassType(reference)
-        context.getLocalClassType(reference) ?: throw AssertionError("callable reference ${reference.dump()} has no name in context")
+        ?: throw AssertionError("callable reference ${reference.dump()} has no name in context")
 
     private val capturedParameters: Map<CapturedParamDesc, IrValueParameter> =
         reference.getArgumentsWithIr().associate { (param, _) ->
-            capturedParamDesc(param.name.asString(), typeMapper.mapType(param.type)) to param
+            capturedParamDesc(param.name.asString(), codegen.typeMapper.mapType(param.type)) to param
         }
 
     override val capturedVars: List<CapturedParamDesc> = capturedParameters.keys.toList()
 
-    private val loweredMethod = methodSignatureMapper.mapAsmMethod(function)
+    private val loweredMethod = codegen.methodSignatureMapper.mapAsmMethod(function)
 
     val capturedParamsInDesc: List<Type> = if (isBoundCallableReference) {
         loweredMethod.argumentTypes.take(1)

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

@@ -7,7 +7,6 @@ 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.ir.allParameters
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.irBlock
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
@@ -15,9 +14,9 @@ import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator
 import org.jetbrains.kotlin.backend.jvm.ir.createJvmIrBuilder
 import org.jetbrains.kotlin.backend.jvm.ir.irArray
-import org.jetbrains.kotlin.codegen.AsmUtil.BOUND_REFERENCE_RECEIVER
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.builders.declarations.addExtensionReceiver
 import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
 import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.*
@@ -86,6 +85,7 @@ private class InlineCallableReferenceToLambdaTransformer(
     private fun expandInlineFieldReferenceToLambda(expression: IrPropertyReference, field: IrField): IrExpression {
         val irBuilder = context.createJvmIrBuilder(currentScope!!.scope.scopeOwnerSymbol, expression.startOffset, expression.endOffset)
         return irBuilder.irBlock(expression, IrStatementOrigin.LAMBDA) {
+            val boundReceiver = expression.dispatchReceiver ?: expression.extensionReceiver
             val function = context.irFactory.buildFun {
                 setSourceRange(expression)
                 origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
@@ -95,15 +95,11 @@ private class InlineCallableReferenceToLambdaTransformer(
                 isSuspend = false
             }.apply {
                 parent = currentDeclarationParent ?: error("No current declaration parent at ${expression.dump()}")
-                val boundReceiver = expression.dispatchReceiver ?: expression.extensionReceiver
+                val receiver = when {
-
+                    field.isStatic -> null
-                val receiver =
+                    boundReceiver != null -> irGet(addExtensionReceiver(boundReceiver.type))
-                    when {
+                    else -> irGet(addValueParameter("receiver", field.parentAsClass.defaultType))
-                        field.isStatic -> null
+                }
-                        boundReceiver != null -> irGet(irTemporary(boundReceiver, BOUND_REFERENCE_RECEIVER))
-                        else -> irGet(addValueParameter("receiver", field.parentAsClass.defaultType))
-                    }
-
                 body = this@InlineCallableReferenceToLambdaTransformer.context.createIrBuilder(symbol).run {
                     irExprBody(irGetField(receiver, field))
                 }
@@ -120,6 +116,7 @@ private class InlineCallableReferenceToLambdaTransformer(
                 origin = IrStatementOrigin.LAMBDA
             ).apply {
                 copyAttributes(expression)
+                extensionReceiver = boundReceiver
             }
         }
     }
@@ -132,10 +129,8 @@ private class InlineCallableReferenceToLambdaTransformer(
             // specific because of approximation. See compiler/testData/codegen/box/callableReference/function/argumentTypes.kt
             val boundReceiver: Pair<IrValueParameter, IrExpression>? = expression.getArgumentsWithIr().singleOrNull()
             val nParams = (expression.type as IrSimpleType).arguments.size - 1
-            var toDropAtStart = 0
+            val toDropAtStart = if (boundReceiver != null) 1 else 0
-            if (boundReceiver != null) toDropAtStart++
+            val argumentTypes = referencedFunction.explicitParameters.drop(toDropAtStart).take(nParams).map { parameter ->
-            if (referencedFunction is IrConstructor) toDropAtStart++
-            val argumentTypes = referencedFunction.allParameters.drop(toDropAtStart).take(nParams).map { parameter ->
                 parameter.type.substitute(
                     referencedFunction.typeParameters,
                     referencedFunction.typeParameters.indices.map { expression.getTypeArgument(it)!! }
@@ -151,6 +146,9 @@ private class InlineCallableReferenceToLambdaTransformer(
                 isSuspend = referencedFunction.isSuspend
             }.apply {
                 parent = currentDeclarationParent!!
+                if (boundReceiver != null) {
+                    addExtensionReceiver(boundReceiver.first.type)
+                }
                 for ((index, argumentType) in argumentTypes.withIndex()) {
                     addValueParameter {
                         name = Name.identifier("p$index")
@@ -172,7 +170,7 @@ private class InlineCallableReferenceToLambdaTransformer(
                         for (parameter in referencedFunction.explicitParameters) {
                             when {
                                 boundReceiver?.first == parameter ->
-                                    irGet(irTemporary(boundReceiver.second))
+                                    irGet(extensionReceiverParameter!!)
                                 parameter.isVararg && unboundIndex < argumentTypes.size && parameter.type == valueParameters[unboundIndex].type ->
                                     irGet(valueParameters[unboundIndex++])
                                 parameter.isVararg && (unboundIndex < argumentTypes.size || !parameter.hasDefaultValue()) ->
@@ -200,6 +198,7 @@ private class InlineCallableReferenceToLambdaTransformer(
                 origin = IrStatementOrigin.LAMBDA
             ).apply {
                 copyAttributes(expression)
+                extensionReceiver = boundReceiver?.second
             }
         }
     }

compiler/testData/codegen/box/builderInference/callableReferenceAndCoercionToUnit.kt

@@ -1,5 +1,4 @@
 // IGNORE_BACKEND: JS_IR
-// IGNORE_BACKEND_FIR: JVM_IR
 // !LANGUAGE: +NewInference
 // !DIAGNOSTICS: -EXPERIMENTAL_API_USAGE_ERROR -UNUSED_EXPRESSION
 // WITH_RUNTIME

compiler/testData/codegen/box/inference/builderInference.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // TARGET_BACKEND: JVM
 // WITH_RUNTIME
 // !LANGUAGE: +NewInference