[K/N] Make classes generated from λs that use generic parameters generic

Previously, a function reference that used generic parameters from its
outer scope was lowered into a top-level non-generic subclass of
`FunctionN`, with `FunctionN` type arguments referencing type parameters
not present in the scope anymore. This sometimes resulted in generating
malformed mangled names.

From now on the generated subclass of `FunctionN` is generic. The needed
type arguments are passed upon instantiation, when the relevant generic
parameters are present in the scope.

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

@@ -6197,6 +6197,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
         }
 
+        @Test
+        @TestMetadata("closureCapturingGenericParam.kt")
+        public void testClosureCapturingGenericParam() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+        }
+
         @Test
         @TestMetadata("closureInsideConstrucor.kt")
         public void testClosureInsideConstrucor() throws Exception {

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/IrTypeParameterRemapper.kt

@@ -13,11 +13,12 @@ import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.types.impl.IrTypeAbbreviationImpl
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 
-/* After moving an IrElement, some type parameter references within it may become out of scope.
+/**
-   This remapper restores validity by redirecting those references to new type parameters.
+ * After moving an [org.jetbrains.kotlin.ir.IrElement], some type parameter references within it may become out of scope.
+ * This remapper restores validity by redirecting those references to new type parameters.
  */
 class IrTypeParameterRemapper(
-    val typeParameterMap: Map<IrTypeParameter, IrTypeParameter>
+    private val typeParameterMap: Map<IrTypeParameter, IrTypeParameter>
 ) : TypeRemapper {
     override fun enterScope(irTypeParametersContainer: IrTypeParametersContainer) {}
     override fun leaveScope() {}
@@ -56,4 +57,4 @@ class IrTypeParameterRemapper(
         ).apply {
             annotations.forEach { it.remapTypes(this@IrTypeParameterRemapper) }
         }
-}
+}

compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt

@@ -0,0 +1,20 @@
+interface IntConvertible {
+    fun toInt(): Int
+}
+
+fun <FooTP> foo(init: Int, v: FooTP, l: Int.(FooTP) -> Int) = init.l(v)
+
+fun <BarTP : IntConvertible> computeSum(array: Array<BarTP>) = foo(0, array) {
+    var res = this
+    for (element in it) res += element.toInt()
+    res
+}
+
+class N(val v: Int) : IntConvertible {
+    override fun toInt() = v
+}
+
+fun box(): String {
+    if (computeSum(arrayOf(N(2), N(14))) != 16) return "Fail"
+    return "OK"
+}

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

@@ -6041,6 +6041,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
         }
 
+        @Test
+        @TestMetadata("closureCapturingGenericParam.kt")
+        public void testClosureCapturingGenericParam() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+        }
+
         @Test
         @TestMetadata("closureInsideConstrucor.kt")
         public void testClosureInsideConstrucor() throws Exception {

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

@@ -6197,6 +6197,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
         }
 
+        @Test
+        @TestMetadata("closureCapturingGenericParam.kt")
+        public void testClosureCapturingGenericParam() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+        }
+
         @Test
         @TestMetadata("closureInsideConstrucor.kt")
         public void testClosureInsideConstrucor() throws Exception {

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

@@ -5277,6 +5277,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
         }
 
+        @TestMetadata("closureCapturingGenericParam.kt")
+        public void testClosureCapturingGenericParam() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+        }
+
         @TestMetadata("closureInsideConstrucor.kt")
         public void testClosureInsideConstrucor() throws Exception {
             runTest("compiler/testData/codegen/box/closures/closureInsideConstrucor.kt");

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

@@ -4547,6 +4547,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
         }
 
+        @Test
+        @TestMetadata("closureCapturingGenericParam.kt")
+        public void testClosureCapturingGenericParam() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+        }
+
         @Test
         @TestMetadata("closureInsideConstrucor.kt")
         public void testClosureInsideConstrucor() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/fir/FirJsCodegenBoxTestGenerated.java

@@ -4607,6 +4607,12 @@ public class FirJsCodegenBoxTestGenerated extends AbstractFirJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
         }
 
+        @Test
+        @TestMetadata("closureCapturingGenericParam.kt")
+        public void testClosureCapturingGenericParam() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+        }
+
         @Test
         @TestMetadata("closureInsideConstrucor.kt")
         public void testClosureInsideConstrucor() throws Exception {

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

@@ -4607,6 +4607,12 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
             runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
         }
 
+        @Test
+        @TestMetadata("closureCapturingGenericParam.kt")
+        public void testClosureCapturingGenericParam() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+        }
+
         @Test
         @TestMetadata("closureInsideConstrucor.kt")
         public void testClosureInsideConstrucor() throws Exception {

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

@@ -4067,6 +4067,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
             runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
         }
 
+        @TestMetadata("closureCapturingGenericParam.kt")
+        public void testClosureCapturingGenericParam() throws Exception {
+            runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+        }
+
         @TestMetadata("closureInsideConstrucor.kt")
         public void testClosureInsideConstrucor() throws Exception {
             runTest("compiler/testData/codegen/box/closures/closureInsideConstrucor.kt");

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

@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.pop
 import org.jetbrains.kotlin.backend.common.push
-import org.jetbrains.kotlin.backend.konan.Context
 import org.jetbrains.kotlin.backend.konan.NativeGenerationState
 import org.jetbrains.kotlin.backend.konan.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.konan.llvm.computeFullName
@@ -29,6 +28,37 @@ import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
+/**
+ * Transforms a function reference into a subclass of `kotlin.native.internal.KFunctionImpl` and `kotlin.FunctionN`,
+ * or `kotlin.native.internal.KSuspendFunctionImpl` and `kotlin.KSuspendFunctionN` (for suspend functions/lambdas),
+ * or `Any` (for simple lamdbas), or a custom superclass (in case of SAM conversion).
+ *
+ * For example, `::bar$lambda$0<BarTP>` in the following code:
+ * ```kotlin
+ * fun <FooTP> foo(v: FooTP, l: (FooTP) -> String): String {
+ *     return l(v)
+ * }
+ *
+ * private fun <T> bar$lambda$0(t: T): String { /* ... */ }
+ *
+ * fun <BarTP> bar(v: BarTP): String {
+ *     return foo(v, ::bar$lambda$0<BarTP>)
+ * }
+ * ```
+ *
+ * is lowered into:
+ * ```kotlin
+ * private class bar$lambda$0$FUNCTION_REFERENCE$0<T> : kotlin.native.internal.KFunctionImpl<String>, kotlin.Function1<T, String> {
+ *     override fun invoke(p1: T): String {
+ *         return bar$lambda$0<T>(p1)
+ *     }
+ * }
+ *
+ * fun <BarTP> bar(v: BarTP): String {
+ *     return foo(v, bar$lambda$0$FUNCTION_REFERENCE$0<BarTP>())
+ * }
+ * ```
+ */
 internal class FunctionReferenceLowering(val generationState: NativeGenerationState) : FileLoweringPass {
     private object DECLARATION_ORIGIN_FUNCTION_REFERENCE_IMPL : IrDeclarationOriginImpl("FUNCTION_REFERENCE_IMPL")
 
@@ -173,18 +203,6 @@ internal class FunctionReferenceLowering(val generationState: NativeGenerationSt
         private val boundFunctionParameters = functionReference.getArgumentsWithIr().map { it.first }
         private val unboundFunctionParameters = functionParameters - boundFunctionParameters
 
-        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()
         private val isKSuspendFunction = functionReference.type.isKSuspendFunction()
@@ -201,12 +219,21 @@ internal class FunctionReferenceLowering(val generationState: NativeGenerationSt
             visibility = DescriptorVisibilities.PRIVATE
         }.apply {
             parent = this@FunctionReferenceBuilder.parent
+            copyTypeParametersFrom(referencedFunction)
             createParameterDeclarations()
 
             // copy the generated name for IrClass, partially solves KT-47194
             generationState.copyLocalClassName(functionReference, this)
         }
 
+        private val typeArguments = functionReferenceClass.typeParameters.map { typeParam ->
+            typeParam.symbol to functionReference.getTypeArgument(typeParam.index)!!
+        }
+
+        private val typeParameterRemapper = IrTypeParameterRemapper(referencedFunction.typeParameters.zip(functionReferenceClass.typeParameters).toMap())
+        private val functionParameterTypes = unboundFunctionParameters.map { typeParameterRemapper.remapType(it.type) }
+        private val functionReturnType = typeParameterRemapper.remapType(referencedFunction.returnType)
+
         private val functionReferenceThis = functionReferenceClass.thisReceiver!!
 
         private val argumentToPropertiesMap = boundFunctionParameters.associateWith {
@@ -220,8 +247,6 @@ internal class FunctionReferenceLowering(val generationState: NativeGenerationSt
             }
         }
 
-        private fun IrClass.getInvokeFunction() = simpleFunctions().single { it.name.asString() == "invoke" }
-
         private val kFunctionImplSymbol = symbols.kFunctionImpl
         private val kFunctionImplConstructorSymbol = kFunctionImplSymbol.constructors.single()
         private val kSuspendFunctionImplSymbol = symbols.kSuspendFunctionImpl
@@ -241,14 +266,15 @@ internal class FunctionReferenceLowering(val generationState: NativeGenerationSt
                 transformedSuperMethod = samSuperClass.functions.single { it.owner.modality == Modality.ABSTRACT }.owner
             } else {
                 val numberOfParameters = unboundFunctionParameters.size
+                val functionParameterAndReturnTypes = functionParameterTypes + functionReturnType
                 if (isKSuspendFunction) {
                     val suspendFunctionClass = symbols.kSuspendFunctionN(numberOfParameters).owner
                     superTypes += suspendFunctionClass.typeWith(functionParameterAndReturnTypes)
-                    transformedSuperMethod = suspendFunctionClass.getInvokeFunction()
+                    transformedSuperMethod = suspendFunctionClass.invokeFun!!
                 } else {
                     val functionClass = (if (isKFunction) symbols.kFunctionN(numberOfParameters) else symbols.functionN(numberOfParameters)).owner
                     superTypes += functionClass.typeWith(functionParameterAndReturnTypes)
-                    transformedSuperMethod = functionClass.getInvokeFunction()
+                    transformedSuperMethod = functionClass.invokeFun!!
                 }
             }
             val originalSuperMethod = context.mapping.functionWithContinuationsToSuspendFunctions[transformedSuperMethod] ?: transformedSuperMethod
@@ -314,6 +340,8 @@ internal class FunctionReferenceLowering(val generationState: NativeGenerationSt
                     ignoredParentSymbols = listOf(transformedSuperMethod.symbol)
             )
 
+            functionReferenceClass.remapTypes(typeParameterRemapper)
+
             return functionReferenceClass
         }
 
@@ -323,6 +351,7 @@ internal class FunctionReferenceLowering(val generationState: NativeGenerationSt
             origin = DECLARATION_ORIGIN_FUNCTION_REFERENCE_IMPL
             isPrimary = true
         }.apply {
+            val typeArgumentsMap = typeArguments.toMap()
             valueParameters += boundFunctionParameters.mapIndexed { index, parameter ->
                 parameter.copyTo(this, DECLARATION_ORIGIN_FUNCTION_REFERENCE_IMPL, index,
                         type = parameter.type.substitute(typeArgumentsMap))
@@ -347,10 +376,11 @@ internal class FunctionReferenceLowering(val generationState: NativeGenerationSt
             val clazz = buildClass()
             val constructor = buildConstructor()
             val arguments = functionReference.getArgumentsWithIr()
+            val typeArguments = typeArguments.map { it.second }
             val expression = if (arguments.isEmpty()) {
-                irBuilder.irConstantObject(clazz, emptyMap())
+                irBuilder.irConstantObject(clazz, emptyMap(), typeArguments)
             } else {
-                irBuilder.irCall(constructor).apply {
+                irBuilder.irCallConstructor(constructor.symbol, typeArguments).apply {
                     arguments.forEachIndexed { index, argument ->
                         putValueArgument(index, argument.second)
                     }
@@ -450,7 +480,7 @@ internal class FunctionReferenceLowering(val generationState: NativeGenerationSt
                             assert(unboundIndex == valueParameters.size) { "Not all arguments of <invoke> are used" }
 
                             referencedFunction.typeParameters.forEach { typeParam ->
-                                putTypeArgument(typeParam.index, functionReference.getTypeArgument(typeParam.index)!!)
+                                putTypeArgument(typeParam.index, functionReferenceClass.typeParameters[typeParam.index].defaultType)
                             }
                         }
                 )

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

@@ -4780,6 +4780,12 @@ public class K2NativeCodegenBoxTestGenerated extends AbstractNativeCodegenBoxTes
                 runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
             }
 
+            @Test
+            @TestMetadata("closureCapturingGenericParam.kt")
+            public void testClosureCapturingGenericParam() throws Exception {
+                runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+            }
+
             @Test
             @TestMetadata("closureInsideConstrucor.kt")
             public void testClosureInsideConstrucor() throws Exception {

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

@@ -4729,6 +4729,12 @@ public class NativeCodegenBoxTestGenerated extends AbstractNativeCodegenBoxTest
                 runTest("compiler/testData/codegen/box/closures/capturedLocalGenericFun.kt");
             }
 
+            @Test
+            @TestMetadata("closureCapturingGenericParam.kt")
+            public void testClosureCapturingGenericParam() throws Exception {
+                runTest("compiler/testData/codegen/box/closures/closureCapturingGenericParam.kt");
+            }
+
             @Test
             @TestMetadata("closureInsideConstrucor.kt")
             public void testClosureInsideConstrucor() throws Exception {