Ruby: sort TRAP output

ruby/extractor/Cargo.toml

@@ -20,3 +20,4 @@ num_cpus = "1.13.0"
 regex = "1.5.5"
 encoding = "0.2"
 lazy_static = "1.4.0"
+indexmap = "1.8"

ruby/extractor/src/extractor.rs

@@ -1,6 +1,6 @@
 use crate::trap;
+use indexmap::IndexMap;
 use node_types::{EntryKind, Field, NodeTypeMap, Storage, TypeName};
-use std::collections::BTreeMap as Map;
 use std::collections::BTreeSet as Set;
 use std::fmt;
 use std::path::Path;
@@ -134,7 +134,6 @@ pub fn extract(
     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));
     let file_label = populate_file(trap_writer, path);
     let mut visitor = Visitor::new(
         source,
@@ -414,7 +413,7 @@ impl<'a> Visitor<'a> {
         child_nodes: &[ChildNode],
         parent_id: trap::Label,
     ) -> Option<Vec<trap::Arg>> {
-        let mut map: Map<&Option<String>, (&Field, Vec<trap::Arg>)> = Map::new();
+        let mut map: IndexMap<&Option<String>, (&Field, Vec<trap::Arg>)> = IndexMap::new();
         for field in fields {
             map.insert(&field.name, (field, Vec::new()));
         }

ruby/extractor/src/trap.rs

@@ -4,29 +4,39 @@ use std::io::{BufWriter, Write};
 use std::path::Path;
 
 use flate2::write::GzEncoder;
+use indexmap::IndexMap;
+
+type Tuple = Vec<Arg>;
 
 pub struct Writer {
-    /// The accumulated trap entries
-    trap_output: Vec<Entry>,
+    /// Labels that should be assigned fresh ids, e.g. `#123=*`.
+    fresh_ids: Vec<Label>,
+
+    /// Labels that should be assigned trap keys, e.g. `#7=@"foo"`.
+    global_keys: IndexMap<String, Label>,
+
+    /// Database relations/tables to emit. Each key is the relation/table name,
+    /// each value is a list of tuples/rows.
+    relations: IndexMap<String, Vec<Tuple>>,
+
     /// 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 {
+            fresh_ids: Vec::new(),
+            relations: IndexMap::new(),
+            global_keys: IndexMap::new(),
             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));
+        self.fresh_ids.push(label);
         label
     }
 
@@ -45,18 +55,15 @@ impl Writer {
         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));
+    /// Add a `tuple` to the table of the given `name`.
+    pub fn add_tuple(&mut self, name: &str, tuple: Tuple) {
+        self.relations
+            .entry(name.to_owned())
+            .or_insert_with(Vec::new)
+            .push(tuple);
     }
 
     pub fn write_to_file(&self, path: &Path, compression: Compression) -> std::io::Result<()> {
@@ -75,45 +82,31 @@ impl Writer {
     }
 
     fn write_trap_entries<W: Write>(&self, file: &mut W) -> std::io::Result<()> {
-        for trap_entry in &self.trap_output {
-            writeln!(file, "{}", trap_entry)?;
+        for label in &self.fresh_ids {
+            writeln!(file, "{}=*", label)?;
+        }
+        for (key, label) in &self.global_keys {
+            writeln!(file, "{}=@\"{}\"", label, key.replace("\"", "\"\""))?;
+        }
+        for (name, tuples) in &self.relations {
+            for tuple in tuples {
+                write!(file, "{}(", name)?;
+                let mut first = true;
+                for arg in tuple {
+                    if first {
+                        first = false;
+                    } else {
+                        write!(file, ",")?;
+                    }
+                    write!(file, "{}", arg)?;
+                }
+                writeln!(file, ")")?;
+            }
         }
         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);
@@ -148,18 +141,6 @@ impl fmt::Display for Arg {
     }
 }
 
-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))
 }