Rework conditional generic extraction to use global state

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

@@ -9,7 +9,7 @@ import java.util.ArrayList
 import java.util.HashSet
 import java.util.zip.GZIPOutputStream
 
-class ExternalClassExtractor(val logger: FileLogger, val invocationTrapFile: String, val sourceFilePath: String, val primitiveTypeMapping: PrimitiveTypeMapping, val pluginContext: IrPluginContext) {
+class ExternalClassExtractor(val logger: FileLogger, val invocationTrapFile: String, val sourceFilePath: String, val primitiveTypeMapping: PrimitiveTypeMapping, val pluginContext: IrPluginContext, val genericSpecialisationsExtracted: MutableSet<String>) {
 
     val externalClassesDone = HashSet<IrClass>()
     val externalClassWorkList = ArrayList<IrClass>()
@@ -43,7 +43,7 @@ class ExternalClassExtractor(val logger: FileLogger, val invocationTrapFile: Str
                                 val binaryPath = getIrClassBinaryPath(irClass)
                                 val ftw = tw.makeFileTrapWriter(binaryPath, true)
 
-                                val fileExtractor = KotlinFileExtractor(logger, ftw, manager, this, primitiveTypeMapping, pluginContext)
+                                val fileExtractor = KotlinFileExtractor(logger, ftw, manager, this, primitiveTypeMapping, pluginContext, genericSpecialisationsExtracted)
                                 fileExtractor.extractClassSource(irClass)
                             }
                         }

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

@@ -42,10 +42,11 @@ class KotlinExtractorExtension(
             FileUtil.logger = logger
             val srcDir = File(System.getenv("CODEQL_EXTRACTOR_JAVA_SOURCE_ARCHIVE_DIR").takeUnless { it.isNullOrEmpty() } ?: "kotlin-extractor/src")
             srcDir.mkdirs()
+            val genericSpecialisationsExtracted = HashSet<String>()
             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, primitiveTypeMapping, pluginContext)
+                doFile(invocationTrapFile, fileTrapWriter, checkTrapIdentical, logCounter, trapDir, srcDir, file, primitiveTypeMapping, pluginContext, genericSpecialisationsExtracted)
             }
             logger.printLimitedWarningCounts()
             // We don't want the compiler to continue and generate class
@@ -92,7 +93,8 @@ fun doFile(invocationTrapFile: String,
            srcDir: File,
            file: IrFile,
            primitiveTypeMapping: PrimitiveTypeMapping,
-           pluginContext: IrPluginContext) {
+           pluginContext: IrPluginContext,
+           genericSpecialisationsExtracted: MutableSet<String>) {
     val filePath = file.path
     val logger = FileLogger(logCounter, fileTrapWriter)
     logger.info("Extracting file $filePath")
@@ -120,8 +122,8 @@ fun doFile(invocationTrapFile: String,
             // Now elevate to a SourceFileTrapWriter, and populate the
             // file information
             val sftw = tw.makeSourceFileTrapWriter(file, true)
-            val externalClassExtractor = ExternalClassExtractor(logger, invocationTrapFile, file.path, primitiveTypeMapping, pluginContext)
-            val fileExtractor = KotlinSourceFileExtractor(logger, sftw, file, externalClassExtractor, primitiveTypeMapping, pluginContext)
+            val externalClassExtractor = ExternalClassExtractor(logger, invocationTrapFile, file.path, primitiveTypeMapping, pluginContext, genericSpecialisationsExtracted)
+            val fileExtractor = KotlinSourceFileExtractor(logger, sftw, file, externalClassExtractor, primitiveTypeMapping, pluginContext, genericSpecialisationsExtracted)
             fileExtractor.extractFileContents(sftw.fileId)
             externalClassExtractor.extractExternalClasses()
         }

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

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.interpreter.toIrConst
 import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
 import org.jetbrains.kotlin.ir.symbols.IrTypeParameterSymbol
 import org.jetbrains.kotlin.ir.types.*
@@ -28,8 +27,9 @@ open class KotlinFileExtractor(
     dependencyCollector: OdasaOutput.TrapFileManager?,
     externalClassExtractor: ExternalClassExtractor,
     primitiveTypeMapping: PrimitiveTypeMapping,
-    pluginContext: IrPluginContext
-): KotlinUsesExtractor(logger, tw, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext) {
+    pluginContext: IrPluginContext,
+    genericSpecialisationsExtracted: MutableSet<String>
+): KotlinUsesExtractor(logger, tw, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext, genericSpecialisationsExtracted) {
 
     fun extractDeclaration(declaration: IrDeclaration, parentId: Label<out DbReftype>) {
         when (declaration) {
@@ -92,15 +92,15 @@ open class KotlinFileExtractor(
         return id
     }
 
-    fun extractClassInstance(c: IrClass, typeArgs: List<IrTypeArgument>, extractFunctionPrototypes: Boolean): Label<out DbClassorinterface> {
+    fun extractClassInstance(c: IrClass, typeArgs: List<IrTypeArgument>): Label<out DbClassorinterface> {
         if (typeArgs.isEmpty()) {
             logger.warn(Severity.ErrorSevere, "Instance without type arguments: " + c.name.asString())
         }
 
-        val results = addClassLabel(c, typeArgs)
-        val id = results.id
+        val classLabelResults = getClassLabel(c, typeArgs)
+        val id = tw.getLabelFor<DbClassorinterface>(classLabelResults.classLabel)
         val pkg = c.packageFqName?.asString() ?: ""
-        val cls = results.shortName
+        val cls = classLabelResults.shortName
         val pkgId = extractPackage(pkg)
         if(c.kind == ClassKind.INTERFACE) {
             @Suppress("UNCHECKED_CAST")
@@ -133,19 +133,19 @@ open class KotlinFileExtractor(
         val locId = tw.getLocation(c)
         tw.writeHasLocation(id, locId)
 
-        if (extractFunctionPrototypes) {
-            val typeParamSubstitution = c.typeParameters.map({ it.symbol }).zip(typeArgs).toMap()
+        return id
+    }
 
-            c.declarations.map {
-                when(it) {
-                    is IrFunction -> extractFunction(it, id, false, typeParamSubstitution)
-                    is IrProperty -> extractProperty(it, id, false, typeParamSubstitution)
-                    else -> {}
-                }
+    fun extractMemberPrototypes(c: IrClass, typeArgs: List<IrTypeArgument>, id: Label<out DbClassorinterface>) {
+        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
     }
 
     private fun extracLocalTypeDeclStmt(c: IrClass, callable: Label<out DbCallable>, parent: Label<out DbStmtparent>, idx: Int) {

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

@@ -17,9 +17,10 @@ class KotlinSourceFileExtractor(
     val file: IrFile,
     externalClassExtractor: ExternalClassExtractor,
     primitiveTypeMapping: PrimitiveTypeMapping,
-    pluginContext: IrPluginContext
+    pluginContext: IrPluginContext,
+    genericSpecialisationsExtracted: MutableSet<String>
 ) :
-  KotlinFileExtractor(logger, tw, null, externalClassExtractor, primitiveTypeMapping, pluginContext) {
+  KotlinFileExtractor(logger, tw, null, externalClassExtractor, primitiveTypeMapping, pluginContext, genericSpecialisationsExtracted) {
 
     val fileClass by lazy {
         extractFileClass(file)

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

@@ -23,7 +23,8 @@ open class KotlinUsesExtractor(
     val dependencyCollector: OdasaOutput.TrapFileManager?,
     val externalClassExtractor: ExternalClassExtractor,
     val primitiveTypeMapping: PrimitiveTypeMapping,
-    val pluginContext: IrPluginContext
+    val pluginContext: IrPluginContext,
+    val genericSpecialisationsExtracted: MutableSet<String>
 ) {
     fun usePackage(pkg: String): Label<out DbPackage> {
         return extractPackage(pkg)
@@ -77,7 +78,7 @@ open class KotlinUsesExtractor(
         if (isExternalDeclaration(cls) || clsFile == null) {
             val newTrapWriter = tw.makeFileTrapWriter(getIrClassBinaryPath(cls), false)
             val newLogger = FileLogger(logger.logCounter, newTrapWriter)
-            return KotlinFileExtractor(newLogger, newTrapWriter, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext)
+            return KotlinFileExtractor(newLogger, newTrapWriter, dependencyCollector, externalClassExtractor, primitiveTypeMapping, pluginContext, genericSpecialisationsExtracted)
         }
 
         if (this is KotlinSourceFileExtractor && this.file == clsFile) {
@@ -86,16 +87,9 @@ open class KotlinUsesExtractor(
 
         val newTrapWriter = tw.makeSourceFileTrapWriter(clsFile, false)
         val newLogger = FileLogger(logger.logCounter, newTrapWriter)
-        return KotlinSourceFileExtractor(newLogger, newTrapWriter, clsFile, externalClassExtractor, primitiveTypeMapping, pluginContext)
+        return KotlinSourceFileExtractor(newLogger, newTrapWriter, clsFile, externalClassExtractor, primitiveTypeMapping, pluginContext, genericSpecialisationsExtracted)
     }
 
-    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")
@@ -150,22 +144,26 @@ open class KotlinUsesExtractor(
     fun addClassLabel(c: IrClass, typeArgs: List<IrTypeArgument>, inReceiverContext: Boolean = false): TypeResult<DbClassorinterface> {
         val classLabelResult = getClassLabel(c, typeArgs)
 
-        var shouldExtractClass = false
+        var instanceSeenBefore = true
 
         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
-            shouldExtractClass = true
+            instanceSeenBefore = false
 
             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)
+            // If this is a generic type instantiation then it has no
+            // source entity, so we need to extract it here
+            val extractorWithCSource by lazy { this.withSourceFileOfClass(c) }
+
+            if (!instanceSeenBefore) {
+                extractorWithCSource.extractClassInstance(c, typeArgs)
+            }
+
+            if (inReceiverContext && genericSpecialisationsExtracted.add(classLabelResult.classLabel)) {
+                extractorWithCSource.extractMemberPrototypes(c, typeArgs, classLabel)
+            }
         }
 
         return TypeResult(