FIR: improve inference of implicit type arguments in casts

Type parameters do not necessarily match one-to-one, or preserve order.
2021-01-06 14:39:22 +01:00
parent acdc1f532b
commit 29f95c7df2
32 changed files with 134 additions and 149 deletions
@@ -74,7 +74,7 @@ class ConeIllegalAnnotationError(val name: Name) : ConeDiagnostic() {
    override val reason: String get() = "Not a legal annotation: $name"
 }

-class ConeWrongNumberOfTypeArgumentsError(val desiredCount: Int, val type: FirRegularClassSymbol) : ConeDiagnostic() {
+class ConeWrongNumberOfTypeArgumentsError(val desiredCount: Int, val type: FirClassLikeSymbol<*>) : ConeDiagnostic() {
    override val reason: String get() = "Wrong number of type arguments"
 }

@@ -450,30 +450,53 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
        override fun <T> shouldRunCompletion(call: T): Boolean where T : FirStatement, T : FirResolvable = false
    }

-    private fun FirTypeRef.withTypeArgumentsForBareType(argument: FirExpression): FirTypeRef {
-        // TODO: Everything should also work for case of checked-type itself is a type alias
-        val type = coneTypeSafe<ConeKotlinType>()
-        if (type !is ConeClassLikeType || type.typeArguments.isNotEmpty()) {
-            return this
+    private fun ConeClassLikeType.inheritTypeArguments(
+        base: FirClassLikeDeclaration<*>,
+        arguments: Array<out ConeTypeProjection>
+    ): Array<out ConeTypeProjection>? {
+        val firClass = lookupTag.toSymbol(session)?.fir ?: return null
+        if (firClass !is FirTypeParameterRefsOwner || firClass.typeParameters.isEmpty()) return arrayOf()
+        return when (firClass) {
+            base -> arguments
+            is FirTypeAlias -> firClass.inheritTypeArguments(firClass.expandedTypeRef, base, arguments)
+            // TODO: if many supertypes, check consistency
+            is FirClass<*> -> firClass.superTypeRefs.mapNotNull { firClass.inheritTypeArguments(it, base, arguments) }.firstOrNull()
+            else -> null
        }
-        val baseTypeArguments =
-            argument.typeRef.coneTypeSafe<ConeKotlinType>()?.fullyExpandedType(session)?.typeArguments
+    }

-        return if (baseTypeArguments?.isEmpty() != false) {
-            this
-        } else {
-            val typeParameters = (type.lookupTag.toSymbol(session)?.fir as? FirTypeParameterRefsOwner)?.typeParameters.orEmpty()
-            if (typeParameters.isEmpty()) {
-                this
-            } else {
-                withReplacedConeType(
-                    type.withArguments(
-                        if (baseTypeArguments.size > typeParameters.size) baseTypeArguments.take(typeParameters.size).toTypedArray()
-                        else baseTypeArguments
-                    )
-                )
-            }
+    private fun FirTypeParameterRefsOwner.inheritTypeArguments(
+        typeRef: FirTypeRef,
+        base: FirClassLikeDeclaration<*>,
+        arguments: Array<out ConeTypeProjection>
+    ): Array<out ConeTypeProjection>? {
+        val type = typeRef.coneTypeSafe<ConeClassLikeType>() ?: return null
+        val indexMapping = typeParameters.map { parameter ->
+            // TODO: if many, check consistency of the result
+            type.typeArguments.indexOfFirst { it is ConeTypeParameterType && it.lookupTag.typeParameterSymbol == parameter.symbol }
        }
+        if (indexMapping.any { it == -1 }) return null
+
+        val typeArguments = type.inheritTypeArguments(base, arguments) ?: return null
+        return Array(typeParameters.size) { typeArguments[indexMapping[it]] }
+    }
+
+    private fun FirTypeRef.withTypeArgumentsForBareType(argument: FirExpression): FirTypeRef {
+        val type = coneTypeSafe<ConeClassLikeType>() ?: return this
+        if (type.typeArguments.isNotEmpty()) return this
+
+        val firClass = type.lookupTag.toSymbol(session)?.fir ?: return this
+        if (firClass !is FirTypeParameterRefsOwner || firClass.typeParameters.isEmpty()) return this
+
+        val baseType = argument.typeRef.coneTypeSafe<ConeClassLikeType>()?.fullyExpandedType(session) ?: return this
+        val baseFirClass = baseType.lookupTag.toSymbol(session)?.fir ?: return this
+
+        val newArguments = type.inheritTypeArguments(baseFirClass, baseType.typeArguments)
+            ?: return buildErrorTypeRef {
+                source = this@withTypeArgumentsForBareType.source
+                diagnostic = ConeWrongNumberOfTypeArgumentsError(firClass.typeParameters.size, firClass.symbol)
+            }
+        return if (newArguments.isEmpty()) this else withReplacedConeType(type.withArguments(newArguments))
    }

    override fun transformTypeOperatorCall(