Kotlin: Use an IrClass for the Java class in the priomitive type handling

This means we can add a label for it, and indicate that we use it,
properly.

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

@@ -3,12 +3,13 @@ package com.github.codeql
 import com.semmle.extractor.java.OdasaOutput
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.util.IdSignature
 import java.io.File
 import java.util.ArrayList
 import java.util.HashSet
 import java.util.zip.GZIPOutputStream
 
-class ExternalClassExtractor(val logger: FileLogger, val sourceFilePath: String, val pluginContext: IrPluginContext) {
+class ExternalClassExtractor(val logger: FileLogger, val sourceFilePath: String, val primitiveTypeMapping: Map<IdSignature.PublicSignature, PrimitiveTypeInfo>, val pluginContext: IrPluginContext) {
 
     val externalClassesDone = HashSet<IrClass>()
     val externalClassWorkList = ArrayList<IrClass>()
@@ -34,7 +35,7 @@ class ExternalClassExtractor(val logger: FileLogger, val sourceFilePath: String,
                             GZIPOutputStream(manager.getFile().outputStream()).bufferedWriter().use { trapFileBW ->
                                 val tw =
                                     FileTrapWriter(TrapLabelManager(), trapFileBW, getIrClassBinaryPath(irClass), true)
-                                val fileExtractor = KotlinFileExtractor(logger, tw, manager, this, pluginContext)
+                                val fileExtractor = KotlinFileExtractor(logger, tw, manager, this, primitiveTypeMapping, pluginContext)
                                 fileExtractor.extractClassSource(irClass)
                             }
                         }

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

@@ -3,6 +3,7 @@ package com.github.codeql
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.util.*
 import java.io.File
 import java.io.FileOutputStream
 import java.nio.file.Files
@@ -35,6 +36,7 @@ class KotlinExtractorExtension(
             val logger = Logger(logCounter, tw)
             logger.info("Extraction started")
             logger.flush()
+            val primitiveTypeMapping = makePrimitiveTypeMapping(logger, pluginContext)
             // FIXME: FileUtil expects a static global logger
             // which should be provided by SLF4J's factory facility. For now we set it here.
             FileUtil.logger = logger
@@ -43,7 +45,7 @@ class KotlinExtractorExtension(
             moduleFragment.files.mapIndexed { index: Int, file: IrFile ->
                 val fileTrapWriter = tw.makeSourceFileTrapWriter(file, true)
                 fileTrapWriter.writeCompilation_compiling_files(compilation, index, fileTrapWriter.fileId)
-                doFile(invocationTrapFile, fileTrapWriter, checkTrapIdentical, logCounter, trapDir, srcDir, file, pluginContext)
+                doFile(invocationTrapFile, fileTrapWriter, checkTrapIdentical, logCounter, trapDir, srcDir, file, primitiveTypeMapping, pluginContext)
             }
             logger.printLimitedWarningCounts()
             // We don't want the compiler to continue and generate class
@@ -89,6 +91,7 @@ fun doFile(invocationTrapFile: String,
            trapDir: File,
            srcDir: File,
            file: IrFile,
+           primitiveTypeMapping: Map<IdSignature.PublicSignature, PrimitiveTypeInfo>,
            pluginContext: IrPluginContext) {
     val filePath = file.path
     val logger = FileLogger(logCounter, fileTrapWriter)
@@ -111,8 +114,8 @@ fun doFile(invocationTrapFile: String,
         trapTmpFile.bufferedWriter().use { trapFileBW ->
             val tw = SourceFileTrapWriter(TrapLabelManager(), trapFileBW, file, true)
             tw.writeComment("Generated by invocation $invocationTrapFile")
-            val externalClassExtractor = ExternalClassExtractor(logger, file.path, pluginContext)
-            val fileExtractor = KotlinSourceFileExtractor(logger, tw, file, externalClassExtractor, pluginContext)
+            val externalClassExtractor = ExternalClassExtractor(logger, file.path, primitiveTypeMapping, pluginContext)
+            val fileExtractor = KotlinSourceFileExtractor(logger, tw, file, externalClassExtractor, primitiveTypeMapping, pluginContext)
             fileExtractor.extractFileContents(tw.fileId)
             externalClassExtractor.extractExternalClasses()
         }

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

@@ -17,8 +17,9 @@ open class KotlinFileExtractor(
     override val tw: FileTrapWriter,
     dependencyCollector: OdasaOutput.TrapFileManager?,
     externalClassExtractor: ExternalClassExtractor,
+    primitiveTypeMapping: Map<IdSignature.PublicSignature, PrimitiveTypeInfo>,
     pluginContext: IrPluginContext
-): KotlinUsesExtractor(logger, tw, dependencyCollector, externalClassExtractor, pluginContext) {
+): KotlinUsesExtractor(logger, tw, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext) {
 
     fun extractDeclaration(declaration: IrDeclaration, parentId: Label<out DbReftype>) {
         when (declaration) {

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

@@ -8,15 +8,17 @@ import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.util.packageFqName
+import org.jetbrains.kotlin.ir.util.IdSignature
 
 class KotlinSourceFileExtractor(
     logger: FileLogger,
     tw: FileTrapWriter,
     val file: IrFile,
     externalClassExtractor: ExternalClassExtractor,
+    primitiveTypeMapping: Map<IdSignature.PublicSignature, PrimitiveTypeInfo>,
     pluginContext: IrPluginContext
 ) :
-  KotlinFileExtractor(logger, tw, null, externalClassExtractor, pluginContext) {
+  KotlinFileExtractor(logger, tw, null, externalClassExtractor, primitiveTypeMapping, pluginContext) {
 
     val fileClass by lazy {
         extractFileClass(file)
@@ -65,4 +67,4 @@ class KotlinSourceFileExtractor(
       return id
   }
 
-}
+}

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

@@ -19,6 +19,7 @@ open class KotlinUsesExtractor(
     open val tw: TrapWriter,
     val dependencyCollector: OdasaOutput.TrapFileManager?,
     val externalClassExtractor: ExternalClassExtractor,
+    val primitiveTypeMapping: Map<IdSignature.PublicSignature, PrimitiveTypeInfo>,
     val pluginContext: IrPluginContext
 ) {
     fun usePackage(pkg: String): Label<out DbPackage> {
@@ -78,7 +79,7 @@ open class KotlinUsesExtractor(
 
         val newLogger = FileLogger(logger.logCounter, newTrapWriter)
 
-        return KotlinFileExtractor(newLogger, newTrapWriter, dependencyCollector, externalClassExtractor, pluginContext)
+        return KotlinFileExtractor(newLogger, newTrapWriter, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext)
     }
 
     fun useClassInstance(c: IrClass, typeArgs: List<IrTypeArgument>): UseClassInstanceResult {
@@ -284,7 +285,7 @@ open class KotlinUsesExtractor(
         }
         fun primitiveType(kotlinClass: IrClass, primitiveName: String?,
                           otherIsPrimitive: Boolean,
-                          javaPackageName: String, javaClassName: String,
+                          javaClass: IrClass,
                           kotlinPackageName: String, kotlinClassName: String): TypeResults {
             val javaResult = if ((context == TypeContext.RETURN || (context == TypeContext.OTHER && otherIsPrimitive)) && !s.hasQuestionMark && primitiveName != null) {
                     val label: Label<DbPrimitive> = tw.getLabelFor("@\"type;$primitiveName\"", {
@@ -292,9 +293,8 @@ open class KotlinUsesExtractor(
                     })
                     TypeResult(label, primitiveName, primitiveName)
                 } else {
-                    val label = makeClass(javaPackageName, javaClassName)
-                    val signature = "$javaPackageName.$javaClassName"
-                    TypeResult(label, signature, javaClassName)
+                    extractClassLaterIfExternal(javaClass)
+                    addClassLabel(javaClass, listOf())
                 }
             val kotlinClassId = useClassInstance(kotlinClass, listOf()).typeResult.id
             val kotlinResult = if (s.hasQuestionMark) {
@@ -332,7 +332,7 @@ XXX delete?
             primitiveInfo != null -> return primitiveType(
                 s.classifier.owner as IrClass,
                 primitiveInfo.primitiveName, primitiveInfo.otherIsPrimitive,
-                primitiveInfo.javaPackageName, primitiveInfo.javaClassName,
+                primitiveInfo.javaClass,
                 primitiveInfo.kotlinPackageName, primitiveInfo.kotlinClassName
             )
 /*

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

@@ -1,32 +1,79 @@
 package com.github.codeql
 
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.types.IdSignatureValues
+import org.jetbrains.kotlin.ir.util.IdSignature
+import org.jetbrains.kotlin.name.FqName
 
 data class PrimitiveTypeInfo(
     val primitiveName: String?,
     val otherIsPrimitive: Boolean,
-    val javaPackageName: String, val javaClassName: String,
+    val javaClass: IrClass,
     val kotlinPackageName: String, val kotlinClassName: String
 )
 
-val primitiveTypeMapping = mapOf(
-    IdSignatureValues._byte to PrimitiveTypeInfo("byte", true, "java.lang", "Byte", "kotlin", "Byte"),
-    IdSignatureValues._short to PrimitiveTypeInfo("short", true, "java.lang", "Short", "kotlin", "Short"),
-    IdSignatureValues._int to PrimitiveTypeInfo("int", true, "java.lang", "Integer", "kotlin", "Int"),
-    IdSignatureValues._long to PrimitiveTypeInfo("long", true, "java.lang", "Long", "kotlin", "Long"),
+private fun findClass(fqName: String, fallback: IrClass, logger: Logger, pluginContext: IrPluginContext): IrClass {
+    val symbol = pluginContext.referenceClass(FqName(fqName))
+    if(symbol == null) {
+        logger.warn("Can't find $fqName")
+        // Do the best we can
+        return fallback
+    } else {
+       return symbol.owner
+    }
+}
 
-    IdSignatureValues.uByte to PrimitiveTypeInfo("byte", true, "kotlin", "UByte", "kotlin", "UByte"),
-    IdSignatureValues.uShort to PrimitiveTypeInfo("short", true, "kotlin", "UShort", "kotlin", "UShort"),
-    IdSignatureValues.uInt to PrimitiveTypeInfo("int", true, "kotlin", "UInt", "kotlin", "UInt"),
-    IdSignatureValues.uLong to PrimitiveTypeInfo("long", true, "kotlin", "ULong", "kotlin", "ULong"),
+fun makePrimitiveTypeMapping(logger: Logger, pluginContext: IrPluginContext): Map<IdSignature.PublicSignature, PrimitiveTypeInfo> {
+    val kotlinByte = pluginContext.irBuiltIns.byteClass.owner
+    val javaLangByte = findClass("java.lang.Byte", kotlinByte, logger, pluginContext)
+    val kotlinShort = pluginContext.irBuiltIns.shortClass.owner
+    val javaLangShort = findClass("java.lang.Short", kotlinShort, logger, pluginContext)
+    val kotlinInt = pluginContext.irBuiltIns.intClass.owner
+    val javaLangInteger = findClass("java.lang.Integer", kotlinInt, logger, pluginContext)
+    val kotlinLong = pluginContext.irBuiltIns.longClass.owner
+    val javaLangLong = findClass("java.lang.Long", kotlinLong, logger, pluginContext)
 
-    IdSignatureValues._double to PrimitiveTypeInfo("double", true, "java.lang", "Double", "kotlin", "Double"),
-    IdSignatureValues._float to PrimitiveTypeInfo("float", true, "java.lang", "Float", "kotlin", "Float"),
+    val kotlinUByte = findClass("kotlin.UByte", kotlinByte, logger, pluginContext)
+    val kotlinUShort = findClass("kotlin.UShort", kotlinShort, logger, pluginContext)
+    val kotlinUInt = findClass("kotlin.UInt", kotlinInt, logger, pluginContext)
+    val kotlinULong = findClass("kotlin.ULong", kotlinLong, logger, pluginContext)
 
-    IdSignatureValues._boolean to PrimitiveTypeInfo("boolean", true, "java.lang", "Boolean", "kotlin", "Boolean"),
+    val kotlinDouble = pluginContext.irBuiltIns.doubleClass.owner
+    val javaLangDouble = findClass("java.lang.Double", kotlinDouble, logger, pluginContext)
+    val kotlinFloat = pluginContext.irBuiltIns.floatClass.owner
+    val javaLangFloat = findClass("java.lang.Float", kotlinFloat, logger, pluginContext)
 
-    IdSignatureValues._char to PrimitiveTypeInfo("char", true, "java.lang", "Character", "kotlin", "Char"),
+    val kotlinBoolean = pluginContext.irBuiltIns.booleanClass.owner
+    val javaLangBoolean = findClass("java.lang.Boolean", kotlinBoolean, logger, pluginContext)
 
-    IdSignatureValues.unit to PrimitiveTypeInfo("void", false, "kotlin", "Unit", "kotlin", "Unit"),
-    IdSignatureValues.nothing to PrimitiveTypeInfo(null, true, "java.lang", "Void", "kotlin", "Nothing"),
-)
+    val kotlinChar = pluginContext.irBuiltIns.charClass.owner
+    val javaLangCharacter = findClass("java.lang.Character", kotlinChar, logger, pluginContext)
+
+    val kotlinUnit = pluginContext.irBuiltIns.unitClass.owner
+
+    val kotlinNothing = pluginContext.irBuiltIns.nothingClass.owner
+    val javaLangVoid = findClass("java.lang.Void", kotlinNothing, logger, pluginContext)
+
+    return mapOf(
+        IdSignatureValues._byte to PrimitiveTypeInfo("byte", true, javaLangByte, "kotlin", "Byte"),
+        IdSignatureValues._short to PrimitiveTypeInfo("short", true, javaLangShort, "kotlin", "Short"),
+        IdSignatureValues._int to PrimitiveTypeInfo("int", true, javaLangInteger, "kotlin", "Int"),
+        IdSignatureValues._long to PrimitiveTypeInfo("long", true, javaLangLong, "kotlin", "Long"),
+
+        IdSignatureValues.uByte to PrimitiveTypeInfo("byte", true, kotlinUByte, "kotlin", "UByte"),
+        IdSignatureValues.uShort to PrimitiveTypeInfo("short", true, kotlinUShort, "kotlin", "UShort"),
+        IdSignatureValues.uInt to PrimitiveTypeInfo("int", true, kotlinUInt, "kotlin", "UInt"),
+        IdSignatureValues.uLong to PrimitiveTypeInfo("long", true, kotlinULong, "kotlin", "ULong"),
+
+        IdSignatureValues._double to PrimitiveTypeInfo("double", true, javaLangDouble, "kotlin", "Double"),
+        IdSignatureValues._float to PrimitiveTypeInfo("float", true, javaLangFloat, "kotlin", "Float"),
+
+        IdSignatureValues._boolean to PrimitiveTypeInfo("boolean", true, javaLangBoolean, "kotlin", "Boolean"),
+
+        IdSignatureValues._char to PrimitiveTypeInfo("char", true, javaLangCharacter, "kotlin", "Char"),
+
+        IdSignatureValues.unit to PrimitiveTypeInfo("void", false, kotlinUnit, "kotlin", "Unit"),
+        IdSignatureValues.nothing to PrimitiveTypeInfo(null, true, javaLangVoid, "kotlin", "Nothing"),
+    )
+}