[IR] Fix nullability of substituted type arguments in IrTypeSubstitutor

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/types/IrTypeSubstitutor.kt

@@ -27,51 +27,39 @@ abstract class AbstractIrTypeSubstitutor(private val irBuiltIns: IrBuiltIns) : T
 
     fun substitute(type: IrType): IrType {
         if (isEmptySubstitution()) return type
-
+        return substituteType(type)
-        return type.typeParameterConstructor()?.let {
-            when (val typeArgument = getSubstitutionArgument(it)) {
-                is IrStarProjection -> irBuiltIns.anyNType // TODO upper bound for T
-                is IrTypeProjection -> typeArgument.type.run { if (type.isMarkedNullable()) makeNullable() else this }
-            }
-        } ?: substituteType(type)
     }
 
     private fun substituteType(irType: IrType): IrType {
+        val substitutedTypeParameter = irType.typeParameterConstructor()?.let {
+            when (val typeArgument = getSubstitutionArgument(it)) {
+                is IrStarProjection -> irBuiltIns.anyNType // TODO upper bound for T
+                is IrTypeProjection -> typeArgument.type.mergeNullability(irType)
+            }
+        }
+        if (substitutedTypeParameter != null) {
+            return substitutedTypeParameter
+        }
+
         return when (irType) {
-            is IrSimpleType ->
+            is IrSimpleType -> with(irType.toBuilder()) {
-                with(irType.toBuilder()) {
+                arguments = irType.arguments.memoryOptimizedMap { substituteTypeArgument(it) }
-                    arguments = irType.arguments.memoryOptimizedMap { substituteTypeArgument(it) }
+                buildSimpleType()
-                    buildSimpleType()
+            }
-                }
             is IrDynamicType,
-            is IrErrorType ->
+            is IrErrorType -> irType
-                irType
+            else -> error("Unexpected type: $irType")
-            else ->
-                throw AssertionError("Unexpected type: $irType")
         }
     }
 
     private fun substituteTypeArgument(typeArgument: IrTypeArgument): IrTypeArgument {
-        when (typeArgument) {
+        return when (typeArgument) {
-            is IrStarProjection -> return typeArgument
+            is IrStarProjection -> typeArgument
-            is IrTypeProjection -> {
+            is IrTypeProjection -> makeTypeProjection(substituteType(typeArgument.type), typeArgument.variance)
-                val type = typeArgument.type
-                if (type is IrSimpleType) {
-                    val classifier = type.classifier
-                    if (classifier is IrTypeParameterSymbol) {
-                        val newArgument = getSubstitutionArgument(classifier)
-                        return if (newArgument is IrTypeProjection) {
-                            makeTypeProjection(newArgument.type, typeArgument.variance)
-                        } else newArgument
-                    }
-                }
-                return makeTypeProjection(substituteType(typeArgument.type), typeArgument.variance)
-            }
         }
     }
 }
 
-
 class IrTypeSubstitutor(
     typeParameters: List<IrTypeParameterSymbol>,
     typeArguments: List<IrTypeArgument>,