Kotlin: stub trap .class files when extracting a class from Kotlin source

java/kotlin-extractor/src/main/java/com/semmle/extractor/java/OdasaOutput.java

@@ -19,9 +19,10 @@
 import java.util.zip.ZipFile;
 
 import com.github.codeql.Logger;
-import static com.github.codeql.ClassNamesKt.getIrDeclBinaryName;
+import static com.github.codeql.ClassNamesKt.getIrElementBinaryName;
 import static com.github.codeql.ClassNamesKt.getIrClassVirtualFile;
 
+import org.jetbrains.kotlin.ir.IrElement;
 import org.jetbrains.kotlin.ir.declarations.IrClass;
 
 import com.intellij.openapi.vfs.VirtualFile;
@@ -212,20 +213,19 @@ private File trapFileFor(File file) {
 				PathTransformer.std().fileAsDatabaseString(file) + ".trap.gz");
 	}
 
-	private File getTrapFileForDecl(IrDeclaration sym, String signature) {
+	private File getTrapFileForDecl(IrElement sym, String signature) {
 		if (currentSpecFileEntry == null)
 			return null;
 		return trapFileForDecl(sym, signature);
 	}
 
-	private File trapFileForDecl(IrDeclaration sym, String signature) {
+	private File trapFileForDecl(IrElement sym, String signature) {
 		return FileUtil.fileRelativeTo(currentSpecFileEntry.getTrapFolder(),
 				trapFilePathForDecl(sym, signature));
 	}
 
-	private String trapFilePathForDecl(IrDeclaration sym, String signature) {
-		String binaryName = getIrDeclBinaryName(sym);
-		String binaryNameWithSignature = binaryName + signature;
+	private String trapFilePathForDecl(IrElement sym, String signature) {
+		String binaryName = getIrElementBinaryName(sym);
 		// TODO: Reinstate this?
 		//if (getTrackClassOrigins())
 		//  classId += "-" + StringDigestor.digest(sym.getSourceFileId());
@@ -241,7 +241,7 @@ private String trapFilePathForDecl(IrDeclaration sym, String signature) {
 	 * Deletion of existing trap files.
 	 */
 
-	private void deleteTrapFileAndDependencies(IrDeclaration sym, String signature) {
+	private void deleteTrapFileAndDependencies(IrElement sym, String signature) {
 		File trap = trapFileForDecl(sym, signature);
 		if (trap.exists()) {
 			trap.delete();
@@ -269,7 +269,7 @@ private void deleteTrapFileAndDependencies(IrDeclaration sym, String signature)
 	 * 		Any unique suffix needed to distinguish `sym` from other declarations with the same name.
 	 * 		For functions for example, this means its parameter signature.
 	 */
-	private TrapFileManager getMembersWriterForDecl(File trap, File trapFileBase, TrapClassVersion trapFileVersion, IrDeclaration sym, String signature) {
+	private TrapFileManager getMembersWriterForDecl(File trap, File trapFileBase, TrapClassVersion trapFileVersion, IrElement sym, String signature) {
 		if (use_trap_locking) {
 			TrapClassVersion currVersion = TrapClassVersion.fromSymbol(sym, log);
 			String shortName = sym instanceof IrDeclarationWithName ? ((IrDeclarationWithName)sym).getName().asString() : "(name unknown)";
@@ -326,7 +326,7 @@ private TrapFileManager getMembersWriterForDecl(File trap, File trapFileBase, Tr
 		return trapWriter(trap, sym, signature);
 	}
 
-	private TrapFileManager trapWriter(File trapFile, IrDeclaration sym, String signature) {
+	private TrapFileManager trapWriter(File trapFile, IrElement sym, String signature) {
 		if (!trapFile.getName().endsWith(".trap.gz"))
 			throw new CatastrophicError("OdasaOutput only supports writing to compressed trap files");
 		String relative = FileUtil.relativePath(trapFile, currentSpecFileEntry.getTrapFolder());
@@ -335,7 +335,7 @@ private TrapFileManager trapWriter(File trapFile, IrDeclaration sym, String sign
 		return concurrentWriter(trapFile, relative, log, sym, signature);
 	}
 
-	private TrapFileManager concurrentWriter(File trapFile, String relative, Logger log, IrDeclaration sym, String signature) {
+	private TrapFileManager concurrentWriter(File trapFile, String relative, Logger log, IrElement sym, String signature) {
 		if (trapFile.exists())
 			return null;
 		return new TrapFileManager(trapFile, relative, true, log, sym, signature);
@@ -345,11 +345,11 @@ public class TrapFileManager implements AutoCloseable {
 
 		private TrapDependencies trapDependenciesForClass;
 		private File trapFile;
-		private IrDeclaration sym;
+		private IrElement sym;
 		private String signature;
 		private boolean hasError = false;
 
-		private TrapFileManager(File trapFile, String relative, boolean concurrentCreation, Logger log, IrDeclaration sym, String signature) {
+		private TrapFileManager(File trapFile, String relative, boolean concurrentCreation, Logger log, IrElement sym, String signature) {
 			trapDependenciesForClass = new TrapDependencies(relative);
 			this.trapFile = trapFile;
 			this.sym = sym;
@@ -360,7 +360,7 @@ public File getFile() {
 			return trapFile;
 		}
 
-		public void addDependency(IrDeclaration dep, String signature) {
+		public void addDependency(IrElement dep, String signature) {
 			trapDependenciesForClass.addDependency(trapFilePathForDecl(dep, signature));
 		}
 
@@ -422,7 +422,7 @@ public void setHasError() {
 	 * previously set by a call to {@link OdasaOutput#setCurrentSourceFile(File)}.
 	 */
 	public TrapLocker getTrapLockerForCurrentSourceFile() {
-		return new TrapLocker((IrClass)null, null);
+		return new TrapLocker((IrClass)null, null, true);
 	}
 
 	/**
@@ -460,19 +460,19 @@ public TrapLocker getTrapLockerForModule(String moduleName) {
 	 *
 	 * @return  a {@link TrapLocker} for the trap file corresponding to the given class symbol.
 	 */
-	public TrapLocker getTrapLockerForDecl(IrDeclaration sym, String signature) {
-		return new TrapLocker(sym, signature);
+	public TrapLocker getTrapLockerForDecl(IrElement sym, String signature, boolean fromSource) {
+		return new TrapLocker(sym, signature, fromSource);
 	}
 
 	public class TrapLocker implements AutoCloseable {
-		private final IrDeclaration sym;
+		private final IrElement sym;
 		private final File trapFile;
 		// trapFileBase is used when doing lockless TRAP file writing.
 		// It is trapFile without the #metadata.trap.gz suffix.
 		private File trapFileBase = null;
 		private TrapClassVersion trapFileVersion = null;
 		private final String signature;
-		private TrapLocker(IrDeclaration decl, String signature) {
+		private TrapLocker(IrElement decl, String signature, boolean fromSource) {
 			this.sym = decl;
 			this.signature = signature;
 			if (sym==null) {
@@ -485,7 +485,10 @@ private TrapLocker(IrDeclaration decl, String signature) {
 				} else {
 					// We encode the metadata into the filename, so that the
 					// TRAP filenames for different metadatas don't overlap.
-					trapFileVersion = TrapClassVersion.fromSymbol(sym, log);
+					if (fromSource)
+						trapFileVersion = new TrapClassVersion(0, 0, 0, "kotlin");
+					else
+						trapFileVersion = TrapClassVersion.fromSymbol(sym, log);
 					String baseName = normalTrapFile.getName().replace(".trap.gz", "");
 					// If a class has lots of inner classes, then we get lots of files
 					// in a single directory. This makes our directory listings later slow.
@@ -717,11 +720,18 @@ private static long getVirtualFileTimeStamp(VirtualFile vf, Logger log) {
             return vf.getTimeStamp();
         }
 
-		private static TrapClassVersion fromSymbol(IrDeclaration sym, Logger log) {
-			VirtualFile vf = sym instanceof IrClass ? getIrClassVirtualFile((IrClass)sym) :
-					sym.getParent() instanceof IrClass ? getIrClassVirtualFile((IrClass)sym.getParent()) :
-					null;
-			if(vf == null)
+		private static VirtualFile getVirtualFileIfClass(IrElement e) {
+			if (e instanceof IrClass)
+				return getIrClassVirtualFile((IrClass)e);
+			else
+				return null;
+		}
+
+		private static TrapClassVersion fromSymbol(IrElement sym, Logger log) {
+			VirtualFile vf = getVirtualFileIfClass(sym);
+			if (vf == null && sym instanceof IrDeclaration)
+				vf = getVirtualFileIfClass(((IrDeclaration)sym).getParent());
+			if (vf == null)
 				return new TrapClassVersion(-1, 0, 0, null);
 
 			final int[] versionStore = new int[1];

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

@@ -1,16 +1,14 @@
 package com.github.codeql
 
-import com.github.codeql.utils.isExternalDeclaration
 import com.github.codeql.utils.isExternalFileClassMember
 import com.semmle.extractor.java.OdasaOutput
 import com.semmle.util.data.StringDigestor
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+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.isFileClass
 import org.jetbrains.kotlin.ir.util.packageFqName
-import org.jetbrains.kotlin.ir.util.parentClassOrNull
-import org.jetbrains.kotlin.name.FqName
+import java.io.BufferedWriter
 import java.io.File
 import java.util.ArrayList
 import java.util.HashSet
@@ -25,87 +23,103 @@ class ExternalDeclExtractor(val logger: FileLogger, val invocationTrapFile: Stri
     val propertySignature = ";property"
     val fieldSignature = ";field"
 
+    val output = OdasaOutput(false, logger).also {
+        it.setCurrentSourceFile(File(sourceFilePath))
+    }
+
     fun extractLater(d: IrDeclarationWithName, signature: String): Boolean {
         if (d !is IrClass && !isExternalFileClassMember(d)) {
             logger.errorElement("External declaration is neither a class, nor a top-level declaration", d)
             return false
         }
-        val declBinaryName = declBinaryNames.getOrPut(d) { getIrDeclBinaryName(d) }
+        val declBinaryName = declBinaryNames.getOrPut(d) { getIrElementBinaryName(d) }
         val ret = externalDeclsDone.add(Pair(declBinaryName, signature))
         if (ret) externalDeclWorkList.add(Pair(d, signature))
         return ret
     }
     fun extractLater(c: IrClass) = extractLater(c, "")
 
+    fun writeStubTrapFile(e: IrElement, signature: String = "") {
+        extractElement(e, signature, true) { trapFileBW, _, _ ->
+            trapFileBW.write("// Trap file stubbed because this declaration was extracted from source in $sourceFilePath\n")
+            trapFileBW.write("// Part of invocation $invocationTrapFile\n")
+        }
+    }
+
+    private fun extractElement(element: IrElement, possiblyLongSignature: String, fromSource: Boolean, extractorFn: (BufferedWriter, String, OdasaOutput.TrapFileManager) -> Unit) {
+        // In order to avoid excessively long signatures which can lead to trap file names longer than the filesystem
+        // limit, we truncate and add a hash to preserve uniqueness if necessary.
+        val signature = if (possiblyLongSignature.length > 100) {
+            possiblyLongSignature.substring(0, 92) + "#" + StringDigestor.digest(possiblyLongSignature).substring(0, 8)
+        } else { possiblyLongSignature }
+        output.getTrapLockerForDecl(element, signature, fromSource).useAC { locker ->
+            locker.trapFileManager.useAC { manager ->
+                val shortName = when(element) {
+                    is IrDeclarationWithName -> element.name.asString()
+                    is IrFile -> element.name
+                    else -> "(unknown name)"
+                }
+                if (manager == null) {
+                    logger.info("Skipping extracting external decl $shortName")
+                } else {
+                    val trapFile = manager.file
+                    val trapTmpFile = File.createTempFile("${trapFile.nameWithoutExtension}.", ".${trapFile.extension}.tmp", trapFile.parentFile)
+                    try {
+                        GZIPOutputStream(trapTmpFile.outputStream()).bufferedWriter().use {
+                            extractorFn(it, signature, manager)
+                        }
+
+                        if (!trapTmpFile.renameTo(trapFile)) {
+                            logger.error("Failed to rename $trapTmpFile to $trapFile")
+                        }
+                    } catch (e: Exception) {
+                        manager.setHasError()
+                        logger.error("Failed to extract '$shortName'. Partial TRAP file location is $trapTmpFile", e)
+                    }
+                }
+            }
+        }
+    }
+
     fun extractExternalClasses() {
-        val output = OdasaOutput(false, logger)
-        output.setCurrentSourceFile(File(sourceFilePath))
         do {
             val nextBatch = ArrayList(externalDeclWorkList)
             externalDeclWorkList.clear()
             nextBatch.forEach { workPair ->
                 val (irDecl, possiblyLongSignature) = workPair
-                // In order to avoid excessively long signatures which can lead to trap file names longer than the filesystem
-                // limit, we truncate and add a hash to preserve uniqueness if necessary.
-                val signature = if (possiblyLongSignature.length > 100) {
-                    possiblyLongSignature.substring(0, 92) + "#" + StringDigestor.digest(possiblyLongSignature).substring(0, 8)
-                } else { possiblyLongSignature }
-                output.getTrapLockerForDecl(irDecl, signature).useAC { locker ->
-                    locker.trapFileManager.useAC { manager ->
-                        val shortName = when(irDecl) {
-                            is IrDeclarationWithName -> irDecl.name.asString()
-                            else -> "(unknown name)"
-                        }
-                        if(manager == null) {
-                            logger.info("Skipping extracting external decl $shortName")
-                        } else {
-                            val trapFile = manager.file
-                            val trapTmpFile = File.createTempFile("${trapFile.nameWithoutExtension}.", ".${trapFile.extension}.tmp", trapFile.parentFile)
-
-                            val containingClass = getContainingClassOrSelf(irDecl)
-                            if (containingClass == null) {
-                                logger.errorElement("Unable to get containing class", irDecl)
-                                return
-                            }
-                            val binaryPath = getIrClassBinaryPath(containingClass)
-                            try {
-                                GZIPOutputStream(trapTmpFile.outputStream()).bufferedWriter().use { trapFileBW ->
-                                    // We want our comments to be the first thing in the file,
-                                    // so start off with a mere TrapWriter
-                                    val tw = TrapWriter(logger.loggerBase, TrapLabelManager(), trapFileBW, diagnosticTrapWriter)
-                                    tw.writeComment("Generated by the CodeQL Kotlin extractor for external dependencies")
-                                    tw.writeComment("Part of invocation $invocationTrapFile")
-                                    if (signature != possiblyLongSignature) {
-                                        tw.writeComment("Function signature abbreviated; full signature is: $possiblyLongSignature")
-                                    }
-                                    // Now elevate to a SourceFileTrapWriter, and populate the
-                                    // file information if needed:
-                                    val ftw = tw.makeFileTrapWriter(binaryPath, true)
+                extractElement(irDecl, possiblyLongSignature, false) { trapFileBW, signature, manager ->
+                    val containingClass = getContainingClassOrSelf(irDecl)
+                    if (containingClass == null) {
+                        logger.errorElement("Unable to get containing class", irDecl)
+                    } else {
+                        val binaryPath = getIrClassBinaryPath(containingClass)
 
-                                    val fileExtractor = KotlinFileExtractor(logger, ftw, null, binaryPath, manager, this, primitiveTypeMapping, pluginContext, KotlinFileExtractor.DeclarationStack(), globalExtensionState)
+                        // We want our comments to be the first thing in the file,
+                        // so start off with a mere TrapWriter
+                        val tw = TrapWriter(logger.loggerBase, TrapLabelManager(), trapFileBW, diagnosticTrapWriter)
+                        tw.writeComment("Generated by the CodeQL Kotlin extractor for external dependencies")
+                        tw.writeComment("Part of invocation $invocationTrapFile")
+                        if (signature != possiblyLongSignature) {
+                            tw.writeComment("Function signature abbreviated; full signature is: $possiblyLongSignature")
+                        }
+                        // Now elevate to a SourceFileTrapWriter, and populate the
+                        // file information if needed:
+                        val ftw = tw.makeFileTrapWriter(binaryPath, true)
 
-                                    if (irDecl is IrClass) {
-                                        // Populate a location and compilation-unit package for the file. This is similar to
-                                        // the beginning of `KotlinFileExtractor.extractFileContents` but without an `IrFile`
-                                        // to start from.
-                                        val pkg = irDecl.packageFqName?.asString() ?: ""
-                                        val pkgId = fileExtractor.extractPackage(pkg)
-                                        ftw.writeHasLocation(ftw.fileId, ftw.getWholeFileLocation())
-                                        ftw.writeCupackage(ftw.fileId, pkgId)
+                        val fileExtractor = KotlinFileExtractor(logger, ftw, null, binaryPath, manager, this, primitiveTypeMapping, pluginContext, KotlinFileExtractor.DeclarationStack(), globalExtensionState)
 
-                                        fileExtractor.extractClassSource(irDecl, extractDeclarations = !irDecl.isFileClass, extractStaticInitializer = false, extractPrivateMembers = false, extractFunctionBodies = false)
-                                    } else {
-                                        fileExtractor.extractDeclaration(irDecl, extractPrivateMembers = false, extractFunctionBodies = false)
-                                    }
-                                }
+                        if (irDecl is IrClass) {
+                            // Populate a location and compilation-unit package for the file. This is similar to
+                            // the beginning of `KotlinFileExtractor.extractFileContents` but without an `IrFile`
+                            // to start from.
+                            val pkg = irDecl.packageFqName?.asString() ?: ""
+                            val pkgId = fileExtractor.extractPackage(pkg)
+                            ftw.writeHasLocation(ftw.fileId, ftw.getWholeFileLocation())
+                            ftw.writeCupackage(ftw.fileId, pkgId)
 
-                                if (!trapTmpFile.renameTo(trapFile)) {
-                                    logger.error("Failed to rename $trapTmpFile to $trapFile")
-                                }
-                            } catch (e: Exception) {
-                                manager.setHasError()
-                                logger.error("Failed to extract '$shortName'. Partial TRAP file location is $trapTmpFile", e)
-                            }
+                            fileExtractor.extractClassSource(irDecl, extractDeclarations = !irDecl.isFileClass, extractStaticInitializer = false, extractPrivateMembers = false, extractFunctionBodies = false)
+                        } else {
+                            fileExtractor.extractDeclaration(irDecl, extractPrivateMembers = false, extractFunctionBodies = false)
                         }
                     }
                 }