Unified: Fix a bunch of extractor failures

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

@@ -63,7 +63,7 @@ impl LanguageExtractor for LanguageSpec {
         trap_writer: &mut trap::Writer,
         path: &std::path::Path,
         source: &[u8],
-    ) {
+    ) -> Result<(), String> {
         crate::extractor::extract_parsed(
             self.parser.as_ref(),
             self.prefix,
@@ -74,7 +74,7 @@ impl LanguageExtractor for LanguageSpec {
             path,
             source,
             self.desugarer.as_ref(),
-        );
+        )
     }
 }
 

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

@@ -27,14 +27,17 @@ pub(crate) trait LanguageExtractor: Sync {
     /// 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`.
+    ///
+    /// A returned error is logged as a failure of this file only; the driver
+    /// archives its source, omits its TRAP, and continues.
     fn extract_file(
         &self,
         schema: &NodeTypeMap,
         diagnostics_writer: &mut diagnostics::LogWriter,
         trap_writer: &mut trap::Writer,
         path: &Path,
         source: &[u8],
-    );
+    ) -> Result<(), String>;
 }
 
 /// Drive extraction over `languages` for every file listed in `file_lists`.
@@ -171,7 +174,7 @@ pub(crate) fn run_extractor<L: LanguageExtractor>(
                             languages_processed[i] = true;
                             let lang = &languages[i];
 
-                            lang.extract_file(
+                            let result = lang.extract_file(
                                 &schemas[i],
                                 &mut diagnostics_writer,
                                 &mut trap_writer,
@@ -180,7 +183,18 @@ pub(crate) fn run_extractor<L: LanguageExtractor>(
                             );
                             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)?;
+                            match result {
+                                Ok(()) => {
+                                    write_trap(trap_dir, &path, &trap_writer, trap_compression)?;
+                                }
+                                Err(error) => {
+                                    tracing::error!(
+                                        file = %path.display(),
+                                        error,
+                                        "Failed to extract file"
+                                    );
+                                }
+                            }
                         }
                     }
                 }
@@ -208,3 +222,82 @@ fn write_trap(
     std::fs::create_dir_all(trap_file.parent().unwrap())?;
     trap_writer.write_to_file(&trap_file, trap_compression)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::Write;
+
+    struct TestLanguage {
+        file_globs: Vec<String>,
+    }
+
+    impl LanguageExtractor for TestLanguage {
+        fn file_globs(&self) -> &[String] {
+            &self.file_globs
+        }
+
+        fn build_schema(&self) -> std::io::Result<NodeTypeMap> {
+            Ok(NodeTypeMap::new())
+        }
+
+        fn extract_file(
+            &self,
+            _schema: &NodeTypeMap,
+            _diagnostics_writer: &mut diagnostics::LogWriter,
+            trap_writer: &mut trap::Writer,
+            _path: &Path,
+            source: &[u8],
+        ) -> Result<(), String> {
+            if source == b"bad" {
+                Err("invalid parser output".to_string())
+            } else {
+                trap_writer.comment("success".to_string());
+                Ok(())
+            }
+        }
+    }
+
+    #[test]
+    fn file_extraction_error_does_not_abort_other_files() {
+        let root = std::env::temp_dir().join(format!("codeql-extractor-{}", rand::random::<u64>()));
+        let source_dir = root.join("input");
+        let source_archive_dir = root.join("source-archive");
+        let trap_dir = root.join("trap");
+        std::fs::create_dir_all(&source_dir).unwrap();
+
+        let good_path = source_dir.join("good.test");
+        let bad_path = source_dir.join("bad.test");
+        std::fs::write(&good_path, b"good").unwrap();
+        std::fs::write(&bad_path, b"bad").unwrap();
+
+        let file_list = root.join("files.txt");
+        let mut file = std::fs::File::create(&file_list).unwrap();
+        writeln!(file, "{}", good_path.display()).unwrap();
+        writeln!(file, "{}", bad_path.display()).unwrap();
+
+        run_extractor(
+            "test",
+            &[TestLanguage {
+                file_globs: vec!["*.test".to_string()],
+            }],
+            &trap_dir,
+            &source_archive_dir,
+            &[file_list],
+            &Ok(trap::Compression::Gzip),
+        )
+        .unwrap();
+
+        let good_trap = file_paths::path_for(&trap_dir, &good_path, "trap.gz", None);
+        let bad_trap = file_paths::path_for(&trap_dir, &bad_path, "trap.gz", None);
+        assert!(good_trap.is_file());
+        assert!(!bad_trap.exists());
+
+        let archived_good = file_paths::path_for(&source_archive_dir, &good_path, "", None);
+        let archived_bad = file_paths::path_for(&source_archive_dir, &bad_path, "", None);
+        assert!(archived_good.is_file());
+        assert!(archived_bad.is_file());
+
+        std::fs::remove_dir_all(root).unwrap();
+    }
+}

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

@@ -419,7 +419,9 @@ fn collect_extras(node: Node<'_>, source: &[u8], out: &mut Vec<ExtraToken>) {
 /// 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.
+/// here; languages that don't desugar use [`extract`] instead. Parse and
+/// desugaring errors are returned to the multi-file driver so it can skip only
+/// this file and continue extracting the rest.
 #[allow(clippy::too_many_arguments)]
 pub fn extract_parsed(
     parse: &(dyn Fn(&[u8]) -> Result<ParsedTree, String> + Send + Sync),
@@ -431,7 +433,7 @@ pub fn extract_parsed(
     path: &Path,
     source: &[u8],
     desugarer: &dyn yeast::Desugarer,
-) {
+) -> Result<(), String> {
     let path_str = file_paths::normalize_and_transform_path(path, transformer);
     let source_root = std::env::current_dir()
         .ok()
@@ -441,6 +443,11 @@ pub fn extract_parsed(
     let _enter = span.enter();
     tracing::debug!("extracting: {}", path_str);
 
+    let parsed = parse(source).map_err(|e| format!("Parsing failed: {e}"))?;
+    let ast = desugarer
+        .run_from_ast(parsed.ast)
+        .map_err(|e| format!("Desugaring failed: {e}"))?;
+
     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(
@@ -453,16 +460,13 @@ pub fn extract_parsed(
         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);
     }
+    Ok(())
 }
 
 /// A lightweight [`AstNode`] over a piece of side-channel `extra` content

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

@@ -32,7 +32,7 @@ impl LanguageExtractor for LanguageSpec {
         trap_writer: &mut trap::Writer,
         path: &std::path::Path,
         source: &[u8],
-    ) {
+    ) -> Result<(), String> {
         crate::extractor::extract(
             &self.ts_language,
             self.prefix,
@@ -44,6 +44,7 @@ impl LanguageExtractor for LanguageSpec {
             source,
             &[],
         );
+        Ok(())
     }
 }
 

shared/yeast/src/lib.rs

@@ -1139,7 +1139,7 @@ impl<C> Rule<C> {
     }
 }
 
-const MAX_REWRITE_DEPTH: usize = 100;
+const MAX_REWRITE_DEPTH: usize = 1000;
 
 /// Index of rules by their root query kind for fast lookup.
 struct RuleIndex<'a, C> {

unified/extractor/Cargo.toml

@@ -14,7 +14,8 @@ rayon = "1.12.0"
 regex = "1.13.1"
 encoding = "0.2"
 lazy_static = "1.5.0"
-serde_json = "1.0.151"
+serde = "1.0.229"
+serde_json = { version = "1.0.151", features = ["unbounded_depth"] }
 
 codeql-extractor = { path = "../../shared/tree-sitter-extractor" }
 yeast = { path = "../../shared/yeast" }