Ruby, Rust: add zstd compression option

Author: Paolo Tranquilli

ruby/codeql-extractor.yml

@@ -27,7 +27,7 @@ options:
         title: Controls compression for the TRAP files written by the extractor.
         description: >
           This option is only intended for use in debugging the extractor. Accepted
-          values are 'gzip' (the default, to write gzip-compressed TRAP) and 'none'
-          (to write uncompressed TRAP).
+          values are 'gzip' (the default, to write gzip-compressed TRAP) 'zstd' (to
+          write Zstandard-compressed TRAP) and 'none' (to write uncompressed TRAP).
         type: string
-        pattern: "^(none|gzip)$"
+        pattern: "^(none|gzip|zstd)$"

rust/codeql-extractor.yml

@@ -23,10 +23,10 @@ options:
         title: Controls compression for the TRAP files written by the extractor.
         description: >
           This option is only intended for use in debugging the extractor. Accepted
-          values are 'gzip' (to write gzip-compressed TRAP) and 'none'
-          (currently the default, to write uncompressed TRAP).
+          values are 'gzip' (the default, to write gzip-compressed TRAP) 'zstd' (to
+          write Zstandard-compressed TRAP) and 'none' (to write uncompressed TRAP).
         type: string
-        pattern: "^(none|gzip)$"
+        pattern: "^(none|gzip|zstd)$"
   cargo_target_dir:
     title: Directory to use for cargo output files.
     description: >

rust/extractor/src/config.rs

@@ -29,13 +29,15 @@ pub enum Compression {
     #[default] // TODO make gzip default
     None,
     Gzip,
+    Zstd,
 }
 
 impl From<Compression> for trap::Compression {
     fn from(val: Compression) -> Self {
         match val {
             Compression::None => Self::None,
             Compression::Gzip => Self::Gzip,
+            Compression::Zstd => Self::Zstd,
         }
     }
 }

rust/ql/integration-tests/hello-project/test_project.py

@@ -20,6 +20,7 @@ def test_do_not_print_env(codeql, rust, rust_edition, cargo, check_env_not_dumpe
 @pytest.mark.ql_test("steps.ql", expected=".cargo.expected")
 @pytest.mark.parametrize(("rust_edition", "compression", "suffix"), [
     pytest.param(2024, "gzip", ".gz", id="gzip"),
+    pytest.param(2024, "zstd", ".zst", id="zstd"),
 ])
 def test_compression(codeql, rust, rust_edition, compression, suffix, cargo, rust_check_diagnostics, cwd):
     codeql.database.create(cleanup=False, _env={

shared/tree-sitter-extractor/src/trap.rs

@@ -96,18 +96,25 @@ impl Writer {
                 self.write_trap_entries(&mut trap_file)
             }
             Compression::Gzip => {
-                let trap_file = GzEncoder::new(trap_file, flate2::Compression::fast());
+                let trap_file = GzEncoder::new(trap_file, Compression::GZIP_LEVEL);
                 let mut trap_file = BufWriter::new(trap_file);
                 self.write_trap_entries(&mut trap_file)
             }
+            Compression::Zstd => {
+                let trap_file = zstd::stream::Encoder::new(trap_file, Compression::ZSTD_LEVEL)?;
+                let mut trap_file = BufWriter::new(trap_file);
+                self.write_trap_entries(&mut trap_file)?;
+                trap_file.into_inner()?.finish()?;
+                Ok(())
+            }
         }
     }
 
     fn write_trap_entries<W: Write>(&self, file: &mut W) -> std::io::Result<()> {
         for trap_entry in &self.trap_output {
             writeln!(file, "{}", trap_entry)?;
         }
-        std::io::Result::Ok(())
+        Ok(())
     }
 }
 
@@ -280,9 +287,13 @@ fn limit_string(string: &str, max_size: usize) -> &str {
 pub enum Compression {
     None,
     Gzip,
+    Zstd,
 }
 
 impl Compression {
+    pub const ZSTD_LEVEL: i32 = 2;
+    pub const GZIP_LEVEL: flate2::Compression = flate2::Compression::fast();
+
     pub fn from_env(var_name: &str) -> Result<Compression, String> {
         match std::env::var(var_name) {
             Ok(method) => match Compression::from_string(&method) {
@@ -298,6 +309,7 @@ impl Compression {
         match s.to_lowercase().as_ref() {
             "none" => Some(Compression::None),
             "gzip" => Some(Compression::Gzip),
+            "zstd" => Some(Compression::Zstd),
             _ => None,
         }
     }
@@ -306,6 +318,7 @@ impl Compression {
         match self {
             Compression::None => "trap",
             Compression::Gzip => "trap.gz",
+            Compression::Zstd => "trap.zst",
         }
     }
 }