Use a single TrapWriter

aibaars · aibaars · commit c568162256f6 · 2021-07-29T12:50:27.000+02:00
The output of two distinct TrapWriters should not be written to the
same TRAP file because this causes name clashes between TRAP labels.
diff --git a/extractor/src/extractor.rs b/extractor/src/extractor.rs
@@ -3,11 +3,13 @@ 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};
 
-struct TrapWriter {
+pub struct TrapWriter {
     /// The accumulated trap entries
     trap_output: Vec<TrapEntry>,
     /// A counter for generating fresh labels
@@ -16,7 +18,7 @@ struct TrapWriter {
     global_keys: std::collections::HashMap<String, Label>,
 }
 
-fn new_trap_writer() -> TrapWriter {
+pub fn new_trap_writer() -> TrapWriter {
     TrapWriter {
         counter: 0,
         trap_output: Vec::new(),
@@ -147,17 +149,22 @@ impl TrapWriter {
     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))
+    }
 }
 
 /// Extracts the source file at `path`, which is assumed to be canonicalized.
 pub fn extract(
     language: Language,
     language_prefix: &str,
     schema: &NodeTypeMap,
+    trap_writer: &mut TrapWriter,
     path: &Path,
     source: &Vec<u8>,
     ranges: &[Range],
-) -> std::io::Result<Program> {
+) -> std::io::Result<()> {
     let span = span!(
         Level::TRACE,
         "extract",
@@ -172,8 +179,7 @@ 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");
-    let mut trap_writer = new_trap_writer();
-    trap_writer.comment(format!("Auto-generated TRAP file for {}", path.display()));
+    &trap_writer.comment(format!("Auto-generated TRAP file for {}", path.display()));
     let file_label = &trap_writer.populate_file(path);
     let mut visitor = Visitor {
         source: &source,
@@ -190,7 +196,7 @@ pub fn extract(
     traverse(&tree, &mut visitor);
 
     parser.reset();
-    Ok(Program(visitor.trap_writer.trap_output))
+    Ok(())
 }
 
 /// Escapes a string for use in a TRAP key, by replacing special characters with
@@ -290,7 +296,7 @@ struct Visitor<'a> {
     /// The source code as a UTF-8 byte array
     source: &'a Vec<u8>,
     /// A TrapWriter to accumulate trap entries
-    trap_writer: TrapWriter,
+    trap_writer: &'a mut TrapWriter,
     /// A counter for tokens
     token_counter: usize,
     /// A counter for top-level child nodes
diff --git a/extractor/src/main.rs b/extractor/src/main.rs
@@ -6,7 +6,7 @@ use clap;
 use flate2::write::GzEncoder;
 use rayon::prelude::*;
 use std::fs;
-use std::io::{BufRead, BufWriter, Write};
+use std::io::{BufRead, BufWriter};
 use std::path::{Path, PathBuf};
 use tree_sitter::{Language, Parser, Range};
 
@@ -135,17 +135,18 @@ fn main() -> std::io::Result<()> {
         let src_archive_file = path_for(&src_archive_dir, &path, "");
         let mut source = std::fs::read(&path)?;
         let code_ranges;
-        let erb_trap;
+        let mut trap_writer = extractor::new_trap_writer();
         if path.extension().map_or(false, |x| x == "erb") {
             tracing::info!("scanning: {}", path.display());
-            erb_trap = Some(extractor::extract(
+            extractor::extract(
                 erb,
                 "erb",
                 &erb_schema,
+                &mut trap_writer,
                 &path,
                 &source,
                 &[],
-            )?);
+            )?;
 
             let (ranges, line_breaks) = scan_erb(erb, &source);
             for i in line_breaks {
@@ -156,27 +157,26 @@ fn main() -> std::io::Result<()> {
             code_ranges = ranges;
         } else {
             code_ranges = vec![];
-            erb_trap = None;
         }
-        let trap = extractor::extract(language, "ruby", &schema, &path, &source, &code_ranges)?;
+        extractor::extract(
+            language,
+            "ruby",
+            &schema,
+            &mut trap_writer,
+            &path,
+            &source,
+            &code_ranges,
+        )?;
         std::fs::create_dir_all(&src_archive_file.parent().unwrap())?;
         std::fs::copy(&path, &src_archive_file)?;
         std::fs::create_dir_all(&trap_file.parent().unwrap())?;
         let trap_file = std::fs::File::create(&trap_file)?;
         let mut trap_file = BufWriter::new(trap_file);
         match trap_compression {
-            TrapCompression::None => {
-                if let Some(erb_trap) = erb_trap {
-                    write!(trap_file, "{}", erb_trap)?;
-                }
-                write!(trap_file, "{}", trap)
-            }
+            TrapCompression::None => trap_writer.output(&mut trap_file),
             TrapCompression::Gzip => {
                 let mut compressed_writer = GzEncoder::new(trap_file, flate2::Compression::fast());
-                if let Some(erb_trap) = erb_trap {
-                    write!(compressed_writer, "{}", erb_trap)?;
-                }
-                write!(compressed_writer, "{}", trap)
+                trap_writer.output(&mut compressed_writer)
             }
         }
     })
