IR: type parameter context for return values in InsertImplicitCasts

compiler/fir/fir2ir/tests/org/jetbrains/kotlin/fir/Fir2IrTextTestGenerated.java

@@ -1735,6 +1735,11 @@ public class Fir2IrTextTestGenerated extends AbstractFir2IrTextTest {
             runTest("compiler/testData/ir/irText/regressions/typeAliasCtorForGenericClass.kt");
         }
 
+        @TestMetadata("typeParametersInImplicitCast.kt")
+        public void testTypeParametersInImplicitCast() throws Exception {
+            runTest("compiler/testData/ir/irText/regressions/typeParametersInImplicitCast.kt");
+        }
+
         @TestMetadata("compiler/testData/ir/irText/regressions/newInference")
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/transformations/InsertImplicitCasts.kt

@@ -36,6 +36,8 @@ import org.jetbrains.kotlin.ir.types.impl.originalKotlinType
 import org.jetbrains.kotlin.ir.util.TypeTranslator
 import org.jetbrains.kotlin.ir.util.coerceToUnit
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.psi2ir.containsNull
 import org.jetbrains.kotlin.psi2ir.generators.GeneratorContext
@@ -64,19 +66,37 @@ internal class InsertImplicitCasts(
 ) : IrElementTransformerVoid() {
 
     private val expectedFunctionExpressionReturnType = hashMapOf<FunctionDescriptor, IrType>()
-    private val returnExpressionsToBePostprocessed = arrayListOf<IrReturn>()
 
     fun run(element: IrElement) {
         element.transformChildrenVoid(this)
-        for (irReturn in returnExpressionsToBePostprocessed) {
-            postprocessReturnExpression(irReturn)
-        }
+        postprocessReturnExpressions(element)
     }
 
-    private fun postprocessReturnExpression(irReturn: IrReturn) {
-        val returnTarget = irReturn.returnTarget
-        val expectedReturnType = expectedFunctionExpressionReturnType[returnTarget] ?: return
-        irReturn.value = irReturn.value.cast(expectedReturnType)
+    private fun postprocessReturnExpressions(element: IrElement) {
+        // We need to re-create type parameter context for casts of postprocessed return values.
+        element.acceptChildrenVoid(object : IrElementVisitorVoid {
+            override fun visitReturn(expression: IrReturn) {
+                super.visitReturn(expression)
+                val expectedReturnType = expectedFunctionExpressionReturnType[expression.returnTarget] ?: return
+                expression.value = expression.value.cast(expectedReturnType)
+            }
+
+            override fun visitClass(declaration: IrClass) {
+                typeTranslator.buildWithScope(declaration) {
+                    super.visitClass(declaration)
+                }
+            }
+
+            override fun visitFunction(declaration: IrFunction) {
+                typeTranslator.buildWithScope(declaration) {
+                    super.visitFunction(declaration)
+                }
+            }
+
+            override fun visitElement(element: IrElement) {
+                element.acceptChildrenVoid(this)
+            }
+        })
     }
 
     private fun KotlinType.toIrType() = typeTranslator.translateType(this)
@@ -172,9 +192,6 @@ internal class InsertImplicitCasts(
             } else {
                 val returnTargetDescriptor = expression.returnTarget
                 val isLambdaReturnValue = returnTargetDescriptor is AnonymousFunctionDescriptor
-                if (isLambdaReturnValue) {
-                    returnExpressionsToBePostprocessed.add(expression)
-                }
                 value.cast(returnTargetDescriptor.returnType, isLambdaReturnValue = isLambdaReturnValue)
             }
         }

compiler/testData/ir/irText/regressions/typeParametersInImplicitCast.fir.txt

@@ -0,0 +1,18 @@
+FILE fqName:<root> fileName:/typeParametersInImplicitCast.kt
+  FUN name:problematic visibility:public modality:FINAL <T> (lss:kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>>) returnType:kotlin.collections.List<T of <root>.problematic>
+    TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?]
+    VALUE_PARAMETER name:lss index:0 type:kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>>
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun problematic <T> (lss: kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>>): kotlin.collections.List<T of <root>.problematic> declared in <root>'
+        CALL 'public final fun flatMap <T, R> (transform: kotlin.Function1<T of kotlin.collections.flatMap, kotlin.collections.Iterable<R of kotlin.collections.flatMap>>): kotlin.collections.List<R of kotlin.collections.flatMap> [inline] declared in kotlin.collections' type=kotlin.collections.List<T of <root>.problematic> origin=null
+          <T>: kotlin.collections.List<T of <root>.problematic>
+          <R>: T of <root>.problematic
+          $receiver: GET_VAR 'lss: kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>> declared in <root>.problematic' type=kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>> origin=null
+          transform: FUN_EXPR type=kotlin.Function1<kotlin.collections.List<T of <root>.problematic>, kotlin.collections.Iterable<T of <root>.problematic>> origin=LAMBDA
+            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.collections.List<T of <root>.problematic>) returnType:kotlin.collections.Iterable<T of <root>.problematic>
+              VALUE_PARAMETER name:it index:0 type:kotlin.collections.List<T of <root>.problematic>
+              BLOCK_BODY
+                RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.collections.List<T of <root>.problematic>): kotlin.collections.Iterable<T of <root>.problematic> declared in <root>.problematic'
+                  CALL 'public/*package*/ open fun id <T> (v: kotlin.collections.List<T of <root>.ListId.id?>?): kotlin.collections.List<T of <root>.ListId.id?> declared in <root>.ListId' type=kotlin.collections.List<T of <root>.problematic?> origin=null
+                    <T>: T of <root>.problematic?
+                    v: GET_VAR 'it: kotlin.collections.List<T of <root>.problematic> declared in <root>.problematic.<anonymous>' type=kotlin.collections.List<T of <root>.problematic> origin=null

compiler/testData/ir/irText/regressions/typeParametersInImplicitCast.kt

@@ -0,0 +1,15 @@
+// WITH_RUNTIME
+// FILE: ListId.java
+import java.util.List;
+import org.jetbrains.annotations.NotNull;
+
+class ListId {
+    @NotNull
+    static <T> List<T> id(List<T> v) {
+        return v;
+    }
+}
+
+// FILE: typeParametersInImplicitCast.kt
+
+fun <T> problematic(lss: List<List<T>>): List<T> = lss.flatMap { ListId.id(it) }

compiler/testData/ir/irText/regressions/typeParametersInImplicitCast.txt

@@ -0,0 +1,19 @@
+FILE fqName:<root> fileName:/typeParametersInImplicitCast.kt
+  FUN name:problematic visibility:public modality:FINAL <T> (lss:kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>>) returnType:kotlin.collections.List<T of <root>.problematic>
+    TYPE_PARAMETER name:T index:0 variance: superTypes:[kotlin.Any?]
+    VALUE_PARAMETER name:lss index:0 type:kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>>
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='public final fun problematic <T> (lss: kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>>): kotlin.collections.List<T of <root>.problematic> declared in <root>'
+        CALL 'public final fun flatMap <T, R> (transform: kotlin.Function1<T of kotlin.collections.flatMap, kotlin.collections.Iterable<R of kotlin.collections.flatMap>>): kotlin.collections.List<R of kotlin.collections.flatMap> [inline] declared in kotlin.collections' type=kotlin.collections.List<T of <root>.problematic?> origin=null
+          <T>: kotlin.collections.List<T of <root>.problematic>
+          <R>: T of <root>.problematic?
+          $receiver: GET_VAR 'lss: kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>> declared in <root>.problematic' type=kotlin.collections.List<kotlin.collections.List<T of <root>.problematic>> origin=null
+          transform: FUN_EXPR type=kotlin.Function1<kotlin.collections.List<T of <root>.problematic>, kotlin.collections.Iterable<T of <root>.problematic?>> origin=LAMBDA
+            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.collections.List<T of <root>.problematic>) returnType:kotlin.collections.Iterable<T of <root>.problematic?>
+              VALUE_PARAMETER name:it index:0 type:kotlin.collections.List<T of <root>.problematic>
+              BLOCK_BODY
+                RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.collections.List<T of <root>.problematic>): kotlin.collections.Iterable<T of <root>.problematic?> declared in <root>.problematic'
+                  TYPE_OP type=kotlin.collections.List<T of <root>.problematic?> origin=IMPLICIT_NOTNULL typeOperand=kotlin.collections.List<T of <root>.problematic?>
+                    CALL 'public/*package*/ open fun id <T> (v: kotlin.collections.List<T of <root>.ListId.id?>?): kotlin.collections.List<T of <root>.ListId.id?> declared in <root>.ListId' type=kotlin.collections.List<T of <root>.problematic?> origin=null
+                      <T>: T of <root>.problematic?
+                      v: GET_VAR 'it: kotlin.collections.List<T of <root>.problematic> declared in <root>.problematic.<anonymous>' type=kotlin.collections.List<T of <root>.problematic> origin=null

compiler/tests/org/jetbrains/kotlin/ir/IrTextTestCaseGenerated.java

@@ -1734,6 +1734,11 @@ public class IrTextTestCaseGenerated extends AbstractIrTextTestCase {
             runTest("compiler/testData/ir/irText/regressions/typeAliasCtorForGenericClass.kt");
         }
 
+        @TestMetadata("typeParametersInImplicitCast.kt")
+        public void testTypeParametersInImplicitCast() throws Exception {
+            runTest("compiler/testData/ir/irText/regressions/typeParametersInImplicitCast.kt");
+        }
+
         @TestMetadata("compiler/testData/ir/irText/regressions/newInference")
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)