JVM_IR KT-44974 fix SAM-converted capturing extension lambda

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -19987,6 +19987,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/invokedynamic/sam/samConversionOnFunctionReference.kt");
             }
 
+            @Test
+            @TestMetadata("samExtFunWithCapturingLambda.kt")
+            public void testSamExtFunWithCapturingLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/invokedynamic/sam/samExtFunWithCapturingLambda.kt");
+            }
+
             @Test
             @TestMetadata("simpleFunInterfaceConstructor.kt")
             public void testSimpleFunInterfaceConstructor() throws Exception {

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

@@ -134,6 +134,8 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
             return super.visitTypeOperator(expression)
         }
 
+        val samSuperType = expression.typeOperand
+
         val invokable = expression.argument
         val reference = if (invokable is IrFunctionReference) {
             invokable
@@ -145,8 +147,6 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
         }
         reference.transformChildrenVoid()
 
-        val samSuperType = expression.typeOperand
-
         if (shouldGenerateIndySamConversions) {
             val lambdaMetafactoryArguments =
                 LambdaMetafactoryArgumentsBuilder(context, crossinlineLambdas)
@@ -169,15 +169,23 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
                 wrapWithIndySamConversion(samType, lambdaMetafactoryArguments)
             }
             is IrBlock -> {
-                val indySamConversion = wrapWithIndySamConversion(samType, lambdaMetafactoryArguments)
+                wrapFunctionReferenceInsideBlockWithIndySamConversion(samType, lambdaMetafactoryArguments, argument)
-                argument.statements[argument.statements.size - 1] = indySamConversion
-                argument.type = indySamConversion.type
-                return argument
             }
             else -> throw AssertionError("Block or function reference expected: ${expression.render()}")
         }
     }
 
+    private fun wrapFunctionReferenceInsideBlockWithIndySamConversion(
+        samType: IrType,
+        lambdaMetafactoryArguments: LambdaMetafactoryArguments,
+        block: IrBlock
+    ): IrExpression {
+        val indySamConversion = wrapWithIndySamConversion(samType, lambdaMetafactoryArguments)
+        block.statements[block.statements.size - 1] = indySamConversion
+        block.type = indySamConversion.type
+        return block
+    }
+
     private val jvmIndyLambdaMetafactoryIntrinsic = context.ir.symbols.indyLambdaMetafactoryIntrinsic
 
     private val specialNullabilityAnnotationsFqNames =
@@ -690,3 +698,5 @@ internal class FunctionReferenceLowering(private val context: JvmBackendContext)
             declaration.parent.let { it is IrClass && it.isFileClass }
     }
 }
+
+

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.ir.allOverridden
+import org.jetbrains.kotlin.backend.common.lower.VariableRemapper
 import org.jetbrains.kotlin.backend.common.lower.parents
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.ir.erasedUpperBound
@@ -14,13 +15,19 @@ import org.jetbrains.kotlin.backend.jvm.ir.isCompiledToJvmDefault
 import org.jetbrains.kotlin.builtins.functions.BuiltInFunctionArity
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.builders.declarations.buildClass
+import org.jetbrains.kotlin.ir.builders.declarations.buildValueParameter
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrFunctionReference
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionReferenceImpl
 import org.jetbrains.kotlin.ir.overrides.buildFakeOverrideMember
 import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.util.dump
+import org.jetbrains.kotlin.ir.util.functions
+import org.jetbrains.kotlin.ir.util.getInlineClassUnderlyingType
+import org.jetbrains.kotlin.ir.util.render
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
 
 class LambdaMetafactoryArguments(
@@ -192,10 +199,12 @@ class LambdaMetafactoryArgumentsBuilder(
 
         adaptLambdaSignature(implLambda, fakeInstanceMethod, signatureAdaptationConstraints)
 
+        val newReference = remapExtensionLambda(implLambda, reference)
+
         if (samMethod.isFakeOverride && nonFakeOverriddenFuns.size == 1) {
-            return LambdaMetafactoryArguments(nonFakeOverriddenFuns.single(), fakeInstanceMethod, reference, listOf())
+            return LambdaMetafactoryArguments(nonFakeOverriddenFuns.single(), fakeInstanceMethod, newReference, listOf())
         }
-        return LambdaMetafactoryArguments(samMethod, fakeInstanceMethod, reference, nonFakeOverriddenFuns)
+        return LambdaMetafactoryArguments(samMethod, fakeInstanceMethod, newReference, nonFakeOverriddenFuns)
     }
 
     private fun adaptLambdaSignature(
@@ -223,8 +232,51 @@ class LambdaMetafactoryArgumentsBuilder(
         if (constraints.returnType == TypeAdaptationConstraint.FORCE_BOXING) {
             lambda.returnType = lambda.returnType.makeNullable()
         }
+
     }
 
+    private fun remapExtensionLambda(lambda: IrSimpleFunction, reference: IrFunctionReference): IrFunctionReference {
+        val oldExtensionReceiver = lambda.extensionReceiverParameter
+            ?: return reference
+
+        val newValueParameters = ArrayList<IrValueParameter>()
+        val oldToNew = HashMap<IrValueParameter, IrValueParameter>()
+        var newParameterIndex = 0
+
+        newValueParameters.add(
+            oldExtensionReceiver.copy(lambda, newParameterIndex++, Name.identifier("\$receiver")).also {
+                oldToNew[oldExtensionReceiver] = it
+            }
+        )
+
+        lambda.valueParameters.mapTo(newValueParameters) { oldParameter ->
+            oldParameter.copy(lambda, newParameterIndex++).also {
+                oldToNew[oldParameter] = it
+            }
+        }
+
+        lambda.body?.transformChildrenVoid(VariableRemapper(oldToNew))
+
+        lambda.extensionReceiverParameter = null
+        lambda.valueParameters = newValueParameters
+
+        return IrFunctionReferenceImpl(
+            reference.startOffset, reference.endOffset, reference.type,
+            lambda.symbol,
+            typeArgumentsCount = 0,
+            valueArgumentsCount = newValueParameters.size,
+            reflectionTarget = null,
+            origin = reference.origin
+        )
+    }
+
+    private fun IrValueParameter.copy(parent: IrSimpleFunction, newIndex: Int, newName: Name = this.name): IrValueParameter =
+        buildValueParameter(parent) {
+            updateFrom(this@copy)
+            index = newIndex
+            name = newName
+        }
+
     private fun adaptFakeInstanceMethodSignature(fakeInstanceMethod: IrSimpleFunction, constraints: SignatureAdaptationConstraints) {
         for ((valueParameter, constraint) in constraints.valueParameters) {
             if (valueParameter.parent != fakeInstanceMethod)
@@ -406,13 +458,13 @@ class LambdaMetafactoryArgumentsBuilder(
 
     private fun IrType.isJvmPrimitiveType() =
         isBoolean() || isChar() || isByte() || isShort() || isInt() || isLong() || isFloat() || isDouble()
+}
 
-    private fun collectValueParameters(irFun: IrFunction): List<IrValueParameter> {
+fun collectValueParameters(irFun: IrFunction): List<IrValueParameter> {
-        if (irFun.extensionReceiverParameter == null)
+    if (irFun.extensionReceiverParameter == null)
-            return irFun.valueParameters
+        return irFun.valueParameters
-        return ArrayList<IrValueParameter>().apply {
+    return ArrayList<IrValueParameter>().apply {
-            add(irFun.extensionReceiverParameter!!)
+        add(irFun.extensionReceiverParameter!!)
-            addAll(irFun.valueParameters)
+        addAll(irFun.valueParameters)
-        }
     }
 }

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

@@ -299,13 +299,14 @@ private class TypeOperatorLowering(private val context: JvmBackendContext) : Fil
                 dynamicCallArguments.add(extensionReceiver)
             }
 
-            val samMethodValueParametersCount = samMethod.valueParameters.size
+            val samMethodValueParametersCount = samMethod.valueParameters.size +
+                    if (samMethod.extensionReceiverParameter != null && irFunRef.extensionReceiver == null) 1 else 0
             val targetFunValueParametersCount = targetFun.valueParameters.size
             for (i in 0 until targetFunValueParametersCount - samMethodValueParametersCount) {
                 val targetFunValueParameter = targetFun.valueParameters[i]
                 addValueParameter("p${syntheticParameterIndex++}", targetFunValueParameter.type)
                 val capturedValueArgument = irFunRef.getValueArgument(i)
-                    ?: fail("Captured value argument #$i (${targetFunValueParameter.name} not provided")
+                    ?: fail("Captured value argument #$i (${targetFunValueParameter.name}) not provided")
                 dynamicCallArguments.add(capturedValueArgument)
             }
         }

compiler/testData/codegen/box/annotations/annotatedLambda/samFunExpression.kt

@@ -17,7 +17,7 @@ class Test {
     }
 }
 
-// FILE: test.kt
+// FILE: samFunExpression.kt
 
 import java.lang.reflect.Method
 import kotlin.test.assertEquals

compiler/testData/codegen/box/annotations/annotatedLambda/samFunReference.kt

@@ -17,7 +17,7 @@ class Test {
     }
 }
 
-// FILE: test.kt
+// FILE: samFunReference.kt
 
 import java.lang.reflect.Method
 import kotlin.test.assertEquals

compiler/testData/codegen/box/invokedynamic/sam/samExtFunWithCapturingLambda.kt

@@ -0,0 +1,15 @@
+// TARGET_BACKEND: JVM
+// JVM_TARGET: 1.8
+// SAM_CONVERSIONS: INDY
+
+fun interface IExtFun {
+    fun String.foo(s: String) : String
+}
+
+fun box(): String {
+    val oChar = 'O'
+    val iExtFun = IExtFun { this + oChar.toString() + it }
+    return with(iExtFun) {
+        "".foo("K")
+    }
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -19987,6 +19987,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/invokedynamic/sam/samConversionOnFunctionReference.kt");
             }
 
+            @Test
+            @TestMetadata("samExtFunWithCapturingLambda.kt")
+            public void testSamExtFunWithCapturingLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/invokedynamic/sam/samExtFunWithCapturingLambda.kt");
+            }
+
             @Test
             @TestMetadata("simpleFunInterfaceConstructor.kt")
             public void testSimpleFunInterfaceConstructor() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -19987,6 +19987,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/invokedynamic/sam/samConversionOnFunctionReference.kt");
             }
 
+            @Test
+            @TestMetadata("samExtFunWithCapturingLambda.kt")
+            public void testSamExtFunWithCapturingLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/invokedynamic/sam/samExtFunWithCapturingLambda.kt");
+            }
+
             @Test
             @TestMetadata("simpleFunInterfaceConstructor.kt")
             public void testSimpleFunInterfaceConstructor() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -16758,6 +16758,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/invokedynamic/sam/samConversionOnFunctionReference.kt");
             }
 
+            @TestMetadata("samExtFunWithCapturingLambda.kt")
+            public void testSamExtFunWithCapturingLambda() throws Exception {
+                runTest("compiler/testData/codegen/box/invokedynamic/sam/samExtFunWithCapturingLambda.kt");
+            }
+
             @TestMetadata("simpleFunInterfaceConstructor.kt")
             public void testSimpleFunInterfaceConstructor() throws Exception {
                 runTest("compiler/testData/codegen/box/invokedynamic/sam/simpleFunInterfaceConstructor.kt");