Merge pull request #9405 from smowton/smowton/fix/restore-wildcard-types

Kotlin: Introduce / restore implied wildcard types

java/kotlin-extractor/src/main/kotlin/KotlinFileExtractor.kt

@@ -19,7 +19,11 @@ import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.load.java.sources.JavaSourceElement
+import org.jetbrains.kotlin.load.java.structure.JavaClass
+import org.jetbrains.kotlin.load.java.structure.JavaTypeParameter
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import java.io.Closeable
 import java.util.*
@@ -171,7 +175,7 @@ open class KotlinFileExtractor(
         }
     }
 
-    fun extractTypeParameter(tp: IrTypeParameter, apparentIndex: Int): Label<out DbTypevariable>? {
+    fun extractTypeParameter(tp: IrTypeParameter, apparentIndex: Int, javaTypeParameter: JavaTypeParameter?): Label<out DbTypevariable>? {
         with("type parameter", tp) {
             val parentId = getTypeParameterParentLabel(tp) ?: return null
             val id = tw.getLabelFor<DbTypevariable>(getTypeParameterLabel(tp))
@@ -183,10 +187,21 @@ open class KotlinFileExtractor(
             val locId = tw.getLocation(tp)
             tw.writeHasLocation(id, locId)
 
+            // Annoyingly, we have no obvious way to pair up the bounds of an IrTypeParameter and a JavaTypeParameter
+            // because JavaTypeParameter provides a Collection not an ordered list, so we can only do our best here:
+            fun tryGetJavaBound(idx: Int) =
+                when(tp.superTypes.size) {
+                    1 -> javaTypeParameter?.upperBounds?.singleOrNull()
+                    else -> (javaTypeParameter?.upperBounds as? List)?.getOrNull(idx)
+                }
+
             tp.superTypes.forEachIndexed { boundIdx, bound ->
                 if(!(bound.isAny() || bound.isNullableAny())) {
                     tw.getLabelFor<DbTypebound>("@\"bound;$boundIdx;{$id}\"") {
-                        tw.writeTypeBounds(it, useType(bound).javaResult.id.cast<DbReftype>(), boundIdx, id)
+                        // Note we don't look for @JvmSuppressWildcards here because it doesn't seem to have any impact
+                        // on kotlinc adding wildcards to type parameter bounds.
+                        val boundWithWildcards = addJavaLoweringWildcards(bound, true, tryGetJavaBound(tp.index))
+                        tw.writeTypeBounds(it, useType(boundWithWildcards).javaResult.id.cast<DbReftype>(), boundIdx, id)
                     }
                 }
             }
@@ -382,7 +397,9 @@ open class KotlinFileExtractor(
 
                 extractEnclosingClass(c, id, locId, listOf())
 
-                c.typeParameters.mapIndexed { idx, param -> extractTypeParameter(param, idx) }
+                val javaClass = (c.source as? JavaSourceElement)?.javaElement as? JavaClass
+
+                c.typeParameters.mapIndexed { idx, param -> extractTypeParameter(param, idx, javaClass?.typeParameters?.getOrNull(idx)) }
                 if (extractDeclarations) {
                     c.declarations.map { extractDeclaration(it, extractPrivateMembers = extractPrivateMembers, extractFunctionBodies = extractFunctionBodies) }
                     if (extractStaticInitializer)
@@ -497,7 +514,9 @@ open class KotlinFileExtractor(
                 else
                     null
             } ?: vp.type
-            val substitutedType = typeSubstitution?.let { it(maybeErasedType, TypeContext.OTHER, pluginContext) } ?: maybeErasedType
+            val javaType = ((vp.parent as? IrFunction)?.let { getJavaMethod(it) })?.valueParameters?.getOrNull(idx)?.type
+            val typeWithWildcards = addJavaLoweringWildcards(maybeErasedType, !hasWildcardSuppressionAnnotation(vp), javaType)
+            val substitutedType = typeSubstitution?.let { it(typeWithWildcards, TypeContext.OTHER, pluginContext) } ?: typeWithWildcards
             val id = useValueParameter(vp, parent)
             if (extractTypeAccess) {
                 extractTypeAccessRecursive(substitutedType, location, id, -1)
@@ -531,7 +550,9 @@ open class KotlinFileExtractor(
                     extensionReceiverParameter = null,
                     functionTypeParameters = listOf(),
                     classTypeArgsIncludingOuterClasses = listOf(),
-                    overridesCollectionsMethod = false
+                    overridesCollectionsMethod = false,
+                    javaSignature = null,
+                    addParameterWildcardsByDefault = false
                 )
                 val clinitId = tw.getLabelFor<DbMethod>(clinitLabel)
                 val returnType = useType(pluginContext.irBuiltIns.unitType, TypeContext.RETURN)
@@ -670,7 +691,8 @@ open class KotlinFileExtractor(
         with("function", f) {
             DeclarationStackAdjuster(f).use {
 
-                getFunctionTypeParameters(f).mapIndexed { idx, tp -> extractTypeParameter(tp, idx) }
+                val javaMethod = getJavaMethod(f)
+                getFunctionTypeParameters(f).mapIndexed { idx, tp -> extractTypeParameter(tp, idx, javaMethod?.typeParameters?.getOrNull(idx)) }
 
                 val id =
                     idOverride
@@ -704,7 +726,7 @@ open class KotlinFileExtractor(
 
                 val paramsSignature = allParamTypes.joinToString(separator = ",", prefix = "(", postfix = ")") { it.javaResult.signature!! }
 
-                val adjustedReturnType = getAdjustedReturnType(f)
+                val adjustedReturnType = addJavaLoweringWildcards(getAdjustedReturnType(f), false, javaMethod?.returnType)
                 val substReturnType = typeSubstitution?.let { it(adjustedReturnType, TypeContext.RETURN, pluginContext) } ?: adjustedReturnType
 
                 val locId = locOverride ?: getLocation(f, classTypeArgsIncludingOuterClasses)
@@ -3744,6 +3766,17 @@ open class KotlinFileExtractor(
         }
     }
 
+    /**
+     * Extracts a single wildcard type access expression with no enclosing callable and statement.
+     */
+    private fun extractWildcardTypeAccess(type: TypeResults, location: Label<DbLocation>, parent: Label<out DbExprparent>, idx: Int): Label<out DbExpr> {
+        val id = tw.getFreshIdLabel<DbWildcardtypeaccess>()
+        tw.writeExprs_wildcardtypeaccess(id, type.javaResult.id, parent, idx)
+        tw.writeExprsKotlinType(id, type.kotlinResult.id)
+        tw.writeHasLocation(id, location)
+        return id
+    }
+
     /**
      * Extracts a single type access expression with no enclosing callable and statement.
      */
@@ -3768,6 +3801,27 @@ open class KotlinFileExtractor(
         return id
     }
 
+    /**
+     * Extracts a type argument type access, introducing a wildcard type access if appropriate, or directly calling
+     * `extractTypeAccessRecursive` if the argument is invariant.
+     * No enclosing callable and statement is extracted, this is useful for type access extraction in field declarations.
+     */
+    private fun extractWildcardTypeAccessRecursive(t: IrTypeArgument, location: Label<DbLocation>, parent: Label<out DbExprparent>, idx: Int) {
+        val typeLabels by lazy { TypeResults(getTypeArgumentLabel(t), TypeResult(fakeKotlinType(), "TODO", "TODO")) }
+        when (t) {
+            is IrStarProjection -> extractWildcardTypeAccess(typeLabels, location, parent, idx)
+            is IrTypeProjection -> when(t.variance) {
+                Variance.INVARIANT -> extractTypeAccessRecursive(t.type, location, parent, idx, TypeContext.GENERIC_ARGUMENT)
+                else -> {
+                    val wildcardLabel = extractWildcardTypeAccess(typeLabels, location, parent, idx)
+                    // Mimic a Java extractor oddity, that it uses the child index to indicate what kind of wildcard this is
+                    val boundChildIdx = if (t.variance == Variance.OUT_VARIANCE) 0 else 1
+                    extractTypeAccessRecursive(t.type, location, wildcardLabel, boundChildIdx, TypeContext.GENERIC_ARGUMENT)
+                }
+            }
+        }
+    }
+
     /**
      * Extracts a type access expression and its child type access expressions in case of a generic type. Nested generics are also handled.
      * No enclosing callable and statement is extracted, this is useful for type access extraction in field declarations.
@@ -3775,8 +3829,8 @@ open class KotlinFileExtractor(
     private fun extractTypeAccessRecursive(t: IrType, location: Label<DbLocation>, parent: Label<out DbExprparent>, idx: Int, typeContext: TypeContext = TypeContext.OTHER): Label<out DbExpr> {
         val typeAccessId = extractTypeAccess(useType(t, typeContext), location, parent, idx)
         if (t is IrSimpleType) {
-            t.arguments.filterIsInstance<IrType>().forEachIndexed { argIdx, arg ->
-                extractTypeAccessRecursive(arg, location, typeAccessId, argIdx, TypeContext.GENERIC_ARGUMENT)
+            t.arguments.forEachIndexed { argIdx, arg ->
+                extractWildcardTypeAccessRecursive(arg, location, typeAccessId, argIdx)
             }
         }
         return typeAccessId

java/kotlin-extractor/src/main/kotlin/KotlinUsesExtractor.kt

@@ -6,6 +6,8 @@ import com.github.codeql.utils.versions.isRawType
 import com.semmle.extractor.java.OdasaOutput
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.backend.common.ir.allOverridden
+import org.jetbrains.kotlin.backend.common.ir.isFinalClass
+import org.jetbrains.kotlin.backend.common.lower.parents
 import org.jetbrains.kotlin.backend.common.lower.parentsWithSelf
 import org.jetbrains.kotlin.backend.jvm.ir.getJvmNameFromAnnotation
 import org.jetbrains.kotlin.backend.jvm.ir.propertyIfAccessor
@@ -20,6 +22,8 @@ import org.jetbrains.kotlin.ir.types.impl.*
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.load.java.BuiltinMethodsWithSpecialGenericSignature
 import org.jetbrains.kotlin.load.java.JvmAbi
+import org.jetbrains.kotlin.load.java.sources.JavaSourceElement
+import org.jetbrains.kotlin.load.java.structure.*
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
@@ -850,6 +854,64 @@ open class KotlinUsesExtractor(
                 (f.name.asString() == "addAll" && overridesFunctionDefinedOn(f, "kotlin.collections", "MutableCollection")) ||
                 (f.name.asString() == "addAll" && overridesFunctionDefinedOn(f, "kotlin.collections", "MutableList"))
 
+
+    private val jvmWildcardAnnotation = FqName("kotlin.jvm.JvmWildcard")
+    private val jvmWildcardSuppressionAnnotaton = FqName("kotlin.jvm.JvmSuppressWildcards")
+
+    private fun wildcardAdditionAllowed(v: Variance, t: IrType, addByDefault: Boolean) =
+        when {
+            t.hasAnnotation(jvmWildcardAnnotation) -> true
+            !addByDefault -> false
+            t.hasAnnotation(jvmWildcardSuppressionAnnotaton) -> false
+            v == Variance.IN_VARIANCE -> !(t.isNullableAny() || t.isAny())
+            v == Variance.OUT_VARIANCE -> ((t as? IrSimpleType)?.classOrNull?.owner?.isFinalClass) != true
+            else -> false
+        }
+
+    private fun addJavaLoweringArgumentWildcards(p: IrTypeParameter, t: IrTypeArgument, addByDefault: Boolean, javaType: JavaType?): IrTypeArgument =
+        (t as? IrTypeProjection)?.let {
+            val newBase = addJavaLoweringWildcards(it.type, addByDefault, javaType)
+            val newVariance =
+                if (it.variance == Variance.INVARIANT &&
+                    p.variance != Variance.INVARIANT &&
+                    // The next line forbids inferring a wildcard type when we have a corresponding Java type with conflicting variance.
+                    // For example, Java might declare f(Comparable<CharSequence> cs), in which case we shouldn't add a `? super ...`
+                    // wildcard. Note if javaType is unknown (e.g. this is a Kotlin source element), we assume wildcards should be added.
+                    (javaType?.let { jt -> jt is JavaWildcardType && jt.isExtends == (p.variance == Variance.OUT_VARIANCE) } != false) &&
+                    wildcardAdditionAllowed(p.variance, it.type, addByDefault))
+                    p.variance
+                else
+                    it.variance
+            if (newBase !== it.type || newVariance != it.variance)
+                makeTypeProjection(newBase, newVariance)
+            else
+                null
+        } ?: t
+
+    fun getJavaTypeArgument(jt: JavaType, idx: Int) =
+        when(jt) {
+            is JavaClassifierType -> jt.typeArguments.getOrNull(idx)
+            is JavaArrayType -> if (idx == 0) jt.componentType else null
+            else -> null
+        }
+
+    fun addJavaLoweringWildcards(t: IrType, addByDefault: Boolean, javaType: JavaType?): IrType =
+        (t as? IrSimpleType)?.let {
+            val typeParams = it.classOrNull?.owner?.typeParameters ?: return t
+            val newArgs = typeParams.zip(it.arguments).mapIndexed { idx, pair ->
+                addJavaLoweringArgumentWildcards(
+                    pair.first,
+                    pair.second,
+                    addByDefault,
+                    javaType?.let { jt -> getJavaTypeArgument(jt, idx) }
+                )
+            }
+            return if (newArgs.zip(it.arguments).all { pair -> pair.first === pair.second })
+                t
+            else
+                it.toBuilder().also { builder -> builder.arguments = newArgs }.buildSimpleType()
+        } ?: t
+
     /*
      * This is the normal getFunctionLabel function to use. If you want
      * to refer to the function in its source class then
@@ -883,6 +945,14 @@ open class KotlinUsesExtractor(
         return otherKeySet.returnType.codeQlWithHasQuestionMark(false)
     }
 
+    @OptIn(ObsoleteDescriptorBasedAPI::class)
+    fun getJavaMethod(f: IrFunction) = (f.descriptor.source as? JavaSourceElement)?.javaElement as? JavaMethod
+
+    fun hasWildcardSuppressionAnnotation(d: IrDeclaration) =
+        d.hasAnnotation(jvmWildcardSuppressionAnnotaton) ||
+        // Note not using `parentsWithSelf` as that only works if `d` is an IrDeclarationParent
+        d.parents.any { (it as? IrAnnotationContainer)?.hasAnnotation(jvmWildcardSuppressionAnnotaton) == true }
+
     /*
      * There are some pairs of classes (e.g. `kotlin.Throwable` and
      * `java.lang.Throwable`) which are really just 2 different names
@@ -903,7 +973,9 @@ open class KotlinUsesExtractor(
             f.extensionReceiverParameter,
             getFunctionTypeParameters(f),
             classTypeArgsIncludingOuterClasses,
-            overridesCollectionsMethodWithAlteredParameterTypes(f)
+            overridesCollectionsMethodWithAlteredParameterTypes(f),
+            getJavaMethod(f),
+            !hasWildcardSuppressionAnnotation(f)
         )
 
     /*
@@ -933,6 +1005,11 @@ open class KotlinUsesExtractor(
         // If true, this method implements a Java Collections interface (Collection, Map or List) and may need
         // parameter erasure to match the way this class will appear to an external consumer of the .class file.
         overridesCollectionsMethod: Boolean,
+        // The Java signature of this callable, if known.
+        javaSignature: JavaMethod?,
+        // If true, Java wildcards implied by Kotlin type parameter variance should be added by default to this function's value parameters' types.
+        // (Return-type wildcard addition is always off by default)
+        addParameterWildcardsByDefault: Boolean,
         // The prefix used in the label. "callable", unless a property label is created, then it's "property".
         prefix: String = "callable"
     ): String {
@@ -956,8 +1033,10 @@ open class KotlinUsesExtractor(
             // Collection.remove(Object) because Collection.remove(Collection::E) in the Kotlin universe.
             // If this has happened, erase the type again to get the correct Java signature.
             val maybeAmendedForCollections = if (overridesCollectionsMethod) eraseCollectionsMethodParameterType(it.value.type, name, it.index) else it.value.type
+            // Add any wildcard types that the Kotlin compiler would add in the Java lowering of this function:
+            val withAddedWildcards = addJavaLoweringWildcards(maybeAmendedForCollections, addParameterWildcardsByDefault, javaSignature?.let { sig -> sig.valueParameters[it.index].type })
             // Now substitute any class type parameters in:
-            val maybeSubbed = maybeAmendedForCollections.substituteTypeAndArguments(substitutionMap, TypeContext.OTHER, pluginContext)
+            val maybeSubbed = withAddedWildcards.substituteTypeAndArguments(substitutionMap, TypeContext.OTHER, pluginContext)
             // Finally, mimic the Java extractor's behaviour by naming functions with type parameters for their erased types;
             // those without type parameters are named for the generic type.
             val maybeErased = if (functionTypeParameters.isEmpty()) maybeSubbed else erase(maybeSubbed)
@@ -969,6 +1048,8 @@ open class KotlinUsesExtractor(
                 pluginContext.irBuiltIns.unitType
             else
                 erase(returnType.substituteTypeAndArguments(substitutionMap, TypeContext.RETURN, pluginContext))
+        // Note that `addJavaLoweringWildcards` is not required here because the return type used to form the function
+        // label is always erased.
         val returnTypeId = useType(labelReturnType, TypeContext.RETURN).javaResult.id
         // This suffix is added to generic methods (and constructors) to match the Java extractor's behaviour.
         // Comments in that extractor indicates it didn't want the label of the callable to clash with the raw
@@ -1425,7 +1506,7 @@ open class KotlinUsesExtractor(
             val returnType = getter?.returnType ?: setter?.valueParameters?.singleOrNull()?.type ?: pluginContext.irBuiltIns.unitType
             val typeParams = getFunctionTypeParameters(func)
 
-            getFunctionLabel(p.parent, parentId, p.name.asString(), listOf(), returnType, ext, typeParams, classTypeArgsIncludingOuterClasses, overridesCollectionsMethod = false, prefix = "property")
+            getFunctionLabel(p.parent, parentId, p.name.asString(), listOf(), returnType, ext, typeParams, classTypeArgsIncludingOuterClasses, overridesCollectionsMethod = false, javaSignature = null, addParameterWildcardsByDefault = false, prefix = "property")
         }
     }