Merge pull request #14390 from igfoo/igfoo/compr

Kotlin: Improve support for TRAP compression options

Author: igfoo

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

@@ -49,6 +49,8 @@
 import com.semmle.util.trap.dependencies.TrapSet;
 import com.semmle.util.trap.pathtransformers.PathTransformer;
 
+import com.github.codeql.Compression;
+
 public class OdasaOutput {
 	private final File trapFolder;
 	private final File sourceArchiveFolder;
@@ -63,16 +65,18 @@ public class OdasaOutput {
 	private final boolean trackClassOrigins;
 
 	private final Logger log;
+	private final Compression compression;
 
 	/** DEBUG only: just use the given file as the root for TRAP, source archive etc */
-	OdasaOutput(File outputRoot, Logger log) {
+	OdasaOutput(File outputRoot, Compression compression, Logger log) {
 		this.trapFolder = new File(outputRoot, "trap");
 		this.sourceArchiveFolder = new File(outputRoot, "src_archive");
 		this.trackClassOrigins = false;
 		this.log = log;
+		this.compression = compression;
 	}
 
-	public OdasaOutput(boolean trackClassOrigins, Logger log) {
+	public OdasaOutput(boolean trackClassOrigins, Compression compression, Logger log) {
 		String trapFolderVar = Env.systemEnv().getFirstNonEmpty("CODEQL_EXTRACTOR_JAVA_TRAP_DIR", Var.TRAP_FOLDER.name());
 		if (trapFolderVar == null) {
 			throw new ResourceError("CODEQL_EXTRACTOR_JAVA_TRAP_DIR was not set");
@@ -85,6 +89,7 @@ public OdasaOutput(boolean trackClassOrigins, Logger log) {
 		this.sourceArchiveFolder = new File(sourceArchiveVar);
 		this.trackClassOrigins = trackClassOrigins;
 		this.log = log;
+		this.compression = compression;
 	}
 
 	public File getTrapFolder() {
@@ -180,18 +185,18 @@ private File getTrapFileForJarFile(File jarFile) {
 			return null;
 		return FileUtil.appendAbsolutePath(
 				currentSpecFileEntry.getTrapFolder(),
-				JARS_DIR + "/" + PathTransformer.std().fileAsDatabaseString(jarFile) + ".trap.gz");
+				JARS_DIR + "/" + PathTransformer.std().fileAsDatabaseString(jarFile) + ".trap" + compression.getExtension());
 	}
 
 	private File getTrapFileForModule(String moduleName) {
 		return FileUtil.appendAbsolutePath(
 				currentSpecFileEntry.getTrapFolder(),
-				MODULES_DIR + "/" + moduleName + ".trap.gz");
+				MODULES_DIR + "/" + moduleName + ".trap" + compression.getExtension());
 	}
 
 	private File trapFileFor(File file) {
 		return FileUtil.appendAbsolutePath(currentSpecFileEntry.getTrapFolder(),
-				PathTransformer.std().fileAsDatabaseString(file) + ".trap.gz");
+				PathTransformer.std().fileAsDatabaseString(file) + ".trap" + compression.getExtension());
 	}
 
 	private File getTrapFileForDecl(IrElement sym, String signature) {
@@ -214,7 +219,7 @@ private String trapFilePathForDecl(IrElement sym, String signature) {
 					binaryName.replace('.', '/') +
 					signature +
 					".members" +
-					".trap.gz";
+					".trap" + compression.getExtension();
 		return result;
 	}
 
@@ -245,7 +250,7 @@ private TrapFileManager getMembersWriterForDecl(File trap, File trapFileBase, Tr
 		// don't need to rewrite it only to rename it
 		// again.
 		File trapFileDir = trap.getParentFile();
-		File trapOld = new File(trapFileDir, trap.getName().replace(".trap.gz", ".trap-old.gz"));
+		File trapOld = new File(trapFileDir, trap.getName().replace(".trap" + compression.getExtension(), ".trap-old" + compression.getExtension()));
 		if (trapOld.exists()) {
 			log.trace("Not rewriting trap file for " + trap.toString() + " as the trap-old exists");
 			return null;
@@ -272,7 +277,7 @@ private TrapFileManager getMembersWriterForDecl(File trap, File trapFileBase, Tr
 	}
 
 	private TrapFileManager trapWriter(File trapFile, IrElement sym, String signature) {
-		if (!trapFile.getName().endsWith(".trap.gz"))
+		if (!trapFile.getName().endsWith(".trap" + compression.getExtension()))
 			throw new CatastrophicError("OdasaOutput only supports writing to compressed trap files");
 		String relative = FileUtil.relativePath(trapFile, currentSpecFileEntry.getTrapFolder());
 		trapFile.getParentFile().mkdirs();
@@ -321,7 +326,7 @@ public void close() {
 			writeTrapDependencies(trapDependenciesForClass);
 		}
 		private void writeTrapDependencies(TrapDependencies trapDependencies) {
-			String dep = trapDependencies.trapFile().replace(".trap.gz", ".dep");
+			String dep = trapDependencies.trapFile().replace(".trap" + compression.getExtension(), ".dep");
 			trapDependencies.save(
 					currentSpecFileEntry.getTrapFolder().toPath().resolve(dep));
 		}
@@ -335,7 +340,7 @@ public void setHasError() {
 	 * Trap file locking.
 	 */
 
-	private final Pattern selectClassVersionComponents = Pattern.compile("(.*)#(-?[0-9]+)\\.(-?[0-9]+)-(-?[0-9]+)-(.*)\\.trap\\.gz");
+	private final Pattern selectClassVersionComponents = Pattern.compile("(.*)#(-?[0-9]+)\\.(-?[0-9]+)-(-?[0-9]+)-(.*)\\.trap.*");
 
 	/**
 	 * <b>CAUTION</b>: to avoid the potential for deadlock between multiple concurrent extractor processes,
@@ -412,12 +417,12 @@ private TrapLocker(IrElement decl, String signature, boolean fromSource) {
 					trapFileVersion = new TrapClassVersion(0, 0, 0, "kotlin");
 				else
 					trapFileVersion = TrapClassVersion.fromSymbol(sym, log);
-				String baseName = normalTrapFile.getName().replace(".trap.gz", "");
+				String baseName = normalTrapFile.getName().replace(".trap" + compression.getExtension(), "");
 				// 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.
 				// To avoid this, rather than using files named .../Foo*, we use .../Foo/Foo*.
 				trapFileBase = new File(new File(normalTrapFile.getParentFile(), baseName), baseName);
-				trapFile = new File(trapFileBase.getPath() + '#' + trapFileVersion.toString() + ".trap.gz");
+				trapFile = new File(trapFileBase.getPath() + '#' + trapFileVersion.toString() + ".trap" + compression.getExtension());
 			}
 		}
 		private TrapLocker(File jarFile) {
@@ -488,7 +493,7 @@ public int compare(Pair<File, TrapClassVersion> p1, Pair<File, TrapClassVersion>
 						for (Pair<File, TrapClassVersion> p: pairs) {
 							if (!latestVersion.equals(p.snd())) {
 								File f = p.fst();
-								File fOld = new File(f.getParentFile(), f.getName().replace(".trap.gz", ".trap-old.gz"));
+								File fOld = new File(f.getParentFile(), f.getName().replace(".trap" + compression.getExtension(), ".trap-old" + compression.getExtension()));
 								// We aren't interested in whether or not this succeeds;
 								// it may fail because a concurrent extractor has already
 								// renamed it.

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

@@ -14,7 +14,7 @@ import java.util.ArrayList
 import java.util.HashSet
 import java.util.zip.GZIPOutputStream
 
-class ExternalDeclExtractor(val logger: FileLogger, val invocationTrapFile: String, val sourceFilePath: String, val primitiveTypeMapping: PrimitiveTypeMapping, val pluginContext: IrPluginContext, val globalExtensionState: KotlinExtractorGlobalState, val diagnosticTrapWriter: DiagnosticTrapWriter) {
+class ExternalDeclExtractor(val logger: FileLogger, val compression: Compression, val invocationTrapFile: String, val sourceFilePath: String, val primitiveTypeMapping: PrimitiveTypeMapping, val pluginContext: IrPluginContext, val globalExtensionState: KotlinExtractorGlobalState, val diagnosticTrapWriter: DiagnosticTrapWriter) {
 
     val declBinaryNames = HashMap<IrDeclaration, String>()
     val externalDeclsDone = HashSet<Pair<String, String>>()
@@ -23,7 +23,7 @@ class ExternalDeclExtractor(val logger: FileLogger, val invocationTrapFile: Stri
     val propertySignature = ";property"
     val fieldSignature = ";field"
 
-    val output = OdasaOutput(false, logger).also {
+    val output = OdasaOutput(false, compression, logger).also {
         it.setCurrentSourceFile(File(sourceFilePath))
     }
 
@@ -65,7 +65,7 @@ class ExternalDeclExtractor(val logger: FileLogger, val invocationTrapFile: Stri
                     val trapFile = manager.file
                     val trapTmpFile = File.createTempFile("${trapFile.nameWithoutExtension}.", ".${trapFile.extension}.tmp", trapFile.parentFile)
                     try {
-                        GZIPOutputStream(trapTmpFile.outputStream()).bufferedWriter().use {
+                        compression.bufferedWriter(trapTmpFile).use {
                             extractorFn(it, signature, manager)
                         }
 

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

@@ -334,7 +334,7 @@ private fun doFile(
                 // Now elevate to a SourceFileTrapWriter, and populate the
                 // file information
                 val sftw = tw.makeSourceFileTrapWriter(srcFile, true)
-                val externalDeclExtractor = ExternalDeclExtractor(logger, invocationTrapFile, srcFilePath, primitiveTypeMapping, pluginContext, globalExtensionState, fileTrapWriter.getDiagnosticTrapWriter())
+                val externalDeclExtractor = ExternalDeclExtractor(logger, compression, invocationTrapFile, srcFilePath, primitiveTypeMapping, pluginContext, globalExtensionState, fileTrapWriter.getDiagnosticTrapWriter())
                 val linesOfCode = LinesOfCode(logger, sftw, srcFile)
                 val fileExtractor = KotlinFileExtractor(logger, sftw, linesOfCode, srcFilePath, null, externalDeclExtractor, primitiveTypeMapping, pluginContext, KotlinFileExtractor.DeclarationStack(), globalExtensionState)
 
@@ -362,7 +362,19 @@ private fun doFile(
     }
 }
 
-enum class Compression { NONE, GZIP }
+enum class Compression(val extension: String) {
+    NONE("") {
+        override fun bufferedWriter(file: File): BufferedWriter {
+            return file.bufferedWriter()
+        }
+    },
+    GZIP(".gz") {
+        override fun bufferedWriter(file: File): BufferedWriter {
+            return GZIPOutputStream(file.outputStream()).bufferedWriter()
+        }
+    };
+    abstract fun bufferedWriter(file: File): BufferedWriter
+}
 
 private fun getTrapFileWriter(compression: Compression, logger: FileLogger, trapFileName: String): TrapFileWriter {
     return when (compression) {

java/ql/integration-tests/all-platforms/kotlin/trap_compression/test.py

@@ -1,55 +1,66 @@
 from create_database_utils import *
 
-def check_extension(directory, expected_extension):
+def check_extensions(directory, counts):
     if platform.system() == 'Windows':
         # It's important that the path is a Unicode path on Windows, so
         # that the right system calls get used.
         directory = u'' + directory
         if not directory.startswith("\\\\?\\"):
             directory = "\\\\?\\" + os.path.abspath(directory)
 
-    if expected_extension == '.trap':
-        # We start TRAP files with a comment
-        expected_start = b'//'
-    elif expected_extension == '.trap.gz':
-        # The GZip magic numbers
-        expected_start = b'\x1f\x8b'
-    else:
-        raise Exception('Unknown expected extension ' + expected_extension)
-    count = check_extension_worker(directory, expected_extension, expected_start)
-    if count != 1:
-        raise Exception('Expected 1 relevant file, but found ' + str(count) + ' in ' + directory)
+    check_extensions_worker(counts, directory)
+    check_counts('non-compressed', counts.expected_none, counts.count_none)
+    check_counts('gzipped', counts.expected_gzip, counts.count_gzip)
 
-def check_extension_worker(directory, expected_extension, expected_start):
-    count = 0
+def check_counts(name, expected, count):
+    if expected == -1:
+        if count < 10:
+            raise Exception('Expected lots of ' + name + ' files, but got ' + str(count))
+    elif expected != count:
+        raise Exception('Expected ' + str(expected) + ' ' + name + ' files, but got ' + str(count))
+
+class Counts:
+    def __init__(self, expected_none, expected_gzip):
+        self.expected_none = expected_none
+        self.expected_gzip = expected_gzip
+        self.count_none = 0
+        self.count_gzip = 0
+
+def check_extensions_worker(counts, directory):
     for f in os.listdir(directory):
         x = os.path.join(directory, f)
         if os.path.isdir(x):
-            count += check_extension_worker(x, expected_extension, expected_start)
-        else:
-            if f.startswith('test.kt') and not f.endswith('.set'):
-                if f.endswith(expected_extension):
-                    with open(x, 'rb') as f_in:
-                        content = f_in.read()
-                    if content.startswith(expected_start):
-                        count += 1
-                    else:
-                        raise Exception('Unexpected start to content of ' + x)
-                else:
-                    raise Exception('Expected test.kt TRAP file to have extension ' + expected_extension + ', but found ' + x)
-    return count
+            check_extensions_worker(counts, x)
+        elif f.endswith('.trap'):
+            counts.count_none += 1
+            if not startsWith(x, b'//'): # We start TRAP files with a comment
+                raise Exception("TRAP file that doesn't start with a comment: " + f)
+        elif f.endswith('.trap.gz'):
+            counts.count_gzip += 1
+            if not startsWith(x, b'\x1f\x8b'): # The GZip magic numbers
+                raise Exception("GZipped TRAP file that doesn't start with GZip magic numbers: " + f)
+
+def startsWith(f, b):
+    with open(f, 'rb') as f_in:
+        content = f_in.read()
+    return content.startswith(b)
 
+# In the counts, we expect lots of files of the compression type chosen
+# (so expected count is -1), but the diagnostic TRAP files will always
+# be uncompressed (so count_none is always 1 or -1) and the
+# sourceLocationPrefix TRAP file is always gzipped (so count_gzip is
+# always 1 or -1).
 run_codeql_database_create(['kotlinc test.kt'], test_db="default-db", db=None, lang="java")
-check_extension('default-db/trap', '.trap.gz')
+check_extensions('default-db/trap', Counts(1, -1))
 os.environ["CODEQL_EXTRACTOR_JAVA_OPTION_TRAP_COMPRESSION"] = "nOnE"
 run_codeql_database_create(['kotlinc test.kt'], test_db="none-db", db=None, lang="java")
-check_extension('none-db/trap', '.trap')
+check_extensions('none-db/trap', Counts(-1, 1))
 os.environ["CODEQL_EXTRACTOR_JAVA_OPTION_TRAP_COMPRESSION"] = "gzip"
 run_codeql_database_create(['kotlinc test.kt'], test_db="gzip-db", db=None, lang="java")
-check_extension('gzip-db/trap', '.trap.gz')
+check_extensions('gzip-db/trap', Counts(1, -1))
 os.environ["CODEQL_EXTRACTOR_JAVA_OPTION_TRAP_COMPRESSION"] = "brotli"
 run_codeql_database_create(['kotlinc test.kt'], test_db="brotli-db", db=None, lang="java")
-check_extension('brotli-db/trap', '.trap.gz')
+check_extensions('brotli-db/trap', Counts(1, -1))
 os.environ["CODEQL_EXTRACTOR_JAVA_OPTION_TRAP_COMPRESSION"] = "invalidValue"
 run_codeql_database_create(['kotlinc test.kt'], test_db="invalid-db", db=None, lang="java")
-check_extension('invalid-db/trap', '.trap.gz')
+check_extensions('invalid-db/trap', Counts(1, -1))