Kotlin: Handle IrExternalPackageFragment when dealing with external decls

Author: igfoo

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

@@ -409,11 +409,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()
             }

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

@@ -274,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}")
         }
     }
 
@@ -317,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

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
 import org.jetbrains.kotlin.ir.util.parentClassOrNull
+import org.jetbrains.kotlin.load.kotlin.FacadeClassSource
 import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 
 // Adapted from Kotlin's interpreter/Utils.kt function 'internalName'
@@ -108,14 +109,40 @@ private fun getRawIrClassBinaryPath(irClass: IrClass) =
 fun getIrClassBinaryPath(irClass: IrClass): String {
   return getRawIrClassBinaryPath(irClass)
   // Otherwise, make up a fake location:
-    ?: "/!unknown-binary-location/${getIrElementBinaryName(irClass).replace(".", "/")}.class"
+    ?: getUnknownBinaryLocation(getIrElementBinaryName(irClass))
 }
 
-fun getContainingClassOrSelf(decl: IrDeclaration): IrClass? {
-    return when(decl) {
-        is IrClass -> decl
-        else -> decl.parentClassOrNull
+fun getIrDeclarationBinaryPath(d: IrDeclaration): String? {
+    if (d is IrClass) {
+        return getIrClassBinaryPath(d)
     }
+    val parentClass = d.parentClassOrNull
+    if (parentClass != null) {
+        return getIrClassBinaryPath(parentClass)
+    }
+    if (d.parent is IrExternalPackageFragment) {
+        // This is in a file class.
+        // Get the name in a similar way to the compiler's ExternalPackageParentPatcherLowering
+        // visitMemberAccess/generateOrGetFacadeClass.
+        if (d is IrMemberWithContainerSource) {
+            val containerSource = d.containerSource
+            if (containerSource is FacadeClassSource) {
+                val facadeClassName = containerSource.facadeClassName
+                if (facadeClassName != null) {
+                    // This is a multifile-class
+                    return getUnknownBinaryLocation(facadeClassName.fqNameForTopLevelClassMaybeWithDollars.asString())
+                } else {
+                    // This is a file-class
+                    return getUnknownBinaryLocation(containerSource.className.fqNameForTopLevelClassMaybeWithDollars.asString())
+                }
+            }
+        }
+    }
+    return null
+}
+
+fun getUnknownBinaryLocation(s: String): String {
+    return "/!unknown-binary-location/${s.replace(".", "/")}.class"
 }
 
 fun getJavaEquivalentClassId(c: IrClass) =

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

@@ -3,6 +3,7 @@ package com.github.codeql.utils
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrExternalPackageFragment
 import org.jetbrains.kotlin.ir.util.isFileClass
 import org.jetbrains.kotlin.ir.util.parentClassOrNull
 
@@ -28,6 +29,10 @@ fun isExternalDeclaration(d: IrDeclaration): Boolean {
                 CONST Int type=kotlin.Int value=-2147483648
     */
     val p = d.parent
+    if (p is IrExternalPackageFragment) {
+        // This is an external declaration in a (multi)file class
+        return true
+    }
     if (p is IrDeclaration) {
         return isExternalDeclaration(p)
     }
@@ -37,5 +42,15 @@ fun isExternalDeclaration(d: IrDeclaration): Boolean {
 /**
  * Returns true if `d` is not itself a class, but is a member of an external file class.
  */
-fun isExternalFileClassMember(d: IrDeclaration) = d !is IrClass && (d.parentClassOrNull?.let { it.isFileClass } ?: false)
+fun isExternalFileClassMember(d: IrDeclaration): Boolean {
+    if (d is IrClass) { return false }
+    val p = d.parent
+    when (p) {
+        is IrClass -> return p.isFileClass
+        is IrExternalPackageFragment ->
+            // This is an external declaration in a (multi)file class
+            return true
+    }
+    return false
+}