Kotlin: Handle Kotlin 2 parents better

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

@@ -88,12 +88,10 @@ class ExternalDeclExtractor(val logger: FileLogger, val invocationTrapFile: Stri
             nextBatch.forEach { workPair ->
                 val (irDecl, possiblyLongSignature) = workPair
                 extractElement(irDecl, possiblyLongSignature, false) { trapFileBW, signature, manager ->
-                    val containingClass = getContainingClassOrSelf(irDecl)
-                    if (containingClass == null) {
-                        logger.errorElement("Unable to get containing class", irDecl)
+                    val binaryPath = getIrDeclarationBinaryPath(irDecl)
+                    if (binaryPath == null) {
+                        logger.errorElement("Unable to get binary path", irDecl)
                     } else {
-                        val binaryPath = getIrClassBinaryPath(containingClass)
-
                         // We want our comments to be the first thing in the file,
                         // so start off with a PlainTrapWriter
                         val tw = PlainTrapWriter(logger.loggerBase, TrapLabelManager(), trapFileBW, diagnosticTrapWriter)

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

@@ -191,7 +191,7 @@ open class KotlinFileExtractor(
                     }
                 }
                 is IrFunction -> {
-                    val parentId = useDeclarationParent(declaration.parent, false)?.cast<DbReftype>()
+                    val parentId = useDeclarationParentOf(declaration, false)?.cast<DbReftype>()
                     if (parentId != null) {
                         extractFunction(declaration, parentId, extractBody = extractFunctionBodies, extractMethodAndParameterTypeAccesses = extractFunctionBodies, extractAnnotations = extractAnnotations, null, listOf())
                     }
@@ -201,21 +201,21 @@ open class KotlinFileExtractor(
                     // Leaving this intentionally empty. init blocks are extracted during class extraction.
                 }
                 is IrProperty -> {
-                    val parentId = useDeclarationParent(declaration.parent, false)?.cast<DbReftype>()
+                    val parentId = useDeclarationParentOf(declaration, false)?.cast<DbReftype>()
                     if (parentId != null) {
                         extractProperty(declaration, parentId, extractBackingField = true, extractFunctionBodies = extractFunctionBodies, extractPrivateMembers = extractPrivateMembers, extractAnnotations = extractAnnotations, null, listOf())
                     }
                     Unit
                 }
                 is IrEnumEntry -> {
-                    val parentId = useDeclarationParent(declaration.parent, false)?.cast<DbReftype>()
+                    val parentId = useDeclarationParentOf(declaration, false)?.cast<DbReftype>()
                     if (parentId != null) {
                         extractEnumEntry(declaration, parentId, extractPrivateMembers, extractFunctionBodies)
                     }
                     Unit
                 }
                 is IrField -> {
-                    val parentId = useDeclarationParent(getFieldParent(declaration), false)?.cast<DbReftype>()
+                    val parentId = useDeclarationParentOf(declaration, false)?.cast<DbReftype>()
                     if (parentId != null) {
                         // For consistency with the Java extractor, enum entries get type accesses only if we're extracting from .kt source (i.e., when `extractFunctionBodies` is set)
                         extractField(declaration, parentId, extractAnnotationEnumTypeAccesses = extractFunctionBodies)
@@ -408,11 +408,10 @@ open class KotlinFileExtractor(
 
     private fun getLocation(decl: IrDeclaration, typeArgs: List<IrTypeArgument>?): Label<DbLocation> {
         return if (typeArgs != null && typeArgs.isNotEmpty()) {
-            val c = getContainingClassOrSelf(decl)
-            if (c == null) {
+            val binaryPath = getIrDeclarationBinaryPath(decl)
+            if (binaryPath == null) {
                 tw.getLocation(decl)
             } else {
-                val binaryPath = getIrClassBinaryPath(c)
                 val newTrapWriter = tw.makeFileTrapWriter(binaryPath, true)
                 newTrapWriter.getWholeFileLocation()
             }
@@ -1285,7 +1284,7 @@ open class KotlinFileExtractor(
             val sourceParentId =
                 maybeSourceParentId ?:
                     if (typeSubstitution != null)
-                        useDeclarationParent(f.parent, false)
+                        useDeclarationParentOf(f, false)
                     else
                         parentId
             if (sourceParentId == null) {
@@ -1617,7 +1616,7 @@ open class KotlinFileExtractor(
                 }
 
                 if (bf != null && extractBackingField) {
-                    val fieldParentId = useDeclarationParent(getFieldParent(bf), false)
+                    val fieldParentId = useDeclarationParentOf(bf, false)
                     if (fieldParentId != null) {
                         val fieldId = extractField(bf, fieldParentId.cast(), extractFunctionBodies)
                         tw.writeKtPropertyBackingFields(id, fieldId)
@@ -2091,7 +2090,7 @@ open class KotlinFileExtractor(
 
     private fun getDefaultsMethodLabel(f: IrFunction): Label<out DbCallable>? {
         val classTypeArgsIncludingOuterClasses = null
-        val parentId = useDeclarationParent(f.parent, false, classTypeArgsIncludingOuterClasses, true)
+        val parentId = useDeclarationParentOf(f, false, classTypeArgsIncludingOuterClasses, true)
         if (parentId == null) {
             logger.errorElement("Couldn't get parent ID for defaults method", f)
             return null
@@ -2154,7 +2153,7 @@ open class KotlinFileExtractor(
         if (overriddenCallTarget.isLocalFunction()) {
             extractTypeAccess(getLocallyVisibleFunctionLabels(overriddenCallTarget).type, locId, id, -1, enclosingCallable, enclosingStmt)
         } else {
-            extractStaticTypeAccessQualifierUnchecked(overriddenCallTarget.parent, id, locId, enclosingCallable, enclosingStmt)
+            extractStaticTypeAccessQualifierUnchecked(overriddenCallTarget, id, locId, enclosingCallable, enclosingStmt)
         }
 
         extractDefaultsCallArguments(id, overriddenCallTarget, enclosingCallable, enclosingStmt, valueArguments, dispatchReceiver, extensionReceiver)
@@ -2380,8 +2379,17 @@ open class KotlinFileExtractor(
         extractValueArguments(argParent, idxOffset)
     }
 
-    private fun extractStaticTypeAccessQualifierUnchecked(parent: IrDeclarationParent, parentExpr: Label<out DbExprparent>, locId: Label<DbLocation>, enclosingCallable: Label<out DbCallable>?, enclosingStmt: Label<out DbStmt>?) {
-        if (parent is IrClass) {
+    private fun extractStaticTypeAccessQualifierUnchecked(target: IrDeclaration, parentExpr: Label<out DbExprparent>, locId: Label<DbLocation>, enclosingCallable: Label<out DbCallable>?, enclosingStmt: Label<out DbStmt>?) {
+        val parent = target.parent
+        if (parent is IrExternalPackageFragment) {
+            // This is in a file class.
+            val fqName = getFileClassFqName(target)
+            if (fqName == null) {
+                logger.error("Can't get FqName for element in external package fragment ${target.javaClass}")
+            } else {
+                extractTypeAccess(useFileClassType(fqName), locId, parentExpr, -1, enclosingCallable, enclosingStmt)
+            }
+        } else if (parent is IrClass) {
             extractTypeAccessRecursive(parent.toRawType(), locId, parentExpr, -1, enclosingCallable, enclosingStmt)
         } else if (parent is IrFile) {
             extractTypeAccess(useFileClassType(parent), locId, parentExpr, -1, enclosingCallable, enclosingStmt)
@@ -2392,7 +2400,7 @@ open class KotlinFileExtractor(
 
     private fun extractStaticTypeAccessQualifier(target: IrDeclaration, parentExpr: Label<out DbExprparent>, locId: Label<DbLocation>, enclosingCallable: Label<out DbCallable>?, enclosingStmt: Label<out DbStmt>?) {
         if (target.shouldExtractAsStatic) {
-            extractStaticTypeAccessQualifierUnchecked(target.parent, parentExpr, locId, enclosingCallable, enclosingStmt)
+            extractStaticTypeAccessQualifierUnchecked(target, parentExpr, locId, enclosingCallable, enclosingStmt)
         }
     }
 
@@ -3841,7 +3849,7 @@ open class KotlinFileExtractor(
                         val type = useType(pluginContext.irBuiltIns.unitType)
                         val locId = tw.getLocation(e)
                         val parentClass = irConstructor.parentAsClass
-                        val parentId = useDeclarationParent(parentClass, false, null, true)
+                        val parentId = useDeclarationParentOf(irConstructor, false, null, true)
                         if (parentId == null) {
                             logger.errorElement("Cannot get parent ID for obinit", e)
                             return

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

@@ -2,6 +2,7 @@ package com.github.codeql
 
 import com.github.codeql.utils.*
 import com.github.codeql.utils.versions.codeQlWithHasQuestionMark
+import com.github.codeql.utils.versions.getFileClassFqName
 import com.github.codeql.utils.versions.getKotlinType
 import com.github.codeql.utils.versions.isRawType
 import com.semmle.extractor.java.OdasaOutput
@@ -71,18 +72,41 @@ open class KotlinUsesExtractor(
         TypeResult(fakeKotlinType(), "", "")
     )
 
+    fun useFileClassType(fqName: FqName) = TypeResults(
+        TypeResult(extractFileClass(fqName), "", ""),
+        TypeResult(fakeKotlinType(), "", "")
+    )
+
+    private fun useFileClassType(pkg: String, jvmName: String) = TypeResults(
+        TypeResult(extractFileClass(pkg, jvmName), "", ""),
+        TypeResult(fakeKotlinType(), "", "")
+    )
+
     fun extractFileClass(f: IrFile): Label<out DbClassorinterface> {
         val pkg = f.fqName.asString()
         val jvmName = getFileClassName(f)
+        val id = extractFileClass(pkg, jvmName)
+        if (tw.lm.fileClassLocationsExtracted.add(f)) {
+            val fileId = tw.mkFileId(f.path, false)
+            val locId = tw.getWholeFileLocation(fileId)
+            tw.writeHasLocation(id, locId)
+        }
+        return id
+    }
+
+    private fun extractFileClass(fqName: FqName): Label<out DbClassorinterface> {
+        val pkg = if (fqName.isRoot()) "" else fqName.parent().asString()
+        val jvmName = fqName.shortName().asString()
+        return extractFileClass(pkg, jvmName)
+    }
+
+    private fun extractFileClass(pkg: String, jvmName: String): Label<out DbClassorinterface> {
         val qualClassName = if (pkg.isEmpty()) jvmName else "$pkg.$jvmName"
         val label = "@\"class;$qualClassName\""
         val id: Label<DbClassorinterface> = tw.getLabelFor(label) {
-            val fileId = tw.mkFileId(f.path, false)
-            val locId = tw.getWholeFileLocation(fileId)
             val pkgId = extractPackage(pkg)
             tw.writeClasses_or_interfaces(it, jvmName, pkgId, it)
             tw.writeFile_class(it)
-            tw.writeHasLocation(it, locId)
 
             addModifiers(it, "public", "final")
         }
@@ -250,7 +274,11 @@ open class KotlinUsesExtractor(
             is IrClass -> extractExternalClassLater(parent)
             is IrFunction -> extractExternalEnclosingClassLater(parent)
             is IrFile -> logger.error("extractExternalEnclosingClassLater but no enclosing class.")
-            else -> logger.error("Unrecognised extractExternalEnclosingClassLater: " + d.javaClass)
+            is IrExternalPackageFragment -> {
+                // The parent is a (multi)file class. We don't need
+                // extract it separately.
+            }
+            else -> logger.error("Unrecognised extractExternalEnclosingClassLater ${parent.javaClass} for ${d.javaClass}")
         }
     }
 
@@ -293,7 +321,17 @@ open class KotlinUsesExtractor(
 
     private fun extractFunctionLaterIfExternalFileMember(f: IrFunction) {
         if (isExternalFileClassMember(f)) {
-            extractExternalClassLater(f.parentAsClass)
+            val p = f.parent
+            when (p) {
+                is IrClass -> extractExternalClassLater(p)
+                is IrExternalPackageFragment -> {
+                    // The parent is a (multi)file class. We don't need
+                    // extract it separately.
+                }
+                else -> {
+                    logger.warn("Unexpected parent type ${p.javaClass} for external file class member")
+                }
+            }
             (f as? IrSimpleFunction)?.correspondingPropertySymbol?.let {
                 extractPropertyLaterIfExternalFileMember(it.owner)
                 // No need to extract the function specifically, as the property's
@@ -761,6 +799,41 @@ open class KotlinUsesExtractor(
         }
     }
 
+    private fun parentOf(d: IrDeclaration): IrDeclarationParent {
+        if (d is IrField) {
+            return getFieldParent(d)
+        }
+        return d.parent
+    }
+
+    fun useDeclarationParentOf(
+        // The declaration
+        d: IrDeclaration,
+        // Whether the type of entity whose parent this is can be a
+        // top-level entity in the JVM's eyes. If so, then its parent may
+        // be a file; otherwise, if dp is a file foo.kt, then the parent
+        // is really the JVM class FooKt.
+        canBeTopLevel: Boolean,
+        classTypeArguments: List<IrTypeArgument>? = null,
+        inReceiverContext: Boolean = false):
+        Label<out DbElement>? {
+
+        val parent = parentOf(d)
+        if (parent is IrExternalPackageFragment) {
+            // This is in a file class.
+            val fqName = getFileClassFqName(d)
+            if (fqName == null) {
+                logger.error("Can't get FqName for element in external package fragment ${d.javaClass}")
+                return null
+            }
+            return extractFileClass(fqName)
+        }
+        return useDeclarationParent(parent, canBeTopLevel, classTypeArguments, inReceiverContext)
+    }
+
+    // Generally, useDeclarationParentOf should be used instead of
+    // calling this directly, as this cannot handle
+    // IrExternalPackageFragment
     fun useDeclarationParent(
         // The declaration parent according to Kotlin
         dp: IrDeclarationParent,
@@ -792,8 +865,7 @@ open class KotlinUsesExtractor(
                 }
             is IrFunction -> useFunction(dp)
             is IrExternalPackageFragment -> {
-                // TODO
-                logger.error("Unhandled IrExternalPackageFragment")
+                logger.error("Unable to handle IrExternalPackageFragment as an IrDeclarationParent")
                 null
             }
             else -> {
@@ -1035,7 +1107,7 @@ open class KotlinUsesExtractor(
      * in.
      */
     fun getFunctionLabel(f: IrFunction, classTypeArgsIncludingOuterClasses: List<IrTypeArgument>?): String? {
-        val parentId = useDeclarationParent(f.parent, false, classTypeArgsIncludingOuterClasses, true)
+        val parentId = useDeclarationParentOf(f, false, classTypeArgsIncludingOuterClasses, true)
         if (parentId == null) {
             logger.error("Couldn't get parent ID for function label")
             return null
@@ -1332,7 +1404,7 @@ open class KotlinUsesExtractor(
             return ids.function.cast<T>()
         }
         val javaFun = kotlinFunctionToJavaEquivalent(f, noReplace)
-        val parentId = useDeclarationParent(javaFun.parent, false, classTypeArgsIncludingOuterClasses, true)
+        val parentId = useDeclarationParentOf(javaFun, false, classTypeArgsIncludingOuterClasses, true)
         if (parentId == null) {
             logger.error("Couldn't find parent ID for function ${f.name.asString()}")
             return null
@@ -1639,7 +1711,7 @@ open class KotlinUsesExtractor(
         val overriddenParentAttributes = (declarationParent as? IrFunction)?.let {
             (this as? KotlinFileExtractor)?.declarationStack?.findOverriddenAttributes(it)
         }
-        val parentId = parent ?: overriddenParentAttributes?.id ?: useDeclarationParent(declarationParent, false)
+        val parentId = parent ?: overriddenParentAttributes?.id ?: useDeclarationParentOf(vp, false)
 
         val idxBase = overriddenParentAttributes?.valueParameters?.indexOf(vp) ?: vp.index
         val idxOffset = if (declarationParent is IrFunction && declarationParent.extensionReceiverParameter != null)
@@ -1667,7 +1739,7 @@ open class KotlinUsesExtractor(
             it.isConst || it.isLateinit
         } ?: false
 
-    fun getFieldParent(f: IrField) =
+    private fun getFieldParent(f: IrField) =
         f.parentClassOrNull?.let {
             if (it.isCompanion && isDirectlyExposableCompanionObjectField(f))
                 it.parent
@@ -1684,7 +1756,7 @@ open class KotlinUsesExtractor(
         }
 
     fun getFieldLabel(f: IrField): String {
-        val parentId = useDeclarationParent(getFieldParent(f), false)
+        val parentId = useDeclarationParentOf(f, false)
         // Distinguish backing fields of properties based on their extension receiver type;
         // otherwise two extension properties declared in the same enclosing context will get
         // clashing trap labels. These are always private, so we can just make up a label without
@@ -1697,7 +1769,7 @@ open class KotlinUsesExtractor(
         tw.getLabelFor<DbField>(getFieldLabel(f)).also { extractFieldLaterIfExternalFileMember(f) }
 
     fun getPropertyLabel(p: IrProperty): String? {
-        val parentId = useDeclarationParent(p.parent, false)
+        val parentId = useDeclarationParentOf(p, false)
         if (parentId == null) {
             return null
         } else {
@@ -1729,15 +1801,15 @@ open class KotlinUsesExtractor(
         }
 
     fun getEnumEntryLabel(ee: IrEnumEntry): String {
-        val parentId = useDeclarationParent(ee.parent, false)
+        val parentId = useDeclarationParentOf(ee, false)
         return "@\"field;{$parentId};${ee.name.asString()}\""
     }
 
     fun useEnumEntry(ee: IrEnumEntry): Label<out DbField> =
         tw.getLabelFor(getEnumEntryLabel(ee))
 
     fun getTypeAliasLabel(ta: IrTypeAlias): String {
-        val parentId = useDeclarationParent(ta.parent, true)
+        val parentId = useDeclarationParentOf(ta, true)
         return "@\"type_alias;{$parentId};${ta.name.asString()}\""
     }
 

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

@@ -48,6 +48,15 @@ class TrapLabelManager {
      * duplication.
      */
     val genericSpecialisationsExtracted = HashSet<String>()
+
+    /**
+     * Sometimes, when we extract a file class we don't have the IrFile
+     * for it, so we are not able to give it a location. This means that
+     * the location is written outside of the label creation.
+     * This allows us to keep track of whether we've written the location
+     * already in this TRAP file, to avoid duplication.
+     */
+    val fileClassLocationsExtracted = HashSet<IrFile>()
 }
 
 /**