tree-sitter-extractor: Split direct and desugaring extractors

Previously the `simple` multi-language extractor carried an optional
desugarer, so every language (including plain tree-sitter ones such as
ql, dbscheme, json and blame) went through the same desugaring-aware
extraction path.

This commit splits it into two front-ends that share a private driver:

  - `simple`: pure tree-sitter extraction with no desugaring. Comments
    and other `extra` nodes are emitted inline as tokens. (The extractor
    then extracts these as usual.)
  - `desugaring`: parses source into a `ParsedTree` (a yeast AST plus
    side-channel `extra` tokens) and rewrites the AST through a
    `yeast::Desugarer` before extraction. The parser is a closure, so
    both tree-sitter grammars (via `tree_sitter_parser`) and custom
    parsers plug in the same way.

The shared multi-file plumbing (threading, glob matching, source-archive
copying, TRAP writing) lives in a new private `driver` module behind a
`LanguageExtractor` trait, so neither front-end duplicates it.

`extract` no longer takes an optional desugarer (it always walks the
parse tree directly); `extract_parsed` takes a required desugarer. ql
and ruby use the direct path; the unified Swift extractor uses the
desugaring path.

Also rename the new side-channel identifiers from "trivia" to "extra"
(ExtraToken, ParsedTree.extras, emit_extra, ...) to match tree-sitter's
own `is_extra()` terminology. The pre-existing `*_trivia_tokeninfo`
relation is left unchanged for a separate change.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

ql/extractor/src/extractor.rs

@@ -29,28 +29,24 @@ pub fn run(options: Options) -> std::io::Result<()> {
                 prefix: "ql",
                 ts_language: tree_sitter_ql::LANGUAGE.into(),
                 node_types: tree_sitter_ql::NODE_TYPES,
-                desugar: None,
                 file_globs: vec!["*.ql".into(), "*.qll".into()],
             },
             simple::LanguageSpec {
                 prefix: "dbscheme",
                 ts_language: tree_sitter_ql_dbscheme::LANGUAGE.into(),
                 node_types: tree_sitter_ql_dbscheme::NODE_TYPES,
-                desugar: None,
                 file_globs: vec!["*.dbscheme".into()],
             },
             simple::LanguageSpec {
                 prefix: "json",
                 ts_language: tree_sitter_json::LANGUAGE.into(),
                 node_types: tree_sitter_json::NODE_TYPES,
-                desugar: None,
                 file_globs: vec!["*.json".into(), "*.jsonl".into(), "*.jsonc".into()],
             },
             simple::LanguageSpec {
                 prefix: "blame",
                 ts_language: tree_sitter_blame::LANGUAGE.into(),
                 node_types: tree_sitter_blame::NODE_TYPES,
-                desugar: None,
                 file_globs: vec!["*.blame".into()],
             },
         ],

ruby/extractor/src/extractor.rs

@@ -94,7 +94,9 @@ pub fn run(options: Options) -> std::io::Result<()> {
         node_types::read_node_types_str("erb", tree_sitter_embedded_template::NODE_TYPES)?;
     let lines: std::io::Result<Vec<String>> = std::io::BufReader::new(file_list).lines().collect();
     let lines = lines?;
-    let source_root = std::env::current_dir().ok().and_then(|d| d.canonicalize().ok());
+    let source_root = std::env::current_dir()
+        .ok()
+        .and_then(|d| d.canonicalize().ok());
     lines
         .par_iter()
         .try_for_each(|line| {
@@ -126,7 +128,6 @@ pub fn run(options: Options) -> std::io::Result<()> {
                     &path,
                     &source,
                     &[],
-                    None,
                 );
 
                 let (ranges, line_breaks) = scan_erb(
@@ -215,7 +216,6 @@ pub fn run(options: Options) -> std::io::Result<()> {
                 &path,
                 &source,
                 &code_ranges,
-                None,
             );
             std::fs::create_dir_all(src_archive_file.parent().unwrap())?;
             if needs_conversion {

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

@@ -0,0 +1,104 @@
+//! Extraction for languages that rewrite their syntax tree before extraction.
+//!
+//! Unlike [`crate::extractor::simple`] (direct tree-sitter extraction), a
+//! desugaring language parses source into a [`ParsedTree`] — a `yeast::Ast`
+//! plus side-channel `extra` tokens (comments and similar) — and rewrites the
+//! AST through a [`yeast::Desugarer`] before emitting TRAP. The parser is a
+//! closure, so both tree-sitter grammars (via
+//! [`crate::extractor::tree_sitter_parser`]) and fully custom parsers plug in
+//! the same way.
+
+use crate::trap;
+use std::path::PathBuf;
+
+use crate::diagnostics;
+use crate::extractor::ParsedTree;
+use crate::extractor::driver::{self, LanguageExtractor};
+use crate::node_types::{self, NodeTypeMap};
+
+/// A parser turns source bytes into a [`ParsedTree`]. Tree-sitter grammars plug
+/// in via [`crate::extractor::tree_sitter_parser`]; custom (non-tree-sitter)
+/// parsers supply their own closure.
+pub type Parser = Box<dyn Fn(&[u8]) -> Result<ParsedTree, String> + Send + Sync>;
+
+pub struct LanguageSpec {
+    pub prefix: &'static str,
+    /// The parser: source -> `yeast::Ast` + `extra` tokens (see [`Parser`]).
+    pub parser: Parser,
+    /// Fallback TRAP schema, used only when `desugarer` does not supply its own
+    /// output schema (via [`yeast::Desugarer::output_node_types_yaml`]). May be
+    /// empty for a custom parser whose desugarer always provides the schema.
+    pub node_types: &'static str,
+    /// The desugarer applied to the parsed AST before extraction. Its
+    /// `output_node_types_yaml()` (when set) provides the TRAP schema.
+    ///
+    /// `Box<dyn yeast::Desugarer>` so the shared extractor is agnostic to the
+    /// user-defined context type the desugarer uses internally.
+    pub desugarer: Box<dyn yeast::Desugarer>,
+    pub file_globs: Vec<String>,
+}
+
+impl LanguageExtractor for LanguageSpec {
+    fn file_globs(&self) -> &[String] {
+        &self.file_globs
+    }
+
+    fn build_schema(&self) -> std::io::Result<NodeTypeMap> {
+        let effective_node_types: String = match self.desugarer.output_node_types_yaml() {
+            Some(yaml) => yeast::node_types_yaml::convert(yaml).map_err(|e| {
+                std::io::Error::other(format!(
+                    "Failed to convert YAML node-types to JSON for {}: {e}",
+                    self.prefix
+                ))
+            })?,
+            None => self.node_types.to_string(),
+        };
+        node_types::read_node_types_str(self.prefix, &effective_node_types)
+    }
+
+    fn extract_file(
+        &self,
+        schema: &NodeTypeMap,
+        diagnostics_writer: &mut diagnostics::LogWriter,
+        trap_writer: &mut trap::Writer,
+        path: &std::path::Path,
+        source: &[u8],
+    ) {
+        crate::extractor::extract_parsed(
+            self.parser.as_ref(),
+            self.prefix,
+            schema,
+            diagnostics_writer,
+            trap_writer,
+            None,
+            path,
+            source,
+            self.desugarer.as_ref(),
+        );
+    }
+}
+
+pub struct Extractor {
+    pub prefix: String,
+    pub languages: Vec<LanguageSpec>,
+    pub trap_dir: PathBuf,
+    pub source_archive_dir: PathBuf,
+    pub file_lists: Vec<PathBuf>,
+    // Typically constructed via `trap::Compression::from_env`.
+    // This allow us to report the error using our diagnostics system
+    // without exposing it to consumers.
+    pub trap_compression: Result<trap::Compression, String>,
+}
+
+impl Extractor {
+    pub fn run(&self) -> std::io::Result<()> {
+        driver::run_extractor(
+            &self.prefix,
+            &self.languages,
+            &self.trap_dir,
+            &self.source_archive_dir,
+            &self.file_lists,
+            &self.trap_compression,
+        )
+    }
+}

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

@@ -0,0 +1,210 @@
+//! Shared multi-file extraction driver.
+//!
+//! The `simple` (direct tree-sitter) and `desugaring` (parse + desugar)
+//! extractors differ only in how a language's schema is built and how a single
+//! file is extracted. Everything else — threading, matching files to languages
+//! by glob, writing TRAP, and copying into the source archive — is identical
+//! and lives here, parameterised over the [`LanguageExtractor`] trait.
+
+use globset::{GlobBuilder, GlobSetBuilder};
+use rayon::prelude::*;
+use std::fs::File;
+use std::io::BufRead;
+use std::path::{Path, PathBuf};
+
+use crate::diagnostics;
+use crate::file_paths;
+use crate::node_types::NodeTypeMap;
+use crate::trap;
+
+/// A language that [`run_extractor`] can process: it knows its file globs, its
+/// TRAP schema, and how to extract a single file. Implemented by
+/// [`super::simple::LanguageSpec`] (direct tree-sitter extraction) and
+/// [`super::desugaring::LanguageSpec`] (parse into an AST and desugar it).
+pub(crate) trait LanguageExtractor: Sync {
+    /// The file-name globs that select files for this language.
+    fn file_globs(&self) -> &[String];
+    /// Build the TRAP node-type schema used to validate emitted tuples.
+    fn build_schema(&self) -> std::io::Result<NodeTypeMap>;
+    /// Extract a single file's `source` into `trap_writer`.
+    fn extract_file(
+        &self,
+        schema: &NodeTypeMap,
+        diagnostics_writer: &mut diagnostics::LogWriter,
+        trap_writer: &mut trap::Writer,
+        path: &Path,
+        source: &[u8],
+    );
+}
+
+/// Drive extraction over `languages` for every file listed in `file_lists`.
+///
+/// Sets up the thread pool, builds a combined glob set, and for each input file
+/// dispatches to the matching language's [`LanguageExtractor::extract_file`],
+/// writing the resulting TRAP and a source-archive copy.
+pub(crate) fn run_extractor<L: LanguageExtractor>(
+    prefix: &str,
+    languages: &[L],
+    trap_dir: &Path,
+    source_archive_dir: &Path,
+    file_lists: &[PathBuf],
+    trap_compression: &Result<trap::Compression, String>,
+) -> std::io::Result<()> {
+    tracing::info!("Extraction started");
+    let diagnostics = diagnostics::DiagnosticLoggers::new(prefix);
+    let mut main_thread_logger = diagnostics.logger();
+    let num_threads = match crate::options::num_threads() {
+        Ok(num) => num,
+        Err(e) => {
+            main_thread_logger.write(
+                main_thread_logger
+                    .new_entry("configuration-error", "Configuration error")
+                    .message(
+                        "{}; defaulting to 1 thread.",
+                        &[diagnostics::MessageArg::Code(&e)],
+                    )
+                    .severity(diagnostics::Severity::Warning),
+            );
+            1
+        }
+    };
+    tracing::info!(
+        "Using {} {}",
+        num_threads,
+        if num_threads == 1 {
+            "thread"
+        } else {
+            "threads"
+        }
+    );
+    let trap_compression = match trap_compression {
+        Ok(x) => *x,
+        Err(e) => {
+            main_thread_logger.write(
+                main_thread_logger
+                    .new_entry("configuration-error", "Configuration error")
+                    .message("{}; using gzip.", &[diagnostics::MessageArg::Code(e)])
+                    .severity(diagnostics::Severity::Warning),
+            );
+            trap::Compression::Gzip
+        }
+    };
+    drop(main_thread_logger);
+
+    rayon::ThreadPoolBuilder::new()
+        .num_threads(num_threads)
+        .build_global()
+        .unwrap();
+
+    let file_lists: Vec<File> = file_lists
+        .iter()
+        .map(|file_list| {
+            File::open(file_list)
+                .unwrap_or_else(|_| panic!("Unable to open file list at {file_list:?}"))
+        })
+        .collect();
+
+    let mut schemas = vec![];
+    for lang in languages {
+        schemas.push(lang.build_schema()?);
+    }
+
+    // Construct a single globset containing all language globs,
+    // and a mapping from glob index to language index.
+    let (globset, glob_language_mapping) = {
+        let mut builder = GlobSetBuilder::new();
+        let mut glob_lang_mapping = vec![];
+        for (i, lang) in languages.iter().enumerate() {
+            for glob_str in lang.file_globs() {
+                let glob = GlobBuilder::new(glob_str)
+                    .literal_separator(true)
+                    .build()
+                    .expect("invalid glob");
+                builder.add(glob);
+                glob_lang_mapping.push(i);
+            }
+        }
+        (
+            builder.build().expect("failed to build globset"),
+            glob_lang_mapping,
+        )
+    };
+
+    let path_transformer = file_paths::load_path_transformer()?;
+
+    let lines: std::io::Result<Vec<String>> = file_lists
+        .iter()
+        .flat_map(|file_list| std::io::BufReader::new(file_list).lines())
+        .collect();
+    let lines = lines?;
+
+    lines
+        .par_iter()
+        .try_for_each(|line| {
+            let mut diagnostics_writer = diagnostics.logger();
+            let path = PathBuf::from(line).canonicalize()?;
+            let src_archive_file = crate::file_paths::path_for(
+                source_archive_dir,
+                &path,
+                "",
+                path_transformer.as_ref(),
+            );
+            let source = std::fs::read(&path)?;
+            let mut trap_writer = trap::Writer::new();
+
+            match path.file_name() {
+                None => {
+                    tracing::error!(?path, "No file name found, skipping file.");
+                }
+                Some(filename) => {
+                    let matches = globset.matches(filename);
+                    if matches.is_empty() {
+                        tracing::error!(?path, "No matching language found, skipping file.");
+                    } else {
+                        let mut languages_processed = vec![false; languages.len()];
+
+                        for m in matches {
+                            let i = glob_language_mapping[m];
+                            if languages_processed[i] {
+                                continue;
+                            }
+                            languages_processed[i] = true;
+                            let lang = &languages[i];
+
+                            lang.extract_file(
+                                &schemas[i],
+                                &mut diagnostics_writer,
+                                &mut trap_writer,
+                                &path,
+                                &source,
+                            );
+                            std::fs::create_dir_all(src_archive_file.parent().unwrap())?;
+                            std::fs::copy(&path, &src_archive_file)?;
+                            write_trap(trap_dir, &path, &trap_writer, trap_compression)?;
+                        }
+                    }
+                }
+            }
+            Ok(()) as std::io::Result<()>
+        })
+        .expect("failed to extract files");
+
+    let path = PathBuf::from("extras");
+    let mut trap_writer = trap::Writer::new();
+    crate::extractor::populate_empty_location(&mut trap_writer);
+
+    let res = write_trap(trap_dir, &path, &trap_writer, trap_compression);
+    tracing::info!("Extraction complete");
+    res
+}
+
+fn write_trap(
+    trap_dir: &Path,
+    path: &Path,
+    trap_writer: &trap::Writer,
+    trap_compression: trap::Compression,
+) -> std::io::Result<()> {
+    let trap_file = crate::file_paths::path_for(trap_dir, path, trap_compression.extension(), None);
+    std::fs::create_dir_all(trap_file.parent().unwrap())?;
+    trap_writer.write_to_file(&trap_file, trap_compression)
+}

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

@@ -16,6 +16,8 @@ use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
 use tree_sitter::{Language, Node, Parser, Range, Tree};
 
+pub mod desugaring;
+mod driver;
 pub mod simple;
 
 /// Trait abstracting over tree-sitter and yeast node types for extraction.
@@ -294,6 +296,10 @@ pub fn location_label(writer: &mut trap::Writer, location: trap::Location) -> tr
 /// caller's responsibility, allowing it to be done once and shared across
 /// files.
 #[allow(clippy::too_many_arguments)]
+/// Extract a file with a tree-sitter grammar, walking the parse tree directly.
+/// Comments and other `extra` nodes are emitted inline as tokens. This is the
+/// path for languages that don't desugar their syntax tree (and hence have no
+/// `extra` side channel); desugaring languages use [`extract_parsed`] instead.
 pub fn extract(
     language: &Language,
     language_prefix: &str,
@@ -304,7 +310,6 @@ pub fn extract(
     path: &Path,
     source: &[u8],
     ranges: &[Range],
-    desugarer: Option<&dyn yeast::Desugarer>,
 ) {
     let path_str = file_paths::normalize_and_transform_path(path, transformer);
     let source_root = std::env::current_dir()
@@ -337,21 +342,173 @@ pub fn extract(
         schema,
     );
 
-    if let Some(desugarer) = desugarer {
-        let ast = desugarer
-            .run_from_tree(&tree, source)
-            .unwrap_or_else(|e| panic!("Desugaring failed for {path_str}: {e}"));
-        traverse_yeast(&ast, &mut visitor);
-        // Comments and other `extra` nodes are not represented in the desugared
-        // AST, so recover them directly from the original parse tree.
-        traverse_extras(&tree, &mut visitor);
-    } else {
-        traverse(&tree, &mut visitor);
-    }
+    traverse(&tree, &mut visitor);
 
     parser.reset();
 }
 
+/// A source tree produced by a parser: the raw `yeast::Ast` (pre-desugaring)
+/// plus side-channel `extra` tokens (comments and similar) that are not
+/// attached to the AST proper.
+pub struct ParsedTree {
+    pub ast: yeast::Ast,
+    pub extras: Vec<ExtraToken>,
+}
+
+/// A piece of side-channel `extra` content (a comment or unexpected text)
+/// recovered by a parser. `kind` is a language-defined id written verbatim into
+/// the `<lang>_trivia_tokeninfo` table (the tree-sitter parser uses the
+/// grammar's node kind id here; a custom parser supplies its own stable
+/// enumeration).
+pub struct ExtraToken {
+    pub kind: usize,
+    pub range: yeast::Range,
+    pub text: String,
+}
+
+/// Build a parser closure that parses `source` with a tree-sitter `language`,
+/// producing a [`ParsedTree`]: a `yeast::Ast` (via [`yeast::Ast::from_tree`])
+/// plus the `extra` nodes (comments and similar) recovered as side-channel
+/// tokens. This lets a tree-sitter language plug into the same
+/// `ParsedTree`-producing interface as a custom parser.
+pub fn tree_sitter_parser(
+    language: tree_sitter::Language,
+) -> impl Fn(&[u8]) -> Result<ParsedTree, String> + Send + Sync {
+    move |source: &[u8]| {
+        let mut parser = Parser::new();
+        parser
+            .set_language(&language)
+            .map_err(|e| format!("failed to set tree-sitter language: {e}"))?;
+        let tree = parser
+            .parse(source, None)
+            .ok_or_else(|| "tree-sitter failed to parse".to_string())?;
+        let ast = yeast::Ast::from_tree_with_schema_and_source(
+            yeast::schema::from_language(&language),
+            &tree,
+            &language,
+            source.to_vec(),
+        );
+        let mut extras = Vec::new();
+        collect_extras(tree.root_node(), source, &mut extras);
+        Ok(ParsedTree { ast, extras })
+    }
+}
+
+/// Collect `extra` nodes (comments and similar) under `node` into `out` as
+/// [`ExtraToken`]s, keyed by the grammar's node kind id, for a desugaring
+/// language whose rewritten AST won't retain them.
+fn collect_extras(node: Node<'_>, source: &[u8], out: &mut Vec<ExtraToken>) {
+    let mut cursor = node.walk();
+    for child in node.children(&mut cursor) {
+        if child.is_extra() {
+            let text = String::from_utf8_lossy(&source[child.byte_range()]).into_owned();
+            out.push(ExtraToken {
+                kind: child.kind_id() as usize,
+                range: child.range().into(),
+                text,
+            });
+        } else {
+            collect_extras(child, source, out);
+        }
+    }
+}
+
+/// Extract a file from a [`ParsedTree`]-producing parser that desugars. The
+/// parser yields a `yeast::Ast` plus side-channel `extra` tokens; the AST is
+/// rewritten through `desugarer` ([`yeast::Desugarer::run_from_ast`]) before
+/// TRAP extraction, and the `extra` tokens (comments and similar, which the
+/// desugared AST does not carry) are emitted from the side channel. Both
+/// tree-sitter grammars (via [`tree_sitter_parser`]) and custom parsers plug in
+/// here; languages that don't desugar use [`extract`] instead.
+#[allow(clippy::too_many_arguments)]
+pub fn extract_parsed(
+    parse: &(dyn Fn(&[u8]) -> Result<ParsedTree, String> + Send + Sync),
+    language_prefix: &str,
+    schema: &NodeTypeMap,
+    diagnostics_writer: &mut diagnostics::LogWriter,
+    trap_writer: &mut trap::Writer,
+    transformer: Option<&file_paths::PathTransformer>,
+    path: &Path,
+    source: &[u8],
+    desugarer: &dyn yeast::Desugarer,
+) {
+    let path_str = file_paths::normalize_and_transform_path(path, transformer);
+    let source_root = std::env::current_dir()
+        .ok()
+        .and_then(|d| d.canonicalize().ok());
+    let diagnostics_path = file_paths::relativize_for_diagnostic(path, source_root.as_deref());
+    let span = tracing::span!(tracing::Level::TRACE, "extract", file = %path_str);
+    let _enter = span.enter();
+    tracing::debug!("extracting: {}", path_str);
+
+    trap_writer.comment(format!("Auto-generated TRAP file for {path_str}"));
+    let file_label = populate_file(trap_writer, path, transformer);
+    let mut visitor = Visitor::new(
+        source,
+        diagnostics_writer,
+        trap_writer,
+        &diagnostics_path,
+        file_label,
+        language_prefix,
+        schema,
+    );
+
+    let parsed = parse(source).unwrap_or_else(|e| panic!("Parsing failed for {path_str}: {e}"));
+    let ast = desugarer
+        .run_from_ast(parsed.ast)
+        .unwrap_or_else(|e| panic!("Desugaring failed for {path_str}: {e}"));
+    traverse_yeast(&ast, &mut visitor);
+    // Comments and other `extra` tokens are not part of the desugared AST; emit
+    // them directly from the parser's side channel.
+    for extra in &parsed.extras {
+        visitor.emit_extra(extra);
+    }
+}
+
+/// A lightweight [`AstNode`] over a piece of side-channel `extra` content
+/// recovered by a custom parser, so it can reuse the tree-sitter location
+/// machinery.
+struct ExtraNode {
+    range: yeast::Range,
+    text: String,
+}
+
+impl AstNode for ExtraNode {
+    fn kind(&self) -> &str {
+        "extra"
+    }
+    fn is_named(&self) -> bool {
+        false
+    }
+    fn is_missing(&self) -> bool {
+        false
+    }
+    fn is_error(&self) -> bool {
+        false
+    }
+    fn is_extra(&self) -> bool {
+        true
+    }
+    fn start_position(&self) -> tree_sitter::Point {
+        tree_sitter::Point {
+            row: self.range.start_point.row,
+            column: self.range.start_point.column,
+        }
+    }
+    fn end_position(&self) -> tree_sitter::Point {
+        tree_sitter::Point {
+            row: self.range.end_point.row,
+            column: self.range.end_point.column,
+        }
+    }
+    fn byte_range(&self) -> std::ops::Range<usize> {
+        self.range.start_byte..self.range.end_byte
+    }
+    fn opt_string_content(&self) -> Option<String> {
+        Some(self.text.clone())
+    }
+}
+
 struct ChildNode {
     field_name: Option<&'static str>,
     label: trap::Label,
@@ -415,10 +572,11 @@ impl<'a> Visitor<'a> {
         }
     }
 
-    /// Emits a `TriviaToken` for the given `extra` node (e.g. a comment) from
-    /// the original parse tree. Trivia tokens carry a location and their source
-    /// text, but are not attached to a parent in the (possibly desugared) AST.
-    fn emit_trivia_token(&mut self, node: &Node) {
+    /// Emit an `extra` token (a location row plus a `trivia_tokeninfo` row) for
+    /// any [`AstNode`], with an explicit `kind` id. Shared by the tree-sitter
+    /// path (kind = the grammar's node kind id) and the custom-parser path
+    /// (kind = a language-defined `extra` kind id).
+    fn emit_extra_from<N: AstNode>(&mut self, node: &N, kind: usize) {
         let id = self.trap_writer.fresh_id();
         let loc = location_for(self, self.file_label, node);
         let loc_label = location_label(self.trap_writer, loc);
@@ -430,12 +588,21 @@ impl<'a> Visitor<'a> {
             &self.trivia_tokeninfo_table_name,
             vec![
                 trap::Arg::Label(id),
-                trap::Arg::Int(node.kind_id() as usize),
+                trap::Arg::Int(kind),
                 sliced_source_arg(self.source, node),
             ],
         );
     }
 
+    /// Emit an `extra` token recovered from a parser's side channel.
+    fn emit_extra(&mut self, extra: &ExtraToken) {
+        let node = ExtraNode {
+            range: extra.range,
+            text: extra.text.clone(),
+        };
+        self.emit_extra_from(&node, extra.kind);
+    }
+
     fn record_parse_error(&mut self, loc: trap::Label, mesg: &diagnostics::DiagnosticMessage) {
         self.diagnostics_writer.write(mesg);
         let id = self.trap_writer.fresh_id();
@@ -871,24 +1038,6 @@ fn traverse(tree: &Tree, visitor: &mut Visitor) {
     }
 }
 
-/// Walks the original tree-sitter tree and emits a `TriviaToken` for every
-/// `extra` node (e.g. a comment). Used to preserve comments that would
-/// otherwise be lost after a desugaring pass rewrites the tree.
-fn traverse_extras(tree: &Tree, visitor: &mut Visitor) {
-    emit_extras_in(visitor, tree.root_node());
-}
-
-fn emit_extras_in(visitor: &mut Visitor, node: Node<'_>) {
-    let mut cursor = node.walk();
-    for child in node.children(&mut cursor) {
-        if child.is_extra() {
-            visitor.emit_trivia_token(&child);
-        } else {
-            emit_extras_in(visitor, child);
-        }
-    }
-}
-
 fn traverse_yeast(tree: &yeast::Ast, visitor: &mut Visitor) {
     let mut cursor = tree.walk();
     visitor.enter_node(cursor.node());

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

@@ -1,29 +1,52 @@
-use crate::{file_paths, trap};
-use globset::{GlobBuilder, GlobSetBuilder};
-use rayon::prelude::*;
-use std::fs::File;
-use std::io::BufRead;
-use std::path::{Path, PathBuf};
+use crate::trap;
+use std::path::PathBuf;
 
 use crate::diagnostics;
-use crate::node_types;
-use yeast;
+use crate::extractor::driver::{self, LanguageExtractor};
+use crate::node_types::{self, NodeTypeMap};
 
+/// A tree-sitter language extracted directly from its parse tree, with no
+/// desugaring. Comments and other `extra` nodes are emitted inline as tokens.
+/// Languages that rewrite their syntax tree use
+/// [`crate::extractor::desugaring`] instead.
 pub struct LanguageSpec {
     pub prefix: &'static str,
     pub ts_language: tree_sitter::Language,
     pub node_types: &'static str,
-    /// Optional desugarer. When set, the parsed tree is rewritten through
-    /// the desugarer before TRAP extraction. The desugarer's
-    /// `output_node_types_yaml()` (if set) provides the schema used both
-    /// at runtime (for the rewriter) and for TRAP validation.
-    ///
-    /// `Box<dyn yeast::Desugarer>` so the shared extractor is agnostic to
-    /// the user-defined context type the desugarer uses internally.
-    pub desugar: Option<Box<dyn yeast::Desugarer>>,
     pub file_globs: Vec<String>,
 }
 
+impl LanguageExtractor for LanguageSpec {
+    fn file_globs(&self) -> &[String] {
+        &self.file_globs
+    }
+
+    fn build_schema(&self) -> std::io::Result<NodeTypeMap> {
+        node_types::read_node_types_str(self.prefix, self.node_types)
+    }
+
+    fn extract_file(
+        &self,
+        schema: &NodeTypeMap,
+        diagnostics_writer: &mut diagnostics::LogWriter,
+        trap_writer: &mut trap::Writer,
+        path: &std::path::Path,
+        source: &[u8],
+    ) {
+        crate::extractor::extract(
+            &self.ts_language,
+            self.prefix,
+            schema,
+            diagnostics_writer,
+            trap_writer,
+            None,
+            path,
+            source,
+            &[],
+        );
+    }
+}
+
 pub struct Extractor {
     pub prefix: String,
     pub languages: Vec<LanguageSpec>,
@@ -38,182 +61,13 @@ pub struct Extractor {
 
 impl Extractor {
     pub fn run(&self) -> std::io::Result<()> {
-        tracing::info!("Extraction started");
-        let diagnostics = diagnostics::DiagnosticLoggers::new(&self.prefix);
-        let mut main_thread_logger = diagnostics.logger();
-        let num_threads = match crate::options::num_threads() {
-            Ok(num) => num,
-            Err(e) => {
-                main_thread_logger.write(
-                    main_thread_logger
-                        .new_entry("configuration-error", "Configuration error")
-                        .message(
-                            "{}; defaulting to 1 thread.",
-                            &[diagnostics::MessageArg::Code(&e)],
-                        )
-                        .severity(diagnostics::Severity::Warning),
-                );
-                1
-            }
-        };
-        tracing::info!(
-            "Using {} {}",
-            num_threads,
-            if num_threads == 1 {
-                "thread"
-            } else {
-                "threads"
-            }
-        );
-        let trap_compression = match &self.trap_compression {
-            Ok(x) => *x,
-            Err(e) => {
-                main_thread_logger.write(
-                    main_thread_logger
-                        .new_entry("configuration-error", "Configuration error")
-                        .message("{}; using gzip.", &[diagnostics::MessageArg::Code(e)])
-                        .severity(diagnostics::Severity::Warning),
-                );
-                trap::Compression::Gzip
-            }
-        };
-        drop(main_thread_logger);
-
-        rayon::ThreadPoolBuilder::new()
-            .num_threads(num_threads)
-            .build_global()
-            .unwrap();
-
-        let file_lists: Vec<File> = self
-            .file_lists
-            .iter()
-            .map(|file_list| {
-                File::open(file_list)
-                    .unwrap_or_else(|_| panic!("Unable to open file list at {file_list:?}"))
-            })
-            .collect();
-
-        let mut schemas = vec![];
-        for lang in &self.languages {
-            let effective_node_types: String = match lang
-                .desugar
-                .as_ref()
-                .and_then(|d| d.output_node_types_yaml())
-            {
-                Some(yaml) => yeast::node_types_yaml::convert(yaml).map_err(|e| {
-                    std::io::Error::other(format!(
-                        "Failed to convert YAML node-types to JSON for {}: {e}",
-                        lang.prefix
-                    ))
-                })?,
-                None => lang.node_types.to_string(),
-            };
-            let schema = node_types::read_node_types_str(lang.prefix, &effective_node_types)?;
-            schemas.push(schema);
-        }
-
-        // Construct a single globset containing all language globs,
-        // and a mapping from glob index to language index.
-        let (globset, glob_language_mapping) = {
-            let mut builder = GlobSetBuilder::new();
-            let mut glob_lang_mapping = vec![];
-            for (i, lang) in self.languages.iter().enumerate() {
-                for glob_str in &lang.file_globs {
-                    let glob = GlobBuilder::new(glob_str)
-                        .literal_separator(true)
-                        .build()
-                        .expect("invalid glob");
-                    builder.add(glob);
-                    glob_lang_mapping.push(i);
-                }
-            }
-            (
-                builder.build().expect("failed to build globset"),
-                glob_lang_mapping,
-            )
-        };
-
-        let path_transformer = file_paths::load_path_transformer()?;
-
-        let lines: std::io::Result<Vec<String>> = file_lists
-            .iter()
-            .flat_map(|file_list| std::io::BufReader::new(file_list).lines())
-            .collect();
-        let lines = lines?;
-
-        lines
-            .par_iter()
-            .try_for_each(|line| {
-                let mut diagnostics_writer = diagnostics.logger();
-                let path = PathBuf::from(line).canonicalize()?;
-                let src_archive_file = crate::file_paths::path_for(
-                    &self.source_archive_dir,
-                    &path,
-                    "",
-                    path_transformer.as_ref(),
-                );
-                let source = std::fs::read(&path)?;
-                let mut trap_writer = trap::Writer::new();
-
-                match path.file_name() {
-                    None => {
-                        tracing::error!(?path, "No file name found, skipping file.");
-                    }
-                    Some(filename) => {
-                        let matches = globset.matches(filename);
-                        if matches.is_empty() {
-                            tracing::error!(?path, "No matching language found, skipping file.");
-                        } else {
-                            let mut languages_processed = vec![false; self.languages.len()];
-
-                            for m in matches {
-                                let i = glob_language_mapping[m];
-                                if languages_processed[i] {
-                                    continue;
-                                }
-                                languages_processed[i] = true;
-                                let lang = &self.languages[i];
-
-                                crate::extractor::extract(
-                                    &lang.ts_language,
-                                    lang.prefix,
-                                    &schemas[i],
-                                    &mut diagnostics_writer,
-                                    &mut trap_writer,
-                                    None,
-                                    &path,
-                                    &source,
-                                    &[],
-                                    lang.desugar.as_deref(),
-                                );
-                                std::fs::create_dir_all(src_archive_file.parent().unwrap())?;
-                                std::fs::copy(&path, &src_archive_file)?;
-                                write_trap(&self.trap_dir, &path, &trap_writer, trap_compression)?;
-                            }
-                        }
-                    }
-                }
-                Ok(()) as std::io::Result<()>
-            })
-            .expect("failed to extract files");
-
-        let path = PathBuf::from("extras");
-        let mut trap_writer = trap::Writer::new();
-        crate::extractor::populate_empty_location(&mut trap_writer);
-
-        let res = write_trap(&self.trap_dir, &path, &trap_writer, trap_compression);
-        tracing::info!("Extraction complete");
-        res
+        driver::run_extractor(
+            &self.prefix,
+            &self.languages,
+            &self.trap_dir,
+            &self.source_archive_dir,
+            &self.file_lists,
+            &self.trap_compression,
+        )
     }
 }
-
-fn write_trap(
-    trap_dir: &Path,
-    path: &Path,
-    trap_writer: &trap::Writer,
-    trap_compression: trap::Compression,
-) -> std::io::Result<()> {
-    let trap_file = crate::file_paths::path_for(trap_dir, path, trap_compression.extension(), None);
-    std::fs::create_dir_all(trap_file.parent().unwrap())?;
-    trap_writer.write_to_file(&trap_file, trap_compression)
-}

shared/tree-sitter-extractor/tests/integration_test.rs

@@ -13,7 +13,6 @@ fn simple_extractor() {
         prefix: "ql",
         ts_language: tree_sitter_ql::LANGUAGE.into(),
         node_types: tree_sitter_ql::NODE_TYPES,
-        desugar: None,
         file_globs: vec!["*.qll".into()],
     };
 

shared/tree-sitter-extractor/tests/multiple_languages.rs

@@ -13,14 +13,12 @@ fn multiple_language_extractor() {
         prefix: "ql",
         ts_language: tree_sitter_ql::LANGUAGE.into(),
         node_types: tree_sitter_ql::NODE_TYPES,
-        desugar: None,
         file_globs: vec!["*.qll".into()],
     };
     let lang_json = simple::LanguageSpec {
         prefix: "json",
         ts_language: tree_sitter_json::LANGUAGE.into(),
         node_types: tree_sitter_json::NODE_TYPES,
-        desugar: None,
         file_globs: vec!["*.json".into(), "*Jsonfile".into()],
     };
 

unified/extractor/src/extractor.rs

@@ -2,7 +2,7 @@ use clap::Args;
 use std::path::PathBuf;
 
 use crate::languages;
-use codeql_extractor::extractor::simple;
+use codeql_extractor::extractor::desugaring;
 use codeql_extractor::trap;
 
 #[derive(Args)]
@@ -31,7 +31,7 @@ pub fn run(options: Options) -> std::io::Result<()> {
         lang.prefix = "unified";
     }
 
-    let extractor = simple::Extractor {
+    let extractor = desugaring::Extractor {
         prefix: "unified".to_string(),
         languages,
         trap_dir: options.output_dir,

unified/extractor/src/languages/mod.rs

@@ -1,4 +1,4 @@
-use codeql_extractor::extractor::simple;
+use codeql_extractor::extractor::desugaring;
 
 #[path = "swift/swift.rs"]
 mod swift;
@@ -15,6 +15,6 @@ pub mod swift_adapter;
 /// Shared YEAST output AST schema for all languages.
 pub(crate) const OUTPUT_AST_SCHEMA: &str = include_str!("../../ast_types.yml");
 
-pub fn all_language_specs() -> Vec<simple::LanguageSpec> {
+pub fn all_language_specs() -> Vec<desugaring::LanguageSpec> {
     vec![swift::language_spec(OUTPUT_AST_SCHEMA)]
 }

unified/extractor/src/languages/swift/swift.rs

@@ -1,4 +1,4 @@
-use codeql_extractor::extractor::simple;
+use codeql_extractor::extractor::desugaring;
 use yeast::{ConcreteDesugarer, DesugaringConfig, PhaseKind, Rule, rule, tree};
 
 /// User context propagated from outer rules down to the inner rules that
@@ -1249,18 +1249,18 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
     ]
 }
 
-pub fn language_spec(desugared_ast_schema: &'static str) -> simple::LanguageSpec {
+pub fn language_spec(desugared_ast_schema: &'static str) -> desugaring::LanguageSpec {
     let ts_language: tree_sitter::Language = tree_sitter_swift::LANGUAGE.into();
     let config = DesugaringConfig::<SwiftContext>::new()
         .add_phase("translate", PhaseKind::OneShot, translation_rules())
         .with_output_node_types_yaml(desugared_ast_schema);
     let desugarer = ConcreteDesugarer::new(ts_language.clone(), config)
         .expect("failed to build Swift desugarer");
-    simple::LanguageSpec {
+    desugaring::LanguageSpec {
         prefix: "swift",
-        ts_language,
+        parser: Box::new(codeql_extractor::extractor::tree_sitter_parser(ts_language)),
         node_types: tree_sitter_swift::NODE_TYPES,
         file_globs: vec!["*.swift".into(), "*.swiftinterface".into()],
-        desugar: Some(Box::new(desugarer)),
+        desugarer: Box::new(desugarer),
     }
 }

unified/extractor/tests/corpus_tests.rs

@@ -1,8 +1,8 @@
 use std::fs;
 use std::path::Path;
 
-use codeql_extractor::extractor::simple;
-use yeast::{Runner, dump::dump_ast, dump::dump_ast_with_type_errors};
+use codeql_extractor::extractor::desugaring;
+use yeast::{dump::dump_ast, dump::dump_ast_with_type_errors};
 
 #[path = "../src/languages/mod.rs"]
 mod languages;
@@ -59,40 +59,21 @@ fn render_corpus(case: &CorpusCase) -> String {
     )
 }
 
-fn run_desugaring(lang: &simple::LanguageSpec, input: &str) -> Result<yeast::Ast, String> {
-    match lang.desugar.as_deref() {
-        Some(desugarer) => {
-            // Parse the input ourselves so we don't depend on the desugarer
-            // knowing about the language.
-            let mut parser = tree_sitter::Parser::new();
-            parser
-                .set_language(&lang.ts_language)
-                .map_err(|e| format!("Failed to set language: {e}"))?;
-            let tree = parser
-                .parse(input, None)
-                .ok_or_else(|| "Failed to parse input".to_string())?;
-            desugarer
-                .run_from_tree(&tree, input.as_bytes())
-                .map_err(|e| format!("Desugaring failed: {e}"))
-        }
-        None => {
-            let runner: Runner = Runner::new(lang.ts_language.clone(), &[]);
-            runner
-                .run(input)
-                .map_err(|e| format!("Failed to parse input: {e}"))
-        }
-    }
+/// Parse `input` through the language's parser and desugar it, returning the
+/// mapped AST.
+fn run_desugaring(lang: &desugaring::LanguageSpec, input: &str) -> Result<yeast::Ast, String> {
+    let parsed = (lang.parser)(input.as_bytes())?;
+    lang.desugarer
+        .run_from_ast(parsed.ast)
+        .map_err(|e| format!("Desugaring failed: {e}"))
 }
 
-/// Produce the raw tree-sitter parse tree dump for `input`, with no
-/// desugaring rules applied. Uses a `Runner` with an empty phase list and
-/// the input grammar's own schema.
-fn dump_raw_parse(lang: &simple::LanguageSpec, input: &str) -> Result<String, String> {
-    let runner: Runner = Runner::new(lang.ts_language.clone(), &[]);
-    let ast = runner
-        .run(input)
-        .map_err(|e| format!("Failed to parse input: {e}"))?;
-    Ok(dump_ast(&ast, ast.get_root(), input))
+/// Produce the raw (pre-desugar) parse tree dump for `input` — the `yeast::Ast`
+/// the language's parser builds, before any desugaring rules. Useful for seeing
+/// what the mapping rules operate on.
+fn dump_raw_parse(lang: &desugaring::LanguageSpec, input: &str) -> Result<String, String> {
+    let parsed = (lang.parser)(input.as_bytes())?;
+    Ok(dump_ast(&parsed.ast, parsed.ast.get_root(), input))
 }
 
 /// Collect the set of corpus test "stems" (paths without an extension) under
@@ -122,11 +103,8 @@ fn test_corpus() {
     let corpus_dir = Path::new("tests/corpus");
 
     for lang in all_languages {
-        let output_schema = yeast::node_types_yaml::schema_from_yaml_with_language(
-            languages::OUTPUT_AST_SCHEMA,
-            &lang.ts_language,
-        )
-        .expect("Failed to parse OUTPUT_AST_SCHEMA YAML");
+        let output_schema = yeast::node_types_yaml::schema_from_yaml(languages::OUTPUT_AST_SCHEMA)
+            .expect("Failed to parse OUTPUT_AST_SCHEMA YAML");
 
         let lang_corpus_dir = corpus_dir.join(&lang.prefix);
         if !lang_corpus_dir.exists() {
@@ -236,8 +214,7 @@ fn test_corpus() {
                         );
                         if update_mode {
                             case.expected = actual_dump.trim().to_string();
-                        } else if output_path.exists()
-                            && case.expected.trim() != actual_dump.trim()
+                        } else if output_path.exists() && case.expected.trim() != actual_dump.trim()
                         {
                             failures.push(format!(
                                 "Test failed in {}:\nEXPECTED:\n\n{}\n\nACTUAL:\n\n{}",

unified/extractor/tests/swift_syntax_pipeline.rs

@@ -22,7 +22,7 @@ fn swift_syntax_json_runs_through_the_desugarer() {
         .into_iter()
         .find(|l| l.file_globs.iter().any(|g| g.contains("swift")))
         .expect("swift language spec");
-    let desugarer = lang.desugar.as_deref().expect("swift desugarer");
+    let desugarer = lang.desugarer.as_ref();
 
     // Adapt the swift-syntax JSON into a yeast AST (pure Rust, no Swift FFI).
     let adapted =