Extract generic method prototypes

These feature substituted types according to their declaring generic specialisation, with wildcards that reach top-level being converted to their upper or lower bound depending on usage context.

This commit also includes an incidental fix such that constructors declare their return-type as unit, consistent with the Java extractor.

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

@@ -1,5 +1,8 @@
 package com.github.codeql
 
+import com.github.codeql.utils.lowerBound
+import com.github.codeql.utils.substituteTypeAndArguments
+import com.github.codeql.utils.upperBound
 import com.semmle.extractor.java.OdasaOutput
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.descriptors.ClassKind
@@ -9,8 +12,10 @@ import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
+import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.types.Variance
 
 open class KotlinFileExtractor(
     override val logger: FileLogger,
@@ -43,7 +48,7 @@ open class KotlinFileExtractor(
             is IrClass -> return getClassLabel(element, listOf()).classLabel
             is IrTypeParameter -> return getTypeParameterLabel(element)
             is IrFunction -> return getFunctionLabel(element)
-            is IrValueParameter -> return getValueParameterLabel(element)
+            is IrValueParameter -> return getValueParameterLabel(element, null)
             is IrProperty -> return getPropertyLabel(element)
             is IrField -> return getFieldLabel(element)
             is IrEnumEntry -> return getEnumEntryLabel(element)
@@ -81,7 +86,7 @@ open class KotlinFileExtractor(
         return id
     }
 
-    fun extractClassInstance(c: IrClass, typeArgs: List<IrTypeArgument>): Label<out DbClassorinterface> {
+    fun extractClassInstance(c: IrClass, typeArgs: List<IrTypeArgument>, extractFunctionPrototypes: Boolean): Label<out DbClassorinterface> {
         if (typeArgs.isEmpty()) {
             logger.warn(Severity.ErrorSevere, "Instance without type arguments: " + c.name.asString())
         }
@@ -122,6 +127,18 @@ open class KotlinFileExtractor(
         val locId = tw.getLocation(c)
         tw.writeHasLocation(id, locId)
 
+        if (extractFunctionPrototypes) {
+            val typeParamSubstitution = c.typeParameters.map({ it.symbol }).zip(typeArgs).toMap()
+
+            c.declarations.map {
+                when(it) {
+                    is IrFunction -> extractFunction(it, id, false, typeParamSubstitution)
+                    is IrProperty -> extractProperty(it, id, false, typeParamSubstitution)
+                    else -> {}
+                }
+            }
+        }
+
         return id
     }
 
@@ -254,9 +271,12 @@ open class KotlinFileExtractor(
         return FieldResult(instanceId, instanceName)
     }
 
-    fun extractValueParameter(vp: IrValueParameter, parent: Label<out DbCallable>, idx: Int): TypeResults {
-        val id = useValueParameter(vp)
-        val type = useType(vp.type)
+    fun extractValueParameter(vp: IrValueParameter, parent: Label<out DbCallable>, idx: Int, typeSubstitutionMap: Map<IrTypeParameterSymbol, IrTypeArgument>?): TypeResults {
+        val id = useValueParameter(vp, parent)
+        val substitutedType = vp.type.substituteTypeAndArguments(typeSubstitutionMap) {
+            it.lowerBound(pluginContext)
+        }
+        val type = useType(substitutedType)
         val locId = tw.getLocation(vp)
         tw.writeParams(id, type.javaResult.id, type.kotlinResult.id, idx, parent, id)
         tw.writeHasLocation(id, locId)
@@ -331,7 +351,7 @@ open class KotlinFileExtractor(
         }
     }
 
-    fun extractFunction(f: IrFunction, parentId: Label<out DbReftype>): Label<out DbCallable> {
+    fun extractFunction(f: IrFunction, parentId: Label<out DbReftype>, extractBody: Boolean = true, typeSubstitutionMap: Map<IrTypeParameterSymbol, IrTypeArgument>? = null): Label<out DbCallable> {
         currentFunction = f
 
         f.typeParameters.map { extractTypeParameter(it) }
@@ -342,19 +362,21 @@ open class KotlinFileExtractor(
             if (f.isLocalFunction())
                 getLocalFunctionLabels(f).function
             else
-                useFunction<DbCallable>(f)
+                // TODO: figure out whether to standardise on naming top-level functions for the file-class
+                //   or (as temporarily done here) for their containing package.
+                useFunction<DbCallable>(f, if (f.parent is IrFile) useDeclarationParent(f.parent) else parentId)
 
         val extReceiver = f.extensionReceiverParameter
         val idxOffset = if (extReceiver != null) 1 else 0
         val paramTypes = f.valueParameters.mapIndexed { i, vp ->
-            extractValueParameter(vp, id, i + idxOffset)
+            extractValueParameter(vp, id, i + idxOffset, typeSubstitutionMap)
         }
         val allParamTypes = if (extReceiver != null) {
             val extendedType = useType(extReceiver.type)
             @Suppress("UNCHECKED_CAST")
             tw.writeKtExtensionFunctions(id as Label<DbMethod>, extendedType.javaResult.id, extendedType.kotlinResult.id)
 
-            val t = extractValueParameter(extReceiver, id, 0)
+            val t = extractValueParameter(extReceiver, id, 0, null)
             val l = mutableListOf(t)
             l.addAll(paramTypes)
             l
@@ -364,13 +386,21 @@ open class KotlinFileExtractor(
 
         val paramsSignature = allParamTypes.joinToString(separator = ",", prefix = "(", postfix = ")") { it.javaResult.signature!! }
 
+        val substReturnType = f.returnType.substituteTypeAndArguments(typeSubstitutionMap) {
+            it.upperBound(pluginContext)
+        }
+
         if (f.symbol is IrConstructorSymbol) {
-            val returnType = useType(erase(f.returnType), TypeContext.RETURN)
-            val shortName = if (f.returnType.isAnonymous) "" else f.returnType.classFqName?.shortName()?.asString() ?: f.name.asString()
+            val unitType = useType(pluginContext.irBuiltIns.unitType, TypeContext.RETURN)
+            val shortName = when {
+                f.returnType.isAnonymous -> ""
+                typeSubstitutionMap != null -> useType(substReturnType).javaResult.shortName
+                else -> f.returnType.classFqName?.shortName()?.asString() ?: f.name.asString()
+            }
             @Suppress("UNCHECKED_CAST")
-            tw.writeConstrs(id as Label<DbConstructor>, shortName, "$shortName$paramsSignature", returnType.javaResult.id, returnType.kotlinResult.id, parentId, id)
+            tw.writeConstrs(id as Label<DbConstructor>, shortName, "$shortName$paramsSignature", unitType.javaResult.id, unitType.kotlinResult.id, parentId, id)
         } else {
-            val returnType = useType(f.returnType, TypeContext.RETURN)
+            val returnType = useType(substReturnType, TypeContext.RETURN)
             val shortName = f.name.asString()
             @Suppress("UNCHECKED_CAST")
             tw.writeMethods(id as Label<DbMethod>, shortName, "$shortName$paramsSignature", returnType.javaResult.id, returnType.kotlinResult.id, parentId, id)
@@ -378,7 +408,9 @@ open class KotlinFileExtractor(
 
         tw.writeHasLocation(id, locId)
         val body = f.body
-        if(body != null) {
+        if(body != null && extractBody) {
+            // Type substitution should only be used to extract a prototype, not the body as well:
+            assert(typeSubstitutionMap == null)
             extractBody(body, id)
         }
 
@@ -395,8 +427,8 @@ open class KotlinFileExtractor(
         return id
     }
 
-    fun extractProperty(p: IrProperty, parentId: Label<out DbReftype>) {
-        val id = useProperty(p)
+    fun extractProperty(p: IrProperty, parentId: Label<out DbReftype>, extractBackingField: Boolean = true, typeSubstitutionMap: Map<IrTypeParameterSymbol, IrTypeArgument>? = null) {
+        val id = useProperty(p, parentId)
         val locId = tw.getLocation(p)
         tw.writeKtProperties(id, p.name.asString())
         tw.writeHasLocation(id, locId)
@@ -407,7 +439,7 @@ open class KotlinFileExtractor(
 
         if(getter != null) {
             @Suppress("UNCHECKED_CAST")
-            val getterId = extractFunction(getter, parentId) as Label<out DbMethod>
+            val getterId = extractFunction(getter, parentId, extractBackingField, typeSubstitutionMap) as Label<out DbMethod>
             tw.writeKtPropertyGetters(id, getterId)
         } else {
             if (p.modality != Modality.FINAL || !isExternalDeclaration(p)) {
@@ -420,7 +452,7 @@ open class KotlinFileExtractor(
                 logger.warnElement(Severity.ErrorSevere, "!isVar property with a setter", p)
             }
             @Suppress("UNCHECKED_CAST")
-            val setterId = extractFunction(setter, parentId) as Label<out DbMethod>
+            val setterId = extractFunction(setter, parentId, extractBackingField, typeSubstitutionMap) as Label<out DbMethod>
             tw.writeKtPropertySetters(id, setterId)
         } else {
             if (p.isVar && !isExternalDeclaration(p)) {
@@ -428,7 +460,7 @@ open class KotlinFileExtractor(
             }
         }
 
-        if(bf != null) {
+        if(bf != null && extractBackingField) {
             val fieldId = extractField(bf, parentId)
             tw.writeKtPropertyBackingFields(id, fieldId)
         }
@@ -684,17 +716,16 @@ open class KotlinFileExtractor(
                    isFunction("kotlin", "Double", fName)
         }
 
-        fun extractMethodAccess(callTarget: IrFunction, extractTypeArguments: Boolean = true){
+        fun extractMethodAccess(callTarget: IrFunction, extractMethodTypeArguments: Boolean = true, extractClassTypeArguments: Boolean = false) {
             val id = tw.getFreshIdLabel<DbMethodaccess>()
             val type = useType(c.type)
             val locId = tw.getLocation(c)
-
             tw.writeExprs_methodaccess(id, type.javaResult.id, type.kotlinResult.id, parent, idx)
             tw.writeHasLocation(id, locId)
             tw.writeCallableEnclosingExpr(id, callable)
             tw.writeStatementEnclosingExpr(id, enclosingStmt)
 
-            if (extractTypeArguments) {
+            if (extractMethodTypeArguments) {
                 // type arguments at index -2, -3, ...
                 extractTypeArguments(c, id, callable, enclosingStmt, -2, true)
             }
@@ -713,10 +744,28 @@ open class KotlinFileExtractor(
                 tw.writeCallableBinding(idNewexpr, ids.constructor)
 
             } else {
-                val methodId = useFunction<DbMethod>(callTarget)
+                val dr = c.dispatchReceiver
+
+                // Returns true if type is C<T1, T2, ...> where C is declared `class <T1, T2, ...> C { ... }`
+                fun isUnspecialised(type: IrSimpleType) =
+                    type.classifier.owner is IrClass &&
+                            (type.classifier.owner as IrClass).typeParameters.zip(type.arguments).all { paramAndArg ->
+                                (paramAndArg.second as? IrTypeProjection)?.let {
+                                    // Type arg refers to the class' own type parameter?
+                                    it.variance == Variance.INVARIANT &&
+                                            it.type.classifierOrNull?.owner === paramAndArg.first
+                                } ?: false
+                            }
+
+                val drType = dr?.type
+                val methodId =
+                    if (drType != null && extractClassTypeArguments && drType is IrSimpleType && !isUnspecialised(drType))
+                        useFunction<DbCallable>(callTarget, drType.arguments)
+                    else
+                        useFunction<DbCallable>(callTarget)
+
                 tw.writeCallableBinding(id, methodId)
 
-                val dr = c.dispatchReceiver
                 if (dr != null) {
                     extractExpressionExpr(dr, callable, id, -1, enclosingStmt)
                 }
@@ -1038,7 +1087,7 @@ open class KotlinFileExtractor(
                 }
             }
             else -> {
-                extractMethodAccess(c.symbol.owner)
+                extractMethodAccess(c.symbol.owner, true, true)
             }
         }
     }

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

@@ -87,7 +87,14 @@ open class KotlinUsesExtractor(
         return KotlinSourceFileExtractor(newLogger, newTrapWriter, clsFile, externalClassExtractor, primitiveTypeMapping, pluginContext)
     }
 
-    fun useClassInstance(c: IrClass, typeArgs: List<IrTypeArgument>): UseClassInstanceResult {
+    private fun anyDeclarationExtracted(c: IrClass, id: Label<out DbClassorinterface>) =
+        c.declarations.any {
+            it is IrFunction &&
+            tw.getExistingLabelFor<DbCallable>(getFunctionLabel(
+                id, it.name.asString(), it.valueParameters, it.returnType, it.extensionReceiverParameter)) != null
+        }
+
+    fun useClassInstance(c: IrClass, typeArgs: List<IrTypeArgument>, inReceiverContext: Boolean = false): UseClassInstanceResult {
         if (c.isAnonymousObject) {
             logger.warn(Severity.ErrorSevere, "Unexpected access to anonymous class instance")
         }
@@ -100,7 +107,8 @@ open class KotlinUsesExtractor(
 
         val extractClass = substituteClass ?: c
 
-        val classTypeResult = addClassLabel(extractClass, typeArgs)
+        val classTypeResult = addClassLabel(extractClass, typeArgs, inReceiverContext)
+
         // Extract both the Kotlin and equivalent Java classes, so that we have database entries
         // for both even if all internal references to the Kotlin type are substituted.
         if(c != extractClass) {
@@ -137,17 +145,27 @@ open class KotlinUsesExtractor(
         externalClassExtractor.extractLater(c)
     }
 
-    fun addClassLabel(c: IrClass, typeArgs: List<IrTypeArgument>): TypeResult<DbClassorinterface> {
+    fun addClassLabel(c: IrClass, typeArgs: List<IrTypeArgument>, inReceiverContext: Boolean = false): TypeResult<DbClassorinterface> {
         val classLabelResult = getClassLabel(c, typeArgs)
+
+        var shouldExtractClass = false
+
         val classLabel : Label<out DbClassorinterface> = tw.getLabelFor(classLabelResult.classLabel, {
             // If this is a generic type instantiation then it has no
             // source entity, so we need to extract it here
-            if (typeArgs.isNotEmpty()) {
-                this.withSourceFileOfClass(c).extractClassInstance(c, typeArgs)
-            }
+            shouldExtractClass = true
 
             extractClassLaterIfExternal(c)
         })
+
+        if (typeArgs.isNotEmpty()) {
+            // Extract again if we've already extracted the class itself but not its declared functions:
+            // This might happen e.g. if we see it for the first time in the context of a parameter type (which doesn't
+            // require method prototype extraction), then later as a function receiver (which does).
+            if (shouldExtractClass || (inReceiverContext && !anyDeclarationExtracted(c, classLabel)))
+                this.withSourceFileOfClass(c).extractClassInstance(c, typeArgs, inReceiverContext)
+        }
+
         return TypeResult(
             classLabel,
             c.fqNameWhenAvailable?.asString(),
@@ -441,10 +459,10 @@ class X {
         }
     }
 
-    fun useDeclarationParent(dp: IrDeclarationParent): Label<out DbElement> =
+    fun useDeclarationParent(dp: IrDeclarationParent, classTypeArguments: List<IrTypeArgument>? = null, inReceiverContext: Boolean = false): Label<out DbElement> =
         when(dp) {
             is IrFile -> usePackage(dp.fqName.asString())
-            is IrClass -> useClassSource(dp)
+            is IrClass -> if (classTypeArguments != null) useClassInstance(dp, classTypeArguments, inReceiverContext).typeResult.id else useClassSource(dp)
             is IrFunction -> useFunction(dp)
             else -> {
                 logger.warn(Severity.ErrorSevere, "Unrecognised IrDeclarationParent: " + dp.javaClass)
@@ -452,8 +470,8 @@ class X {
             }
         }
 
-    fun getFunctionLabel(f: IrFunction) : String {
-        return getFunctionLabel(f.parent, f.name.asString(), f.valueParameters, f.returnType, f.extensionReceiverParameter)
+    fun getFunctionLabel(f: IrFunction, classTypeArguments: List<IrTypeArgument>? = null) : String {
+        return getFunctionLabel(f.parent, f.name.asString(), f.valueParameters, f.returnType, f.extensionReceiverParameter, classTypeArguments)
     }
 
     fun getFunctionLabel(
@@ -461,6 +479,21 @@ class X {
         name: String,
         parameters: List<IrValueParameter>,
         returnType: IrType,
+        extensionReceiverParameter: IrValueParameter?,
+        classTypeArguments: List<IrTypeArgument>? = null
+    ): String {
+        val parentId = useDeclarationParent(parent, classTypeArguments, true)
+        return getFunctionLabel(parentId, name, parameters, returnType, extensionReceiverParameter)
+    }
+
+    fun getFunctionLabel(f: IrFunction, parentId: Label<out DbElement>) =
+        getFunctionLabel(parentId, f.name.asString(), f.valueParameters, f.returnType, f.extensionReceiverParameter)
+
+    fun getFunctionLabel(
+        parentId: Label<out DbElement>,
+        name: String,
+        parameters: List<IrValueParameter>,
+        returnType: IrType,
         extensionReceiverParameter: IrValueParameter?
     ): String {
         val allParams = if (extensionReceiverParameter == null) {
@@ -472,7 +505,6 @@ class X {
                         }
         val paramTypeIds = allParams.joinToString { "{${useType(erase(it.type)).javaResult.id}}" }
         val returnTypeId = useType(erase(returnType)).javaResult.id
-        val parentId = useDeclarationParent(parent)
         return "@\"callable;{$parentId}.$name($paramTypeIds){$returnTypeId}\""
     }
 
@@ -505,21 +537,27 @@ class X {
         return res
     }
 
-    fun <T: DbCallable> useFunction(f: IrFunction): Label<out T> {
-        if (f.isLocalFunction()) {
-            val ids = getLocalFunctionLabels(f)
-            @Suppress("UNCHECKED_CAST")
-            return ids.function as Label<out T>
-        }
-
-        val label = getFunctionLabel(f)
+    fun <T: DbCallable> useFunctionCommon(f: IrFunction, label: String): Label<out T> {
         val id: Label<T> = tw.getLabelFor(label)
-        if(isExternalDeclaration(f)) {
+        if (isExternalDeclaration(f)) {
             extractExternalEnclosingClassLater(f)
         }
         return id
     }
 
+    fun <T: DbCallable> useFunction(f: IrFunction, classTypeArguments: List<IrTypeArgument>? = null): Label<out T> {
+        if (f.isLocalFunction()) {
+            val ids = getLocalFunctionLabels(f)
+            @Suppress("UNCHECKED_CAST")
+            return ids.function as Label<out T>
+        } else {
+            return useFunctionCommon<T>(f, getFunctionLabel(f, classTypeArguments))
+        }
+    }
+
+    fun <T: DbCallable> useFunction(f: IrFunction, parentId: Label<out DbElement>) =
+        useFunctionCommon<T>(f, getFunctionLabel(f, parentId))
+
     fun getTypeArgumentLabel(
         arg: IrTypeArgument
     ): TypeResult<DbReftype> {
@@ -696,7 +734,7 @@ class X {
 
     fun useValueDeclaration(d: IrValueDeclaration): Label<out DbVariable> =
         when(d) {
-            is IrValueParameter -> useValueParameter(d)
+            is IrValueParameter -> useValueParameter(d, null)
             is IrVariable -> useVariable(d)
             else -> {
                 logger.warn(Severity.ErrorSevere, "Unrecognised IrValueDeclaration: " + d.javaClass)
@@ -729,9 +767,12 @@ class X {
     fun eraseTypeParameter(t: IrTypeParameter) =
         erase(t.superTypes[0])
 
-    fun getValueParameterLabel(vp: IrValueParameter): String {
-        @Suppress("UNCHECKED_CAST")
-        val parentId: Label<out DbMethod> = useDeclarationParent(vp.parent) as Label<out DbMethod>
+    /**
+     * Gets the label for `vp` in the context of function instance `parent`, or in that of its declaring function if
+     * `parent` is null.
+     */
+    fun getValueParameterLabel(vp: IrValueParameter, parent: Label<out DbCallable>?): String {
+        val parentId = parent ?: useDeclarationParent(vp.parent)
         val idx = vp.index
         if (idx < 0) {
             // We're not extracting this and this@TYPE parameters of functions:
@@ -740,8 +781,9 @@ class X {
         return "@\"params;{$parentId};$idx\""
     }
 
-    fun useValueParameter(vp: IrValueParameter): Label<out DbParam> =
-        tw.getLabelFor(getValueParameterLabel(vp))
+
+    fun useValueParameter(vp: IrValueParameter, parent: Label<out DbCallable>?): Label<out DbParam> =
+        tw.getLabelFor(getValueParameterLabel(vp, parent))
 
     fun getFieldLabel(f: IrField): String {
         val parentId = useDeclarationParent(f.parent)
@@ -751,14 +793,18 @@ class X {
     fun useField(f: IrField): Label<out DbField> =
         tw.getLabelFor(getFieldLabel(f))
 
-    fun getPropertyLabel(p: IrProperty): String {
-        val parentId = useDeclarationParent(p.parent)
-        return "@\"property;{$parentId};${p.name.asString()}\""
-    }
+    fun getPropertyLabel(p: IrProperty) =
+        getPropertyLabel(p, useDeclarationParent(p.parent))
+
+    fun getPropertyLabel(p: IrProperty, parentId: Label<out DbElement>) =
+        "@\"property;{$parentId};${p.name.asString()}\""
 
     fun useProperty(p: IrProperty): Label<out DbKt_property> =
         tw.getLabelFor(getPropertyLabel(p))
 
+    fun useProperty(p: IrProperty, parentId: Label<out DbElement>): Label<out DbKt_property> =
+        tw.getLabelFor(getPropertyLabel(p, parentId))
+
     fun getEnumEntryLabel(ee: IrEnumEntry): String {
         val parentId = useDeclarationParent(ee.parent)
         return "@\"field;{$parentId};${ee.name.asString()}\""

java/kotlin-extractor/src/main/kotlin/utils/TypeSubstitution.kt

@@ -1,5 +1,6 @@
 package com.github.codeql.utils
 
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.ir.declarations.IrTypeParameter
 import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
 import org.jetbrains.kotlin.ir.types.*
@@ -76,3 +77,34 @@ fun IrSimpleType.substituteTypeArguments(substitutionMap: Map<IrTypeParameterSym
         annotations
     )
 }
+
+fun IrTypeArgument.upperBound(context: IrPluginContext) =
+    when(this) {
+        is IrStarProjection -> context.irBuiltIns.anyNType
+        is IrTypeProjection -> when(this.variance) {
+            Variance.INVARIANT -> this.type
+            Variance.IN_VARIANCE -> if (this.type.isNullable()) context.irBuiltIns.anyNType else context.irBuiltIns.anyType
+            Variance.OUT_VARIANCE -> this.type
+        }
+        else -> context.irBuiltIns.anyNType
+    }
+
+fun IrTypeArgument.lowerBound(context: IrPluginContext) =
+    when(this) {
+        is IrStarProjection -> context.irBuiltIns.nothingType
+        is IrTypeProjection -> when(this.variance) {
+            Variance.INVARIANT -> this.type
+            Variance.IN_VARIANCE -> this.type
+            Variance.OUT_VARIANCE -> if (this.type.isNullable()) context.irBuiltIns.nothingNType else context.irBuiltIns.nothingType
+        }
+        else -> context.irBuiltIns.nothingType
+    }
+
+fun IrType.substituteTypeAndArguments(substitutionMap: Map<IrTypeParameterSymbol, IrTypeArgument>?, topLevelMatchHandler: (IrTypeArgument) -> IrType): IrType =
+    substitutionMap?.let { substMap ->
+        this.classifierOrNull?.let { typeClassifier ->
+            substMap[typeClassifier]?.let {
+                topLevelMatchHandler(it)
+            } ?: (this as IrSimpleType).substituteTypeArguments(substMap)
+        } ?: this
+    } ?: this