[Common IR] Forward the extension receiver in the reified wrapper

^KT-53672 Fixed

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

@@ -4032,6 +4032,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceArguments.kt");
             }
 
+            @Test
+            @TestMetadata("genericCallableReferenceWithReifiedTypeParam.kt")
+            public void testGenericCallableReferenceWithReifiedTypeParam() throws Exception {
+                runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceWithReifiedTypeParam.kt");
+            }
+
             @Test
             @TestMetadata("genericCallableReferencesWithNullableTypes.kt")
             public void testGenericCallableReferencesWithNullableTypes() throws Exception {

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/WrapInlineDeclarationsWithReifiedTypeParametersLowering.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.ir.util.typeSubstitutionMap
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 // Replace callable reference on inline function with reified parameter
 // with callable reference on new non inline function with substituted types
@@ -67,6 +68,15 @@ class WrapInlineDeclarationsWithReifiedTypeParametersLowering(val context: Backe
                 }.apply {
                     parent = container as IrDeclarationParent
                     val irBuilder = context.createIrBuilder(symbol, SYNTHETIC_OFFSET, SYNTHETIC_OFFSET)
+                    val forwardExtensionReceiverAsParam = owner.extensionReceiverParameter?.let { extensionReceiver ->
+                        runIf(expression.extensionReceiver == null) {
+                            addValueParameter(
+                                extensionReceiver.name,
+                                typeSubstitutor.substitute(extensionReceiver.type)
+                            )
+                            true
+                        }
+                    } ?: false
                     owner.valueParameters.forEach { valueParameter ->
                         addValueParameter(
                             valueParameter.name,
@@ -80,9 +90,15 @@ class WrapInlineDeclarationsWithReifiedTypeParametersLowering(val context: Backe
                         statements.add(
                             irBuilder.irReturn(
                                 irBuilder.irCall(owner.symbol).also { call ->
+                                    val (extensionReceiver, forwardedParams) = if (forwardExtensionReceiverAsParam) {
+                                        irBuilder.irGet(valueParameters.first()) to valueParameters.subList(1, valueParameters.size)
+                                    } else {
+                                        expression.extensionReceiver to valueParameters
+                                    }
+                                    call.extensionReceiver = extensionReceiver
                                     call.dispatchReceiver = expression.dispatchReceiver
-                                    call.extensionReceiver = expression.extensionReceiver
-                                    valueParameters.forEachIndexed { index, valueParameter ->
+
+                                    forwardedParams.forEachIndexed { index, valueParameter ->
                                         call.putValueArgument(index, irBuilder.irGet(valueParameter))
                                     }
                                     for (i in 0 until expression.typeArgumentsCount) {
@@ -110,4 +126,4 @@ class WrapInlineDeclarationsWithReifiedTypeParametersLowering(val context: Backe
             }
         })
     }
-}
+}

compiler/testData/codegen/box/callableReference/function/genericCallableReferenceWithReifiedTypeParam.kt

@@ -0,0 +1,114 @@
+// WITH_STDLIB
+
+import kotlin.test.assertEquals
+
+
+inline fun <reified T> funNoArgs() = "OK" as? T
+
+fun testFunctionNoArgs() {
+    val callable: () -> String? = ::funNoArgs
+    assertEquals(callable(), "OK")
+}
+
+inline fun <reified T> funWithArgs(x: T, y: T) = x to y
+
+fun testFunctionWithArgs() {
+    val callable: (String, String) -> Pair<String, String> = ::funWithArgs
+    assertEquals(callable("O", "K"), "O" to "K")
+}
+
+inline fun <reified T> funWithVarargs(vararg i: T) = i.toList()
+
+fun testFunctionWithVarargs() {
+    val callable: (Array<Int>) -> List<Int> = ::funWithVarargs
+    assertEquals(callable(arrayOf(1, 2, 3)), listOf(1, 2, 3))
+}
+
+inline fun <reified T> T.funWithExtensionNoArgs() = this
+
+fun testFunctionWithExtensionNoArgs() {
+    val callable1 = String::funWithExtensionNoArgs
+    assertEquals(callable1("OK1"), "OK1")
+
+    val callable2 = "OK2"::funWithExtensionNoArgs
+    assertEquals(callable2(), "OK2")
+
+    val callable3 = callable1::funWithExtensionNoArgs
+    assertEquals(callable3()("OK3"), "OK3")
+
+    val callable4 = with("OK4") { ::funWithExtensionNoArgs }
+    assertEquals(callable4(), "OK4")
+}
+
+inline fun <reified T> T.funWithExtensionAndArgs(x: Int, y: Int) = this to (x + y)
+
+fun testFunctionWithExtensionAndArgs() {
+    val callable1 = String::funWithExtensionAndArgs
+    assertEquals(callable1("OK1", 1, 2), "OK1" to 3)
+
+    val callable2 = "OK2"::funWithExtensionAndArgs
+    assertEquals(callable2(3, 4), "OK2" to 7)
+
+    val callable3 = callable1::funWithExtensionAndArgs
+    val (cb, s) = callable3(5, 6)
+    assertEquals(s, 11)
+    assertEquals(cb("OK3", 7, 8), "OK3" to 15)
+
+    val callable4 = with("OK4") { ::funWithExtensionAndArgs }
+    assertEquals(callable4(9, 10), "OK4" to 19)
+}
+
+inline fun <reified T> T.funWithExtensionAndVarargs(vararg i: Int) = this to i.sum()
+
+fun testFunctionWithExtensionAndVararg() {
+    val callable1 = String::funWithExtensionAndVarargs
+    assertEquals(callable1("OK1", arrayOf(1, 2, 3).toIntArray()), "OK1" to 6)
+
+    val callable2 = "OK2"::funWithExtensionAndVarargs
+    assertEquals(callable2(arrayOf(4, 5, 6).toIntArray()), "OK2" to 15)
+
+    val callable3 = callable1::funWithExtensionAndVarargs
+    val (cb, s) = callable3(arrayOf(7, 8).toIntArray())
+    assertEquals(s, 15)
+    assertEquals(cb("OK3", arrayOf(9, 10).toIntArray()), "OK3" to 19)
+
+    val callable4 = with("OK4") { ::funWithExtensionAndVarargs }
+    assertEquals(callable4(arrayOf(11, 12).toIntArray()), "OK4" to 23)
+}
+
+class TestClass(val s: String) {
+    inline fun <reified T> classFunNoArgs() = s as? T
+    inline fun <reified T> classFunWithArgs(x: T) = x to s
+    inline fun <reified T> classFunWithVarargs(vararg i: T) = i.toList() to s
+}
+
+fun testClassFunctionNoArgs() {
+    val callable: () -> String? = with(TestClass("OK1")) { ::classFunNoArgs }
+    assertEquals(callable(), "OK1")
+}
+
+fun testClassFunctionWithArgs() {
+    val callable: (Int) -> Pair<Int, String> = with(TestClass("OK1")) { ::classFunWithArgs }
+    assertEquals(callable(1), 1 to "OK1")
+}
+
+fun testClassFunctionWithVarargs() {
+    val callable: (Array<Int>) -> Pair<List<Int>, String> = with(TestClass("OK1")) { ::classFunWithVarargs }
+    assertEquals(callable(arrayOf(1, 2)), listOf(1, 2) to "OK1")
+}
+
+fun box(): String {
+    testFunctionNoArgs()
+    testFunctionWithArgs()
+    testFunctionWithVarargs()
+
+    testFunctionWithExtensionNoArgs()
+    testFunctionWithExtensionAndArgs()
+    testFunctionWithExtensionAndVararg()
+
+    testClassFunctionNoArgs()
+    testClassFunctionWithArgs()
+    testClassFunctionWithVarargs()
+
+    return "OK"
+}

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

@@ -3918,6 +3918,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceArguments.kt");
             }
 
+            @Test
+            @TestMetadata("genericCallableReferenceWithReifiedTypeParam.kt")
+            public void testGenericCallableReferenceWithReifiedTypeParam() throws Exception {
+                runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceWithReifiedTypeParam.kt");
+            }
+
             @Test
             @TestMetadata("genericCallableReferencesWithNullableTypes.kt")
             public void testGenericCallableReferencesWithNullableTypes() throws Exception {

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

@@ -4032,6 +4032,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceArguments.kt");
             }
 
+            @Test
+            @TestMetadata("genericCallableReferenceWithReifiedTypeParam.kt")
+            public void testGenericCallableReferenceWithReifiedTypeParam() throws Exception {
+                runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceWithReifiedTypeParam.kt");
+            }
+
             @Test
             @TestMetadata("genericCallableReferencesWithNullableTypes.kt")
             public void testGenericCallableReferencesWithNullableTypes() throws Exception {

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

@@ -3428,6 +3428,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceArguments.kt");
             }
 
+            @TestMetadata("genericCallableReferenceWithReifiedTypeParam.kt")
+            public void testGenericCallableReferenceWithReifiedTypeParam() throws Exception {
+                runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceWithReifiedTypeParam.kt");
+            }
+
             @TestMetadata("genericCallableReferencesWithNullableTypes.kt")
             public void testGenericCallableReferencesWithNullableTypes() throws Exception {
                 runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferencesWithNullableTypes.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/JsCodegenBoxTestGenerated.java

@@ -2736,6 +2736,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceArguments.kt");
             }
 
+            @Test
+            @TestMetadata("genericCallableReferenceWithReifiedTypeParam.kt")
+            public void testGenericCallableReferenceWithReifiedTypeParam() throws Exception {
+                runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceWithReifiedTypeParam.kt");
+            }
+
             @Test
             @TestMetadata("genericCallableReferencesWithNullableTypes.kt")
             public void testGenericCallableReferencesWithNullableTypes() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrJsCodegenBoxTestGenerated.java

@@ -2784,6 +2784,12 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceArguments.kt");
             }
 
+            @Test
+            @TestMetadata("genericCallableReferenceWithReifiedTypeParam.kt")
+            public void testGenericCallableReferenceWithReifiedTypeParam() throws Exception {
+                runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceWithReifiedTypeParam.kt");
+            }
+
             @Test
             @TestMetadata("genericCallableReferencesWithNullableTypes.kt")
             public void testGenericCallableReferencesWithNullableTypes() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/testOld/wasm/semantics/IrCodegenBoxWasmTestGenerated.java

@@ -2473,6 +2473,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
                 runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceArguments.kt");
             }
 
+            @TestMetadata("genericCallableReferenceWithReifiedTypeParam.kt")
+            public void testGenericCallableReferenceWithReifiedTypeParam() throws Exception {
+                runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceWithReifiedTypeParam.kt");
+            }
+
             @TestMetadata("genericCallableReferencesWithNullableTypes.kt")
             public void testGenericCallableReferencesWithNullableTypes() throws Exception {
                 runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferencesWithNullableTypes.kt");

native/native.tests/tests-gen/org/jetbrains/kotlin/konan/blackboxtest/NativeCodegenBoxTestGenerated.java

@@ -2850,6 +2850,12 @@ public class NativeCodegenBoxTestGenerated extends AbstractNativeCodegenBoxTest
                     runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceArguments.kt");
                 }
 
+                @Test
+                @TestMetadata("genericCallableReferenceWithReifiedTypeParam.kt")
+                public void testGenericCallableReferenceWithReifiedTypeParam() throws Exception {
+                    runTest("compiler/testData/codegen/box/callableReference/function/genericCallableReferenceWithReifiedTypeParam.kt");
+                }
+
                 @Test
                 @TestMetadata("genericCallableReferencesWithNullableTypes.kt")
                 public void testGenericCallableReferencesWithNullableTypes() throws Exception {