fix import

This reverts commit 388c9ffb74.

ql/extractor/src/trap.rs

@@ -61,15 +61,12 @@ impl Writer {
 
     pub fn write_to_file(&self, path: &Path, compression: Compression) -> std::io::Result<()> {
         let trap_file = std::fs::File::create(path)?;
+        let mut trap_file = BufWriter::new(trap_file);
         match compression {
-            Compression::None => {
+            Compression::None => self.write_trap_entries(&mut trap_file),
-                let mut trap_file = BufWriter::new(trap_file);
-                self.write_trap_entries(&mut trap_file)
-            }
             Compression::Gzip => {
-                let trap_file = GzEncoder::new(trap_file, flate2::Compression::fast());
+                let mut compressed_writer = GzEncoder::new(trap_file, flate2::Compression::fast());
-                let mut trap_file = BufWriter::new(trap_file);
+                self.write_trap_entries(&mut compressed_writer)
-                self.write_trap_entries(&mut trap_file)
             }
         }
     }

ruby/extractor/src/extractor.rs

@@ -1,112 +1,161 @@
-use crate::trap;
 use node_types::{EntryKind, Field, NodeTypeMap, Storage, TypeName};
+use std::borrow::Cow;
 use std::collections::BTreeMap as Map;
 use std::collections::BTreeSet as Set;
 use std::fmt;
+use std::io::Write;
 use std::path::Path;
 
 use tracing::{error, info, span, Level};
 use tree_sitter::{Language, Node, Parser, Range, Tree};
 
-pub fn populate_file(writer: &mut trap::Writer, absolute_path: &Path) -> trap::Label {
+pub struct TrapWriter {
-    let (file_label, fresh) =
+    /// The accumulated trap entries
-        writer.global_id(&trap::full_id_for_file(&normalize_path(absolute_path)));
+    trap_output: Vec<TrapEntry>,
-    if fresh {
+    /// A counter for generating fresh labels
-        writer.add_tuple(
+    counter: u32,
-            "files",
+    /// cache of global keys
-            vec![
+    global_keys: std::collections::HashMap<String, Label>,
-                trap::Arg::Label(file_label),
+}
-                trap::Arg::String(normalize_path(absolute_path)),
+
-            ],
+pub fn new_trap_writer() -> TrapWriter {
-        );
+    TrapWriter {
-        populate_parent_folders(writer, file_label, absolute_path.parent());
+        counter: 0,
+        trap_output: Vec::new(),
+        global_keys: std::collections::HashMap::new(),
     }
-    file_label
 }
 
-fn populate_empty_file(writer: &mut trap::Writer) -> trap::Label {
+impl TrapWriter {
-    let (file_label, fresh) = writer.global_id("empty;sourcefile");
+    ///  Gets a label that will hold the unique ID of the passed string at import time.
-    if fresh {
+    ///  This can be used for incrementally importable TRAP files -- use globally unique
-        writer.add_tuple(
+    ///  strings to compute a unique ID for table tuples.
-            "files",
+    ///
-            vec![
+    ///  Note: You probably want to make sure that the key strings that you use are disjoint
-                trap::Arg::Label(file_label),
+    ///  for disjoint column types; the standard way of doing this is to prefix (or append)
-                trap::Arg::String("".to_string()),
+    ///  the column type name to the ID. Thus, you might identify methods in Java by the
-            ],
+    ///  full ID "methods_com.method.package.DeclaringClass.method(argumentList)".
-        );
+
+    fn fresh_id(&mut self) -> Label {
+        let label = Label(self.counter);
+        self.counter += 1;
+        self.trap_output.push(TrapEntry::FreshId(label));
+        label
     }
-    file_label
-}
 
-pub fn populate_empty_location(writer: &mut trap::Writer) {
+    fn global_id(&mut self, key: &str) -> (Label, bool) {
-    let file_label = populate_empty_file(writer);
+        if let Some(label) = self.global_keys.get(key) {
-    location(writer, file_label, 0, 0, 0, 0);
+            return (*label, false);
-}
+        }
+        let label = Label(self.counter);
+        self.counter += 1;
+        self.global_keys.insert(key.to_owned(), label);
+        self.trap_output
+            .push(TrapEntry::MapLabelToKey(label, key.to_owned()));
+        (label, true)
+    }
 
-pub fn populate_parent_folders(
+    fn add_tuple(&mut self, table_name: &str, args: Vec<Arg>) {
-    writer: &mut trap::Writer,
+        self.trap_output
-    child_label: trap::Label,
+            .push(TrapEntry::GenericTuple(table_name.to_owned(), args))
-    path: Option<&Path>,
+    }
-) {
+
-    let mut path = path;
+    fn populate_file(&mut self, absolute_path: &Path) -> Label {
-    let mut child_label = child_label;
+        let (file_label, fresh) = self.global_id(&full_id_for_file(absolute_path));
-    loop {
+        if fresh {
-        match path {
+            self.add_tuple(
-            None => break,
+                "files",
-            Some(folder) => {
+                vec![
-                let (folder_label, fresh) =
+                    Arg::Label(file_label),
-                    writer.global_id(&trap::full_id_for_folder(&normalize_path(folder)));
+                    Arg::String(normalize_path(absolute_path)),
-                writer.add_tuple(
+                ],
-                    "containerparent",
+            );
-                    vec![
+            self.populate_parent_folders(file_label, absolute_path.parent());
-                        trap::Arg::Label(folder_label),
+        }
-                        trap::Arg::Label(child_label),
+        file_label
-                    ],
+    }
-                );
+
-                if fresh {
+    fn populate_empty_file(&mut self) -> Label {
-                    writer.add_tuple(
+        let (file_label, fresh) = self.global_id("empty;sourcefile");
-                        "folders",
+        if fresh {
-                        vec![
+            self.add_tuple(
-                            trap::Arg::Label(folder_label),
+                "files",
-                            trap::Arg::String(normalize_path(folder)),
+                vec![Arg::Label(file_label), Arg::String("".to_string())],
-                        ],
+            );
+        }
+        file_label
+    }
+
+    pub fn populate_empty_location(&mut self) {
+        let file_label = self.populate_empty_file();
+        self.location(file_label, 0, 0, 0, 0);
+    }
+
+    fn populate_parent_folders(&mut self, child_label: Label, path: Option<&Path>) {
+        let mut path = path;
+        let mut child_label = child_label;
+        loop {
+            match path {
+                None => break,
+                Some(folder) => {
+                    let (folder_label, fresh) = self.global_id(&full_id_for_folder(folder));
+                    self.add_tuple(
+                        "containerparent",
+                        vec![Arg::Label(folder_label), Arg::Label(child_label)],
                     );
-                    path = folder.parent();
+                    if fresh {
-                    child_label = folder_label;
+                        self.add_tuple(
-                } else {
+                            "folders",
-                    break;
+                            vec![
+                                Arg::Label(folder_label),
+                                Arg::String(normalize_path(folder)),
+                            ],
+                        );
+                        path = folder.parent();
+                        child_label = folder_label;
+                    } else {
+                        break;
+                    }
                 }
             }
         }
     }
-}
 
-fn location(
+    fn location(
-    writer: &mut trap::Writer,
+        &mut self,
-    file_label: trap::Label,
+        file_label: Label,
-    start_line: usize,
+        start_line: usize,
-    start_column: usize,
+        start_column: usize,
-    end_line: usize,
+        end_line: usize,
-    end_column: usize,
+        end_column: usize,
-) -> trap::Label {
+    ) -> Label {
-    let (loc_label, fresh) = writer.global_id(&format!(
+        let (loc_label, fresh) = self.global_id(&format!(
-        "loc,{{{}}},{},{},{},{}",
+            "loc,{{{}}},{},{},{},{}",
-        file_label, start_line, start_column, end_line, end_column
+            file_label, start_line, start_column, end_line, end_column
-    ));
+        ));
-    if fresh {
+        if fresh {
-        writer.add_tuple(
+            self.add_tuple(
-            "locations_default",
+                "locations_default",
-            vec![
+                vec![
-                trap::Arg::Label(loc_label),
+                    Arg::Label(loc_label),
-                trap::Arg::Label(file_label),
+                    Arg::Label(file_label),
-                trap::Arg::Int(start_line),
+                    Arg::Int(start_line),
-                trap::Arg::Int(start_column),
+                    Arg::Int(start_column),
-                trap::Arg::Int(end_line),
+                    Arg::Int(end_line),
-                trap::Arg::Int(end_column),
+                    Arg::Int(end_column),
-            ],
+                ],
-        );
+            );
+        }
+        loc_label
+    }
+
+    fn comment(&mut self, text: String) {
+        self.trap_output.push(TrapEntry::Comment(text));
+    }
+
+    pub fn output(self, writer: &mut dyn Write) -> std::io::Result<()> {
+        write!(writer, "{}", Program(self.trap_output))
     }
-    loc_label
 }
 
 /// Extracts the source file at `path`, which is assumed to be canonicalized.
@@ -114,43 +163,71 @@ pub fn extract(
     language: Language,
     language_prefix: &str,
     schema: &NodeTypeMap,
-    trap_writer: &mut trap::Writer,
+    trap_writer: &mut TrapWriter,
     path: &Path,
     source: &[u8],
     ranges: &[Range],
 ) -> std::io::Result<()> {
-    let path_str = format!("{}", path.display());
     let span = span!(
         Level::TRACE,
         "extract",
-        file = %path_str
+        file = %path.display()
     );
 
     let _enter = span.enter();
 
-    info!("extracting: {}", path_str);
+    info!("extracting: {}", path.display());
 
     let mut parser = Parser::new();
     parser.set_language(language).unwrap();
     parser.set_included_ranges(ranges).unwrap();
     let tree = parser.parse(&source, None).expect("Failed to parse file");
-    trap_writer.comment(format!("Auto-generated TRAP file for {}", path_str));
+    trap_writer.comment(format!("Auto-generated TRAP file for {}", path.display()));
-    let file_label = populate_file(trap_writer, path);
+    let file_label = &trap_writer.populate_file(path);
-    let mut visitor = Visitor::new(
+    let mut visitor = Visitor {
         source,
         trap_writer,
         // TODO: should we handle path strings that are not valid UTF8 better?
-        &path_str,
+        path: format!("{}", path.display()),
-        file_label,
+        file_label: *file_label,
+        toplevel_child_counter: 0,
+        stack: Vec::new(),
         language_prefix,
         schema,
-    );
+    };
     traverse(&tree, &mut visitor);
 
     parser.reset();
     Ok(())
 }
 
+/// Escapes a string for use in a TRAP key, by replacing special characters with
+/// HTML entities.
+fn escape_key<'a, S: Into<Cow<'a, str>>>(key: S) -> Cow<'a, str> {
+    fn needs_escaping(c: char) -> bool {
+        matches!(c, '&' | '{' | '}' | '"' | '@' | '#')
+    }
+
+    let key = key.into();
+    if key.contains(needs_escaping) {
+        let mut escaped = String::with_capacity(2 * key.len());
+        for c in key.chars() {
+            match c {
+                '&' => escaped.push_str("&amp;"),
+                '{' => escaped.push_str("&lbrace;"),
+                '}' => escaped.push_str("&rbrace;"),
+                '"' => escaped.push_str("&quot;"),
+                '@' => escaped.push_str("&commat;"),
+                '#' => escaped.push_str("&num;"),
+                _ => escaped.push(c),
+            }
+        }
+        Cow::Owned(escaped)
+    } else {
+        key
+    }
+}
+
 /// Normalizes the path according the common CodeQL specification. Assumes that
 /// `path` has already been canonicalized using `std::fs::canonicalize`.
 fn normalize_path(path: &Path) -> String {
@@ -190,28 +267,34 @@ fn normalize_path(path: &Path) -> String {
     }
 }
 
+fn full_id_for_file(path: &Path) -> String {
+    format!("{};sourcefile", escape_key(&normalize_path(path)))
+}
+
+fn full_id_for_folder(path: &Path) -> String {
+    format!("{};folder", escape_key(&normalize_path(path)))
+}
+
 struct ChildNode {
     field_name: Option<&'static str>,
-    label: trap::Label,
+    label: Label,
     type_name: TypeName,
 }
 
 struct Visitor<'a> {
     /// The file path of the source code (as string)
-    path: &'a str,
+    path: String,
     /// The label to use whenever we need to refer to the `@file` entity of this
     /// source file.
-    file_label: trap::Label,
+    file_label: Label,
     /// The source code as a UTF-8 byte array
     source: &'a [u8],
-    /// A trap::Writer to accumulate trap entries
+    /// A TrapWriter to accumulate trap entries
-    trap_writer: &'a mut trap::Writer,
+    trap_writer: &'a mut TrapWriter,
     /// A counter for top-level child nodes
     toplevel_child_counter: usize,
-    /// Language-specific name of the AST info table
+    /// Language prefix
-    ast_node_info_table_name: String,
+    language_prefix: &'a str,
-    /// Language-specific name of the tokeninfo table
-    tokeninfo_table_name: String,
     /// A lookup table from type name to node types
     schema: &'a NodeTypeMap,
     /// A stack for gathering information from child nodes. Whenever a node is
@@ -220,48 +303,27 @@ struct Visitor<'a> {
     /// node the list containing the child data is popped from the stack and
     /// matched against the dbscheme for the node. If the expectations are met
     /// the corresponding row definitions are added to the trap_output.
-    stack: Vec<(trap::Label, usize, Vec<ChildNode>)>,
+    stack: Vec<(Label, usize, Vec<ChildNode>)>,
 }
 
-impl<'a> Visitor<'a> {
+impl Visitor<'_> {
-    fn new(
-        source: &'a [u8],
-        trap_writer: &'a mut trap::Writer,
-        path: &'a str,
-        file_label: trap::Label,
-        language_prefix: &str,
-        schema: &'a NodeTypeMap,
-    ) -> Visitor<'a> {
-        Visitor {
-            path,
-            file_label,
-            source,
-            trap_writer,
-            toplevel_child_counter: 0,
-            ast_node_info_table_name: format!("{}_ast_node_info", language_prefix),
-            tokeninfo_table_name: format!("{}_tokeninfo", language_prefix),
-            schema,
-            stack: Vec::new(),
-        }
-    }
-
     fn record_parse_error(
         &mut self,
         error_message: String,
         full_error_message: String,
-        loc: trap::Label,
+        loc: Label,
     ) {
         error!("{}", full_error_message);
         let id = self.trap_writer.fresh_id();
         self.trap_writer.add_tuple(
             "diagnostics",
             vec![
-                trap::Arg::Label(id),
+                Arg::Label(id),
-                trap::Arg::Int(40), // severity 40 = error
+                Arg::Int(40), // severity 40 = error
-                trap::Arg::String("parse_error".to_string()),
+                Arg::String("parse_error".to_string()),
-                trap::Arg::String(error_message),
+                Arg::String(error_message),
-                trap::Arg::String(full_error_message),
+                Arg::String(full_error_message),
-                trap::Arg::Label(loc),
+                Arg::Label(loc),
             ],
         );
     }
@@ -273,8 +335,7 @@ impl<'a> Visitor<'a> {
         node: Node,
     ) {
         let (start_line, start_column, end_line, end_column) = location_for(self.source, node);
-        let loc = location(
+        let loc = self.trap_writer.location(
-            self.trap_writer,
             self.file_label,
             start_line,
             start_column,
@@ -313,8 +374,7 @@ impl<'a> Visitor<'a> {
         }
         let (id, _, child_nodes) = self.stack.pop().expect("Vistor: empty stack");
         let (start_line, start_column, end_line, end_column) = location_for(self.source, node);
-        let loc = location(
+        let loc = self.trap_writer.location(
-            self.trap_writer,
             self.file_label,
             start_line,
             start_column,
@@ -342,19 +402,19 @@ impl<'a> Visitor<'a> {
         match &table.kind {
             EntryKind::Token { kind_id, .. } => {
                 self.trap_writer.add_tuple(
-                    &self.ast_node_info_table_name,
+                    &format!("{}_ast_node_info", self.language_prefix),
                     vec![
-                        trap::Arg::Label(id),
+                        Arg::Label(id),
-                        trap::Arg::Label(parent_id),
+                        Arg::Label(parent_id),
-                        trap::Arg::Int(parent_index),
+                        Arg::Int(parent_index),
-                        trap::Arg::Label(loc),
+                        Arg::Label(loc),
                     ],
                 );
                 self.trap_writer.add_tuple(
-                    &self.tokeninfo_table_name,
+                    &format!("{}_tokeninfo", self.language_prefix),
                     vec![
-                        trap::Arg::Label(id),
+                        Arg::Label(id),
-                        trap::Arg::Int(*kind_id),
+                        Arg::Int(*kind_id),
                         sliced_source_arg(self.source, node),
                     ],
                 );
@@ -365,15 +425,15 @@ impl<'a> Visitor<'a> {
             } => {
                 if let Some(args) = self.complex_node(&node, fields, &child_nodes, id) {
                     self.trap_writer.add_tuple(
-                        &self.ast_node_info_table_name,
+                        &format!("{}_ast_node_info", self.language_prefix),
                         vec![
-                            trap::Arg::Label(id),
+                            Arg::Label(id),
-                            trap::Arg::Label(parent_id),
+                            Arg::Label(parent_id),
-                            trap::Arg::Int(parent_index),
+                            Arg::Int(parent_index),
-                            trap::Arg::Label(loc),
+                            Arg::Label(loc),
                         ],
                     );
-                    let mut all_args = vec![trap::Arg::Label(id)];
+                    let mut all_args = vec![Arg::Label(id)];
                     all_args.extend(args);
                     self.trap_writer.add_tuple(table_name, all_args);
                 }
@@ -412,9 +472,9 @@ impl<'a> Visitor<'a> {
         node: &Node,
         fields: &[Field],
         child_nodes: &[ChildNode],
-        parent_id: trap::Label,
+        parent_id: Label,
-    ) -> Option<Vec<trap::Arg>> {
+    ) -> Option<Vec<Arg>> {
-        let mut map: Map<&Option<String>, (&Field, Vec<trap::Arg>)> = Map::new();
+        let mut map: Map<&Option<String>, (&Field, Vec<Arg>)> = Map::new();
         for field in fields {
             map.insert(&field.name, (field, Vec::new()));
         }
@@ -428,9 +488,9 @@ impl<'a> Visitor<'a> {
                     {
                         // We can safely unwrap because type_matches checks the key is in the map.
                         let (int_value, _) = int_mapping.get(&child_node.type_name.kind).unwrap();
-                        values.push(trap::Arg::Int(*int_value));
+                        values.push(Arg::Int(*int_value));
                     } else {
-                        values.push(trap::Arg::Label(child_node.label));
+                        values.push(Arg::Label(child_node.label));
                     }
                 } else if field.name.is_some() {
                     let error_message = format!(
@@ -509,9 +569,9 @@ impl<'a> Visitor<'a> {
                             );
                             break;
                         }
-                        let mut args = vec![trap::Arg::Label(parent_id)];
+                        let mut args = vec![Arg::Label(parent_id)];
                         if *has_index {
-                            args.push(trap::Arg::Int(index))
+                            args.push(Arg::Int(index))
                         }
                         args.push(child_value.clone());
                         self.trap_writer.add_tuple(table_name, args);
@@ -565,9 +625,9 @@ impl<'a> Visitor<'a> {
 }
 
 // Emit a slice of a source file as an Arg.
-fn sliced_source_arg(source: &[u8], n: Node) -> trap::Arg {
+fn sliced_source_arg(source: &[u8], n: Node) -> Arg {
     let range = n.byte_range();
-    trap::Arg::String(String::from_utf8_lossy(&source[range.start..range.end]).into_owned())
+    Arg::String(String::from_utf8_lossy(&source[range.start..range.end]).into_owned())
 }
 
 // Emit a pair of `TrapEntry`s for the provided node, appropriately calibrated.
@@ -639,6 +699,59 @@ fn traverse(tree: &Tree, visitor: &mut Visitor) {
     }
 }
 
+pub struct Program(Vec<TrapEntry>);
+
+impl fmt::Display for Program {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let mut text = String::new();
+        for trap_entry in &self.0 {
+            text.push_str(&format!("{}\n", trap_entry));
+        }
+        write!(f, "{}", text)
+    }
+}
+
+enum TrapEntry {
+    /// Maps the label to a fresh id, e.g. `#123=*`.
+    FreshId(Label),
+    /// Maps the label to a key, e.g. `#7=@"foo"`.
+    MapLabelToKey(Label, String),
+    /// foo_bar(arg*)
+    GenericTuple(String, Vec<Arg>),
+    Comment(String),
+}
+impl fmt::Display for TrapEntry {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            TrapEntry::FreshId(label) => write!(f, "{}=*", label),
+            TrapEntry::MapLabelToKey(label, key) => {
+                write!(f, "{}=@\"{}\"", label, key.replace("\"", "\"\""))
+            }
+            TrapEntry::GenericTuple(name, args) => {
+                write!(f, "{}(", name)?;
+                for (index, arg) in args.iter().enumerate() {
+                    if index > 0 {
+                        write!(f, ",")?;
+                    }
+                    write!(f, "{}", arg)?;
+                }
+                write!(f, ")")
+            }
+            TrapEntry::Comment(line) => write!(f, "// {}", line),
+        }
+    }
+}
+
+#[derive(Debug, Copy, Clone)]
+// Identifiers of the form #0, #1...
+struct Label(u32);
+
+impl fmt::Display for Label {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "#{:x}", self.0)
+    }
+}
+
 // Numeric indices.
 #[derive(Debug, Copy, Clone)]
 struct Index(usize);
@@ -648,3 +761,69 @@ impl fmt::Display for Index {
         write!(f, "{}", self.0)
     }
 }
+
+// Some untyped argument to a TrapEntry.
+#[derive(Debug, Clone)]
+enum Arg {
+    Label(Label),
+    Int(usize),
+    String(String),
+}
+
+const MAX_STRLEN: usize = 1048576;
+
+impl fmt::Display for Arg {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Arg::Label(x) => write!(f, "{}", x),
+            Arg::Int(x) => write!(f, "{}", x),
+            Arg::String(x) => write!(
+                f,
+                "\"{}\"",
+                limit_string(x, MAX_STRLEN).replace("\"", "\"\"")
+            ),
+        }
+    }
+}
+
+/// Limit the length (in bytes) of a string. If the string's length in bytes is
+/// less than or equal to the limit then the entire string is returned. Otherwise
+/// the string is sliced at the provided limit. If there is a multi-byte character
+/// at the limit then the returned slice will be slightly shorter than the limit to
+/// avoid splitting that multi-byte character.
+fn limit_string(string: &str, max_size: usize) -> &str {
+    if string.len() <= max_size {
+        return string;
+    }
+    let p = string.as_bytes();
+    let mut index = max_size;
+    // We want to clip the string at [max_size]; however, the character at that position
+    // may span several bytes. We need to find the first byte of the character. In UTF-8
+    // encoded data any byte that matches the bit pattern 10XXXXXX is not a start byte.
+    // Therefore we decrement the index as long as there are bytes matching this pattern.
+    // This ensures we cut the string at the border between one character and another.
+    while index > 0 && (p[index] & 0b11000000) == 0b10000000 {
+        index -= 1;
+    }
+    &string[0..index]
+}
+
+#[test]
+fn limit_string_test() {
+    assert_eq!("hello", limit_string(&"hello world".to_owned(), 5));
+    assert_eq!("hi ☹", limit_string(&"hi ☹☹".to_owned(), 6));
+    assert_eq!("hi ", limit_string(&"hi ☹☹".to_owned(), 5));
+}
+
+#[test]
+fn escape_key_test() {
+    assert_eq!("foo!", escape_key("foo!"));
+    assert_eq!("foo&lbrace;&rbrace;", escape_key("foo{}"));
+    assert_eq!("&lbrace;&rbrace;", escape_key("{}"));
+    assert_eq!("", escape_key(""));
+    assert_eq!("/path/to/foo.rb", escape_key("/path/to/foo.rb"));
+    assert_eq!(
+        "/path/to/foo&amp;&lbrace;&rbrace;&quot;&commat;&num;.rb",
+        escape_key("/path/to/foo&{}\"@#.rb")
+    );
+}

ruby/extractor/src/main.rs

@@ -1,19 +1,54 @@
 mod extractor;
-mod trap;
 
 #[macro_use]
 extern crate lazy_static;
 extern crate num_cpus;
 
 use clap::arg;
-use encoding::{self};
+use flate2::write::GzEncoder;
 use rayon::prelude::*;
 use std::borrow::Cow;
 use std::fs;
-use std::io::BufRead;
+use std::io::{BufRead, BufWriter};
 use std::path::{Path, PathBuf};
 use tree_sitter::{Language, Parser, Range};
 
+enum TrapCompression {
+    None,
+    Gzip,
+}
+
+impl TrapCompression {
+    fn from_env() -> TrapCompression {
+        match std::env::var("CODEQL_RUBY_TRAP_COMPRESSION") {
+            Ok(method) => match TrapCompression::from_string(&method) {
+                Some(c) => c,
+                None => {
+                    tracing::error!("Unknown compression method '{}'; using gzip.", &method);
+                    TrapCompression::Gzip
+                }
+            },
+            // Default compression method if the env var isn't set:
+            Err(_) => TrapCompression::Gzip,
+        }
+    }
+
+    fn from_string(s: &str) -> Option<TrapCompression> {
+        match s.to_lowercase().as_ref() {
+            "none" => Some(TrapCompression::None),
+            "gzip" => Some(TrapCompression::Gzip),
+            _ => None,
+        }
+    }
+
+    fn extension(&self) -> &str {
+        match self {
+            TrapCompression::None => "trap",
+            TrapCompression::Gzip => "trap.gz",
+        }
+    }
+}
+
 /**
  * Gets the number of threads the extractor should use, by reading the
  * CODEQL_THREADS environment variable and using it as described in the
@@ -101,7 +136,7 @@ fn main() -> std::io::Result<()> {
         .value_of("output-dir")
         .expect("missing --output-dir");
     let trap_dir = PathBuf::from(trap_dir);
-    let trap_compression = trap::Compression::from_env("CODEQL_RUBY_TRAP_COMPRESSION");
+    let trap_compression = TrapCompression::from_env();
 
     let file_list = matches.value_of("file-list").expect("missing --file-list");
     let file_list = fs::File::open(file_list)?;
@@ -125,7 +160,7 @@ fn main() -> std::io::Result<()> {
             let mut source = std::fs::read(&path)?;
             let mut needs_conversion = false;
             let code_ranges;
-            let mut trap_writer = trap::Writer::new();
+            let mut trap_writer = extractor::new_trap_writer();
             if path.extension().map_or(false, |x| x == "erb") {
                 tracing::info!("scanning: {}", path.display());
                 extractor::extract(
@@ -206,25 +241,33 @@ fn main() -> std::io::Result<()> {
             } else {
                 std::fs::copy(&path, &src_archive_file)?;
             }
-            write_trap(&trap_dir, path, &trap_writer, trap_compression)
+            write_trap(&trap_dir, path, trap_writer, &trap_compression)
         })
         .expect("failed to extract files");
 
     let path = PathBuf::from("extras");
-    let mut trap_writer = trap::Writer::new();
+    let mut trap_writer = extractor::new_trap_writer();
-    extractor::populate_empty_location(&mut trap_writer);
+    trap_writer.populate_empty_location();
-    write_trap(&trap_dir, path, &trap_writer, trap_compression)
+    write_trap(&trap_dir, path, trap_writer, &trap_compression)
 }
 
 fn write_trap(
     trap_dir: &Path,
     path: PathBuf,
-    trap_writer: &trap::Writer,
+    trap_writer: extractor::TrapWriter,
-    trap_compression: trap::Compression,
+    trap_compression: &TrapCompression,
 ) -> std::io::Result<()> {
     let trap_file = path_for(trap_dir, &path, trap_compression.extension());
     std::fs::create_dir_all(&trap_file.parent().unwrap())?;
-    trap_writer.write_to_file(&trap_file, trap_compression)
+    let trap_file = std::fs::File::create(&trap_file)?;
+    let mut trap_file = BufWriter::new(trap_file);
+    match trap_compression {
+        TrapCompression::None => trap_writer.output(&mut trap_file),
+        TrapCompression::Gzip => {
+            let mut compressed_writer = GzEncoder::new(trap_file, flate2::Compression::fast());
+            trap_writer.output(&mut compressed_writer)
+        }
+    }
 }
 
 fn scan_erb(

ruby/extractor/src/trap.rs

@@ -1,275 +0,0 @@
-use std::borrow::Cow;
-use std::fmt;
-use std::io::{BufWriter, Write};
-use std::path::Path;
-
-use flate2::write::GzEncoder;
-
-pub struct Writer {
-    /// The accumulated trap entries
-    trap_output: Vec<Entry>,
-    /// A counter for generating fresh labels
-    counter: u32,
-    /// cache of global keys
-    global_keys: std::collections::HashMap<String, Label>,
-}
-
-impl Writer {
-    pub fn new() -> Writer {
-        Writer {
-            counter: 0,
-            trap_output: Vec::new(),
-            global_keys: std::collections::HashMap::new(),
-        }
-    }
-
-    pub fn fresh_id(&mut self) -> Label {
-        let label = Label(self.counter);
-        self.counter += 1;
-        self.trap_output.push(Entry::FreshId(label));
-        label
-    }
-
-    ///  Gets a label that will hold the unique ID of the passed string at import time.
-    ///  This can be used for incrementally importable TRAP files -- use globally unique
-    ///  strings to compute a unique ID for table tuples.
-    ///
-    ///  Note: You probably want to make sure that the key strings that you use are disjoint
-    ///  for disjoint column types; the standard way of doing this is to prefix (or append)
-    ///  the column type name to the ID. Thus, you might identify methods in Java by the
-    ///  full ID "methods_com.method.package.DeclaringClass.method(argumentList)".
-    pub fn global_id(&mut self, key: &str) -> (Label, bool) {
-        if let Some(label) = self.global_keys.get(key) {
-            return (*label, false);
-        }
-        let label = Label(self.counter);
-        self.counter += 1;
-        self.global_keys.insert(key.to_owned(), label);
-        self.trap_output
-            .push(Entry::MapLabelToKey(label, key.to_owned()));
-        (label, true)
-    }
-
-    pub fn add_tuple(&mut self, table_name: &str, args: Vec<Arg>) {
-        self.trap_output
-            .push(Entry::GenericTuple(table_name.to_owned(), args))
-    }
-
-    pub fn comment(&mut self, text: String) {
-        self.trap_output.push(Entry::Comment(text));
-    }
-
-    pub fn write_to_file(&self, path: &Path, compression: Compression) -> std::io::Result<()> {
-        let trap_file = std::fs::File::create(path)?;
-        match compression {
-            Compression::None => {
-                let mut trap_file = BufWriter::new(trap_file);
-                self.write_trap_entries(&mut trap_file)
-            }
-            Compression::Gzip => {
-                let trap_file = GzEncoder::new(trap_file, flate2::Compression::fast());
-                let mut trap_file = BufWriter::new(trap_file);
-                self.write_trap_entries(&mut trap_file)
-            }
-        }
-    }
-
-    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(())
-    }
-}
-
-pub enum Entry {
-    /// Maps the label to a fresh id, e.g. `#123=*`.
-    FreshId(Label),
-    /// Maps the label to a key, e.g. `#7=@"foo"`.
-    MapLabelToKey(Label, String),
-    /// foo_bar(arg*)
-    GenericTuple(String, Vec<Arg>),
-    Comment(String),
-}
-
-impl fmt::Display for Entry {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self {
-            Entry::FreshId(label) => write!(f, "{}=*", label),
-            Entry::MapLabelToKey(label, key) => {
-                write!(f, "{}=@\"{}\"", label, key.replace("\"", "\"\""))
-            }
-            Entry::GenericTuple(name, args) => {
-                write!(f, "{}(", name)?;
-                for (index, arg) in args.iter().enumerate() {
-                    if index > 0 {
-                        write!(f, ",")?;
-                    }
-                    write!(f, "{}", arg)?;
-                }
-                write!(f, ")")
-            }
-            Entry::Comment(line) => write!(f, "// {}", line),
-        }
-    }
-}
-
-#[derive(Debug, Copy, Clone)]
-// Identifiers of the form #0, #1...
-pub struct Label(u32);
-
-impl fmt::Display for Label {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "#{:x}", self.0)
-    }
-}
-
-// Some untyped argument to a TrapEntry.
-#[derive(Debug, Clone)]
-pub enum Arg {
-    Label(Label),
-    Int(usize),
-    String(String),
-}
-
-const MAX_STRLEN: usize = 1048576;
-
-impl fmt::Display for Arg {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self {
-            Arg::Label(x) => write!(f, "{}", x),
-            Arg::Int(x) => write!(f, "{}", x),
-            Arg::String(x) => write!(
-                f,
-                "\"{}\"",
-                limit_string(x, MAX_STRLEN).replace("\"", "\"\"")
-            ),
-        }
-    }
-}
-
-pub struct Program(Vec<Entry>);
-
-impl fmt::Display for Program {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let mut text = String::new();
-        for trap_entry in &self.0 {
-            text.push_str(&format!("{}\n", trap_entry));
-        }
-        write!(f, "{}", text)
-    }
-}
-
-pub fn full_id_for_file(normalized_path: &str) -> String {
-    format!("{};sourcefile", escape_key(normalized_path))
-}
-
-pub fn full_id_for_folder(normalized_path: &str) -> String {
-    format!("{};folder", escape_key(normalized_path))
-}
-
-/// Escapes a string for use in a TRAP key, by replacing special characters with
-/// HTML entities.
-fn escape_key<'a, S: Into<Cow<'a, str>>>(key: S) -> Cow<'a, str> {
-    fn needs_escaping(c: char) -> bool {
-        matches!(c, '&' | '{' | '}' | '"' | '@' | '#')
-    }
-
-    let key = key.into();
-    if key.contains(needs_escaping) {
-        let mut escaped = String::with_capacity(2 * key.len());
-        for c in key.chars() {
-            match c {
-                '&' => escaped.push_str("&amp;"),
-                '{' => escaped.push_str("&lbrace;"),
-                '}' => escaped.push_str("&rbrace;"),
-                '"' => escaped.push_str("&quot;"),
-                '@' => escaped.push_str("&commat;"),
-                '#' => escaped.push_str("&num;"),
-                _ => escaped.push(c),
-            }
-        }
-        Cow::Owned(escaped)
-    } else {
-        key
-    }
-}
-
-/// Limit the length (in bytes) of a string. If the string's length in bytes is
-/// less than or equal to the limit then the entire string is returned. Otherwise
-/// the string is sliced at the provided limit. If there is a multi-byte character
-/// at the limit then the returned slice will be slightly shorter than the limit to
-/// avoid splitting that multi-byte character.
-fn limit_string(string: &str, max_size: usize) -> &str {
-    if string.len() <= max_size {
-        return string;
-    }
-    let p = string.as_bytes();
-    let mut index = max_size;
-    // We want to clip the string at [max_size]; however, the character at that position
-    // may span several bytes. We need to find the first byte of the character. In UTF-8
-    // encoded data any byte that matches the bit pattern 10XXXXXX is not a start byte.
-    // Therefore we decrement the index as long as there are bytes matching this pattern.
-    // This ensures we cut the string at the border between one character and another.
-    while index > 0 && (p[index] & 0b11000000) == 0b10000000 {
-        index -= 1;
-    }
-    &string[0..index]
-}
-
-#[derive(Clone, Copy)]
-pub enum Compression {
-    None,
-    Gzip,
-}
-
-impl Compression {
-    pub fn from_env(var_name: &str) -> Compression {
-        match std::env::var(var_name) {
-            Ok(method) => match Compression::from_string(&method) {
-                Some(c) => c,
-                None => {
-                    tracing::error!("Unknown compression method '{}'; using gzip.", &method);
-                    Compression::Gzip
-                }
-            },
-            // Default compression method if the env var isn't set:
-            Err(_) => Compression::Gzip,
-        }
-    }
-
-    pub fn from_string(s: &str) -> Option<Compression> {
-        match s.to_lowercase().as_ref() {
-            "none" => Some(Compression::None),
-            "gzip" => Some(Compression::Gzip),
-            _ => None,
-        }
-    }
-
-    pub fn extension(&self) -> &str {
-        match self {
-            Compression::None => "trap",
-            Compression::Gzip => "trap.gz",
-        }
-    }
-}
-
-#[test]
-fn limit_string_test() {
-    assert_eq!("hello", limit_string(&"hello world".to_owned(), 5));
-    assert_eq!("hi ☹", limit_string(&"hi ☹☹".to_owned(), 6));
-    assert_eq!("hi ", limit_string(&"hi ☹☹".to_owned(), 5));
-}
-
-#[test]
-fn escape_key_test() {
-    assert_eq!("foo!", escape_key("foo!"));
-    assert_eq!("foo&lbrace;&rbrace;", escape_key("foo{}"));
-    assert_eq!("&lbrace;&rbrace;", escape_key("{}"));
-    assert_eq!("", escape_key(""));
-    assert_eq!("/path/to/foo.rb", escape_key("/path/to/foo.rb"));
-    assert_eq!(
-        "/path/to/foo&amp;&lbrace;&rbrace;&quot;&commat;&num;.rb",
-        escape_key("/path/to/foo&{}\"@#.rb")
-    );
-}