From bbf52ce7a76bb8b620801863e9d2598210a6e285 Mon Sep 17 00:00:00 2001 From: Taus Date: Fri, 17 Jul 2026 12:26:55 +0000 Subject: [PATCH] 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/Cargo.lock | Bin 27305 -> 27434 bytes ql/extractor/src/extractor.rs | 4 - ruby/extractor/src/extractor.rs | 6 +- .../src/extractor/desugaring.rs | 104 ++++++++ .../src/extractor/driver.rs | 210 +++++++++++++++ .../src/extractor/mod.rs | 219 +++++++++++++--- .../src/extractor/simple.rs | 240 ++++-------------- .../tests/integration_test.rs | 1 - .../tests/multiple_languages.rs | 2 - unified/extractor/src/extractor.rs | 4 +- unified/extractor/src/languages/mod.rs | 4 +- .../extractor/src/languages/swift/swift.rs | 10 +- unified/extractor/tests/corpus_tests.rs | 59 ++--- .../extractor/tests/swift_syntax_pipeline.rs | 2 +- 14 files changed, 576 insertions(+), 289 deletions(-) create mode 100644 shared/tree-sitter-extractor/src/extractor/desugaring.rs create mode 100644 shared/tree-sitter-extractor/src/extractor/driver.rs diff --git a/ql/Cargo.lock b/ql/Cargo.lock index 5c65d88de6a7a0adb821f61ffd087c0fce338fce..d6ea9ea343ef14975dbb250774bdbae3436b9f99 100644 GIT binary patch delta 60 zcmZ2^m2uTI#toh6dAf delta 22 ecmZ2=jdA5w#toh6lUXukCvQ(r+1#Bmp9uhZ0}2@c diff --git a/ql/extractor/src/extractor.rs b/ql/extractor/src/extractor.rs index 66096442e85..8383d8424ee 100644 --- a/ql/extractor/src/extractor.rs +++ b/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()], }, ], diff --git a/ruby/extractor/src/extractor.rs b/ruby/extractor/src/extractor.rs index d418c144bfc..3749609e788 100644 --- a/ruby/extractor/src/extractor.rs +++ b/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> = 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 { diff --git a/shared/tree-sitter-extractor/src/extractor/desugaring.rs b/shared/tree-sitter-extractor/src/extractor/desugaring.rs new file mode 100644 index 00000000000..c52659750c7 --- /dev/null +++ b/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 Result + 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` so the shared extractor is agnostic to the + /// user-defined context type the desugarer uses internally. + pub desugarer: Box, + pub file_globs: Vec, +} + +impl LanguageExtractor for LanguageSpec { + fn file_globs(&self) -> &[String] { + &self.file_globs + } + + fn build_schema(&self) -> std::io::Result { + 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, + pub trap_dir: PathBuf, + pub source_archive_dir: PathBuf, + pub file_lists: Vec, + // 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, +} + +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, + ) + } +} diff --git a/shared/tree-sitter-extractor/src/extractor/driver.rs b/shared/tree-sitter-extractor/src/extractor/driver.rs new file mode 100644 index 00000000000..d97d8f4c75c --- /dev/null +++ b/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; + /// 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( + prefix: &str, + languages: &[L], + trap_dir: &Path, + source_archive_dir: &Path, + file_lists: &[PathBuf], + trap_compression: &Result, +) -> 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_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> = 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) +} diff --git a/shared/tree-sitter-extractor/src/extractor/mod.rs b/shared/tree-sitter-extractor/src/extractor/mod.rs index 474dc096a2a..1a8b9e82016 100644 --- a/shared/tree-sitter-extractor/src/extractor/mod.rs +++ b/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, +} + +/// 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 `_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 + 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) { + 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 + 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 { + self.range.start_byte..self.range.end_byte + } + fn opt_string_content(&self) -> Option { + 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(&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()); diff --git a/shared/tree-sitter-extractor/src/extractor/simple.rs b/shared/tree-sitter-extractor/src/extractor/simple.rs index 9ba6f21778c..1c6691fa8cf 100644 --- a/shared/tree-sitter-extractor/src/extractor/simple.rs +++ b/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` so the shared extractor is agnostic to - /// the user-defined context type the desugarer uses internally. - pub desugar: Option>, pub file_globs: Vec, } +impl LanguageExtractor for LanguageSpec { + fn file_globs(&self) -> &[String] { + &self.file_globs + } + + fn build_schema(&self) -> std::io::Result { + 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, @@ -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 = 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> = 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) -} diff --git a/shared/tree-sitter-extractor/tests/integration_test.rs b/shared/tree-sitter-extractor/tests/integration_test.rs index 694eb526f39..2b243ff7945 100644 --- a/shared/tree-sitter-extractor/tests/integration_test.rs +++ b/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()], }; diff --git a/shared/tree-sitter-extractor/tests/multiple_languages.rs b/shared/tree-sitter-extractor/tests/multiple_languages.rs index e345eec5828..2e45e56754a 100644 --- a/shared/tree-sitter-extractor/tests/multiple_languages.rs +++ b/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()], }; diff --git a/unified/extractor/src/extractor.rs b/unified/extractor/src/extractor.rs index 301c6cf533f..82bfe81219b 100644 --- a/unified/extractor/src/extractor.rs +++ b/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, diff --git a/unified/extractor/src/languages/mod.rs b/unified/extractor/src/languages/mod.rs index 52a1bd40ffc..a1f5c88cf3a 100644 --- a/unified/extractor/src/languages/mod.rs +++ b/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 { +pub fn all_language_specs() -> Vec { vec![swift::language_spec(OUTPUT_AST_SCHEMA)] } diff --git a/unified/extractor/src/languages/swift/swift.rs b/unified/extractor/src/languages/swift/swift.rs index 13f1a6beadf..991e95aaa6d 100644 --- a/unified/extractor/src/languages/swift/swift.rs +++ b/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> { ] } -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::::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), } } diff --git a/unified/extractor/tests/corpus_tests.rs b/unified/extractor/tests/corpus_tests.rs index ac93f06622c..d192b19485e 100644 --- a/unified/extractor/tests/corpus_tests.rs +++ b/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 { - 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 { + 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 { - 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 { + 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{}", diff --git a/unified/extractor/tests/swift_syntax_pipeline.rs b/unified/extractor/tests/swift_syntax_pipeline.rs index cdaae1f47c6..fde060f9820 100644 --- a/unified/extractor/tests/swift_syntax_pipeline.rs +++ b/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 =