Devirtualization fails to eliminate boxing in function reference context

^KT-49847 Fixed

Merge-request: KT-MR-6460
Merged-by: Vladimir Sukharev <Vladimir.Sukharev@jetbrains.com>

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

@@ -41912,6 +41912,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/reflection/genericSignature/genericBackingFieldSignature.kt");
             }
 
+            @Test
+            @TestMetadata("genericFunctionReferenceSignature.kt")
+            public void testGenericFunctionReferenceSignature() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/genericSignature/genericFunctionReferenceSignature.kt");
+            }
+
             @Test
             @TestMetadata("genericMethodSignature.kt")
             public void testGenericMethodSignature() throws Exception {

compiler/testData/codegen/box/reflection/genericSignature/genericFunctionReferenceSignature.kt

@@ -0,0 +1,12 @@
+// WITH_REFLECT
+
+package test
+
+fun <T> foo(x: T) = x
+
+fun box(): String {
+    val bar: kotlin.reflect.KFunction1<Int, Int> = ::foo
+    val returnType = bar.returnType
+    if (returnType.toString() != "T") return returnType.toString()
+    return "OK"
+}

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

@@ -41306,6 +41306,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/reflection/genericSignature/genericBackingFieldSignature.kt");
             }
 
+            @Test
+            @TestMetadata("genericFunctionReferenceSignature.kt")
+            public void testGenericFunctionReferenceSignature() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/genericSignature/genericFunctionReferenceSignature.kt");
+            }
+
             @Test
             @TestMetadata("genericMethodSignature.kt")
             public void testGenericMethodSignature() throws Exception {

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

@@ -41912,6 +41912,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/reflection/genericSignature/genericBackingFieldSignature.kt");
             }
 
+            @Test
+            @TestMetadata("genericFunctionReferenceSignature.kt")
+            public void testGenericFunctionReferenceSignature() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/genericSignature/genericFunctionReferenceSignature.kt");
+            }
+
             @Test
             @TestMetadata("genericMethodSignature.kt")
             public void testGenericMethodSignature() throws Exception {

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

@@ -33159,6 +33159,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/reflection/genericSignature/genericBackingFieldSignature.kt");
             }
 
+            @TestMetadata("genericFunctionReferenceSignature.kt")
+            public void testGenericFunctionReferenceSignature() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/genericSignature/genericFunctionReferenceSignature.kt");
+            }
+
             @TestMetadata("genericMethodSignature.kt")
             public void testGenericMethodSignature() throws Exception {
                 runTest("compiler/testData/codegen/box/reflection/genericSignature/genericMethodSignature.kt");

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

@@ -31023,6 +31023,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             public void testAllFilesPresentInGenericSignature() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/reflection/genericSignature"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS, true);
             }
+
+            @Test
+            @TestMetadata("genericFunctionReferenceSignature.kt")
+            public void testGenericFunctionReferenceSignature() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/genericSignature/genericFunctionReferenceSignature.kt");
+            }
         }
 
         @Nested

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

@@ -31125,6 +31125,12 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             public void testAllFilesPresentInGenericSignature() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/reflection/genericSignature"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JS_IR, true);
             }
+
+            @Test
+            @TestMetadata("genericFunctionReferenceSignature.kt")
+            public void testGenericFunctionReferenceSignature() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/genericSignature/genericFunctionReferenceSignature.kt");
+            }
         }
 
         @Nested

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

@@ -27852,6 +27852,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
             public void testAllFilesPresentInGenericSignature() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/reflection/genericSignature"), Pattern.compile("^([^_](.+))\\.kt$"), null, TargetBackend.WASM, true);
             }
+
+            @TestMetadata("genericFunctionReferenceSignature.kt")
+            public void testGenericFunctionReferenceSignature() throws Exception {
+                runTest("compiler/testData/codegen/box/reflection/genericSignature/genericFunctionReferenceSignature.kt");
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/reflection/isInstance")

kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/FunctionReferenceLowering.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.backend.konan.lower
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
-import org.jetbrains.kotlin.backend.common.ir.*
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.pop
 import org.jetbrains.kotlin.backend.common.push
@@ -184,6 +183,14 @@ internal class FunctionReferenceLowering(val context: Context) : FileLoweringPas
         private val typeArgumentsMap = referencedFunction.typeParameters.associate { typeParam ->
             typeParam.symbol to functionReference.getTypeArgument(typeParam.index)!!
         }
+        private val functionParameterAndReturnTypes = (functionReference.type as IrSimpleType).arguments.map {
+            when (it) {
+                is IrTypeProjection -> it.type
+                else -> context.irBuiltIns.anyNType
+            }
+        }
+        private val functionParameterTypes = functionParameterAndReturnTypes.dropLast(1)
+        private val functionReturnType = functionParameterAndReturnTypes.last()
 
         private val isLambda = functionReference.origin.isLambda
         private val isKFunction = functionReference.type.isKFunction()
@@ -241,9 +248,9 @@ internal class FunctionReferenceLowering(val context: Context) : FileLoweringPas
 
         private fun buildClass(): IrClass {
             val superClass = when {
-                isKSuspendFunction -> kSuspendFunctionImplSymbol.typeWith(referencedFunction.returnType)
+                isKSuspendFunction -> kSuspendFunctionImplSymbol.typeWith(functionReturnType)
                 isLambda -> irBuiltIns.anyType
-                else -> kFunctionImplSymbol.typeWith(referencedFunction.returnType)
+                else -> kFunctionImplSymbol.typeWith(functionReturnType)
             }
             val superTypes = mutableListOf(superClass)
             if (samSuperType != null) {
@@ -253,16 +260,15 @@ internal class FunctionReferenceLowering(val context: Context) : FileLoweringPas
                 buildInvokeMethod(sam.owner)
             } else {
                 val numberOfParameters = unboundFunctionParameters.size
-                val functionParameterTypes = unboundFunctionParameters.map { it.type }
                 val functionClass: IrClass?
                 val suspendFunctionClass: IrClass?
                 if (isKSuspendFunction) {
                     functionClass = null
                     suspendFunctionClass = symbols.kSuspendFunctionN(numberOfParameters).owner
-                    superTypes += suspendFunctionClass.typeWith(functionParameterTypes + referencedFunction.returnType)
+                    superTypes += suspendFunctionClass.typeWith(functionParameterAndReturnTypes)
                 } else {
                     functionClass = (if (isKFunction) symbols.kFunctionN(numberOfParameters) else symbols.functionN(numberOfParameters)).owner
-                    superTypes += functionClass.typeWith(functionParameterTypes + referencedFunction.returnType)
+                    superTypes += functionClass.typeWith(functionParameterAndReturnTypes)
                     val lastParameterType = unboundFunctionParameters.lastOrNull()?.type
                     if (lastParameterType?.classifierOrNull != continuationClassSymbol)
                         suspendFunctionClass = null
@@ -442,7 +448,7 @@ internal class FunctionReferenceLowering(val context: Context) : FileLoweringPas
                     superFunction.name,
                     DescriptorVisibilities.PRIVATE,
                     Modality.FINAL,
-                    referencedFunction.returnType,
+                    functionReturnType,
                     isInline = false,
                     isExternal = false,
                     isTailrec = false,
@@ -462,7 +468,7 @@ internal class FunctionReferenceLowering(val context: Context) : FileLoweringPas
 
                 valueParameters += superFunction.valueParameters.mapIndexed { index, parameter ->
                     parameter.copyTo(function, DECLARATION_ORIGIN_FUNCTION_REFERENCE_IMPL, index,
-                            type = parameter.type.substitute(typeArgumentsMap))
+                            type = functionParameterTypes[index])
                 }
 
                 overriddenSymbols += superFunction.symbol
@@ -497,6 +503,10 @@ internal class FunctionReferenceLowering(val context: Context) : FileLoweringPas
                                     }
                                 }
                                 assert(unboundIndex == valueParameters.size) { "Not all arguments of <invoke> are used" }
+
+                                referencedFunction.typeParameters.forEach { typeParam ->
+                                    putTypeArgument(typeParam.index, functionReference.getTypeArgument(typeParam.index)!!)
+                                }
                             }
                     )
                 }

kotlin-native/backend.native/tests/build.gradle

@@ -6152,6 +6152,46 @@ fileCheckTest("filecheck_replace_invoke_with_call") {
     annotatedSource = project.file('filecheck/replace_invoke_with_call.kt')
 }
 
+fileCheckTest("filecheck_kt49847_simple_function_reference") {
+    annotatedSource = project.file('filecheck/kt49847_simple_function_reference.kt')
+    enabled = project.globalTestArgs.contains('-opt')
+}
+
+fileCheckTest("filecheck_kt49847_sam_Any") {
+    annotatedSource = project.file('filecheck/kt49847_sam_Any.kt')
+    enabled = project.globalTestArgs.contains('-opt')
+}
+
+fileCheckTest("filecheck_kt49847_sam_Int") {
+    annotatedSource = project.file('filecheck/kt49847_sam_Int.kt')
+    enabled = project.globalTestArgs.contains('-opt')
+}
+
+fileCheckTest("filecheck_kt49847_sam_Any_generic") {
+    annotatedSource = project.file('filecheck/kt49847_sam_Any_generic.kt')
+    enabled = project.globalTestArgs.contains('-opt')
+}
+
+fileCheckTest("filecheck_kt49847_sam_Int_generic") {
+    annotatedSource = project.file('filecheck/kt49847_sam_Int_generic.kt')
+    enabled = project.globalTestArgs.contains('-opt')
+}
+
+fileCheckTest("filecheck_kt49847_class") {
+    annotatedSource = project.file('filecheck/kt49847_class.kt')
+    enabled = project.globalTestArgs.contains('-opt')
+}
+
+fileCheckTest("filecheck_kt49847_generic") {
+    annotatedSource = project.file('filecheck/kt49847_generic.kt')
+    enabled = project.globalTestArgs.contains('-opt')
+}
+
+fileCheckTest("filecheck_kt49847_generic_receiver") {
+    annotatedSource = project.file('filecheck/kt49847_generic_receiver.kt')
+    enabled = project.globalTestArgs.contains('-opt')
+}
+
 fileCheckTest("filecheck_intrinsics") {
     annotatedSource = project.file('filecheck/intrinsics.kt')
 }

kotlin-native/backend.native/tests/filecheck/kt49847_class.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2022 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.
+ */
+
+class C {
+    fun foo(x: Int) = x
+}
+
+// CHECK: define void @"kfun:#main(){}"()
+// CHECK-NOT: Int-box
+//   TODO Remove next check after fix of https://youtrack.jetbrains.com/issue/KT-53100/Optimization-needed-T-unboxCONSTANTPRIMITIVEx-T-x
+// CHECK: Int-unbox
+// CHECK-NOT: Int-unbox
+// CHECK: ret void
+fun main() {
+    val c = C()
+    val fooref = c::foo
+    if( fooref(42) == 42)
+        println("ok")
+}

kotlin-native/backend.native/tests/filecheck/kt49847_generic.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2022 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.
+ */
+
+class C<T> {
+    fun foo(x: T) = x
+}
+
+// CHECK: define void @"kfun:#main(){}"()
+// CHECK-NOT: Int-box
+//   TODO Remove next check after fix of https://youtrack.jetbrains.com/issue/KT-53100/Optimization-needed-T-unboxCONSTANTPRIMITIVEx-T-x
+// CHECK: Int-unbox
+// CHECK-NOT: Int-unbox
+// CHECK: ret void
+fun main() {
+    val c = C<Int>()
+    val fooref = c::foo
+    if( fooref(42) == 42)
+        println("ok")
+}

kotlin-native/backend.native/tests/filecheck/kt49847_generic_receiver.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2010-2022 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.
+ */
+
+// CHECK: define internal void @"kfun:$foo$FUNCTION_REFERENCE$0.<init>#internal"
+// CHECK-SAME: i32
+
+fun <T> T.foo() { println(this) }
+
+fun main() {
+    println(5::foo)
+}

kotlin-native/backend.native/tests/filecheck/kt49847_sam_Any.kt

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2022 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.
+ */
+
+fun interface Foo {
+    fun bar(x: Int): Any
+}
+
+fun baz(x: Any): Int = x.hashCode()
+
+// CHECK: define void @"kfun:#main(){}"()
+// Boxing/unboxing need to be used now due to non-devirtualized call
+// CHECK: Int-box
+
+//  TODO  Remove two next checks, when advanced optimization of Int-unbox(Int-box(x)) would be done for snippet like:
+//  TODO  VAR IR_TEMPORARY_VARIABLE name:arg0 type:kotlin.Any [val]
+//  TODO    BLOCK type=kotlin.Any origin=null
+//  TODO      CALL <Int-box>
+//  TODO        GET_VAR 'val arg1: kotlin.Int [val]'
+//  TODO  CALL <Int-unbox>
+//  TODO    GET_VAR 'val arg0: kotlin.Any [val]'
+// CHECK: Int-box
+// CHECK: Int-unbox
+
+// CHECK-NOT: Int-box
+// CHECK-NOT: Int-unbox
+// CHECK: ret void
+
+fun main() {
+    val foo: Foo = Foo(::baz)
+    if( foo.bar(42) == 42 )
+        println("passed")
+}

kotlin-native/backend.native/tests/filecheck/kt49847_sam_Any_generic.kt

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2022 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.
+ */
+
+fun interface Foo<T> {
+    fun bar(x: T): Any
+}
+
+fun baz(x: Any): Int = x.hashCode()
+
+// CHECK: define void @"kfun:#main(){}"()
+// Boxing/unboxing need to be used now due to non-devirtualized call
+// CHECK: Int-box
+// CHECK-NOT: Int-box
+// CHECK: Int-unbox
+// CHECK-NOT: Int-unbox
+// CHECK-NOT: Int-box
+// CHECK: ret void
+fun main() {
+    val foo: Foo<Int> = Foo(::baz)
+    if( foo.bar(42) == 42 )
+        println("passed")
+}

kotlin-native/backend.native/tests/filecheck/kt49847_sam_Int.kt

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2010-2022 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.
+ */
+
+fun interface Foo {
+    fun bar(x: Int): Int
+}
+
+fun baz(x: Int): Int = x.hashCode()
+
+// CHECK: define void @"kfun:#main(){}"()
+// CHECK-NOT: Int-box
+// CHECK-NOT: Int-unbox
+// CHECK: ret void
+fun main() {
+    val foo: Foo = Foo(::baz)
+    if( foo.bar(42) == 42 )
+        println("passed")
+}

kotlin-native/backend.native/tests/filecheck/kt49847_sam_Int_generic.kt

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2010-2022 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.
+ */
+
+fun interface Foo<T> {
+    fun bar(x: T): Int
+}
+
+fun baz(x: Any): Int = x.hashCode()
+
+// CHECK: define void @"kfun:#main(){}"()
+// CHECK-NOT: Int-box
+// CHECK-NOT: Int-unbox
+// CHECK: ret void
+fun main() {
+    val foo: Foo<Int> = Foo(::baz)
+    if( foo.bar(42) == 42 )
+        println("passed")
+}

kotlin-native/backend.native/tests/filecheck/kt49847_simple_function_reference.kt

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2010-2022 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.
+ */
+
+fun plus1(x: Int) = x + 1
+
+// CHECK: define void @"kfun:#main(){}"()
+// CHECK-NOT: Int-box
+// CHECK-NOT: Int-unbox
+// CHECK: ret void
+fun main() {
+    val ref = ::plus1
+    var y = 0
+    repeat(100000) {
+        y += ref(it)  // Should be devirtualized and invoked without boxing/unboxing (`Int-box`/`Int-unbox`)
+    }
+    if (y > 999999)
+        println("y > 999999")
+}

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

@@ -34116,6 +34116,12 @@ public class NativeCodegenBoxTestGenerated extends AbstractNativeCodegenBoxTest
                 public void testAllFilesPresentInGenericSignature() throws Exception {
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/reflection/genericSignature"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.NATIVE, true);
                 }
+
+                @Test
+                @TestMetadata("genericFunctionReferenceSignature.kt")
+                public void testGenericFunctionReferenceSignature() throws Exception {
+                    runTest("compiler/testData/codegen/box/reflection/genericSignature/genericFunctionReferenceSignature.kt");
+                }
             }
 
             @Nested