Unified: Add schema checking and corpus-style tests

shared/yeast/src/lib.rs

@@ -193,10 +193,43 @@ impl Ast {
         AstCursor::new(self)
     }
 
+    /// Return all nodes currently allocated in the AST arena.
+    ///
+    /// This includes nodes that are no longer reachable from `get_root()`
+    /// after desugaring rewrites. Use `reachable_node_ids()` for output-level
+    /// validation/traversal semantics.
     pub fn nodes(&self) -> &[Node] {
         &self.nodes
     }
 
+    /// Return node ids reachable from `get_root()` by following child edges.
+    ///
+    /// This reflects the effective AST after desugaring and excludes orphaned
+    /// arena nodes left behind by rewrite operations.
+    pub fn reachable_node_ids(&self) -> Vec<usize> {
+        let mut reachable = Vec::new();
+        let mut stack = vec![self.root];
+        let mut seen = vec![false; self.nodes.len()];
+
+        while let Some(id) = stack.pop() {
+            if id >= self.nodes.len() || seen[id] {
+                continue;
+            }
+            seen[id] = true;
+            reachable.push(id);
+
+            if let Some(node) = self.get_node(id) {
+                for children in node.fields.values() {
+                    for &child in children {
+                        stack.push(child);
+                    }
+                }
+            }
+        }
+
+        reachable
+    }
+
     pub fn get_root(&self) -> Id {
         self.root
     }

shared/yeast/tests/test.rs

@@ -166,6 +166,28 @@ fn test_query_no_match() {
     assert!(!matched);
 }
 
+#[test]
+fn test_reachable_nodes_excludes_orphaned_rewrite_nodes() {
+    let lang: tree_sitter::Language = tree_sitter_ruby::LANGUAGE.into();
+    let schema = yeast::node_types_yaml::schema_from_yaml_with_language(OUTPUT_SCHEMA_YAML, &lang)
+        .unwrap();
+    let rules = vec![yeast::rule!((integer) => (identifier "replaced"))];
+    let runner = Runner::with_schema(lang, &schema, &rules);
+
+    let input = "x = 1";
+    let ast = runner.run(input).unwrap();
+    let reachable_ids = ast.reachable_node_ids();
+
+    assert!(
+        ast.nodes().len() > reachable_ids.len(),
+        "expected rewrite to leave orphaned arena nodes"
+    );
+
+    let dump = dump_ast(&ast, ast.get_root(), input);
+    assert!(dump.contains("identifier \"replaced\""));
+    assert!(!dump.contains("integer \"1\""));
+}
+
 #[test]
 fn test_query_repeated_capture() {
     let runner = Runner::new(tree_sitter_ruby::LANGUAGE.into(), &[]);

unified/AGENTS.md

@@ -20,10 +20,15 @@ grammar source), run `scripts/regenerate-grammar.sh` to:
 it shows the impact of a grammar tweak on the named node kinds, fields,
 and child types in a form much easier to read than the raw JSON.
 
-## Testing
-- If you changed the extractor code, always rebuild it before running tests.
+## Extractor Testing
+- To run extractor tests, run `cargo test` in the `extractor` directory.
 
-- To run all tests, run `codeql test run --search-path extractor-pack ql/test`
+- Do not edit the printed ASTs in `extractor/test/corpus` directly. To regenerate the ASTs, run `scripts/update-corpus.sh`.
+
+## CodeQL Testing
+- If you changed the extractor code, always rebuild it before running CodeQL tests.
+
+- To run all CodeQL tests, run `codeql test run --search-path extractor-pack ql/test`
 
 - Do not edit `.expected` files manually. To update the expected output, pass `--learn` to the `codeql test run` command.
 

unified/extractor/src/extractor.rs

@@ -3,9 +3,7 @@ use std::path::PathBuf;
 
 use codeql_extractor::extractor::simple;
 use codeql_extractor::trap;
-
-#[path = "languages/swift/swift.rs"]
-mod swift;
+use crate::languages;
 
 #[derive(Args)]
 pub struct Options {
@@ -27,9 +25,7 @@ pub fn run(options: Options) -> std::io::Result<()> {
 
     let extractor = simple::Extractor {
         prefix: "unified".to_string(),
-        languages: vec![
-            swift::language_spec(),
-        ],
+        languages: languages::all_language_specs(),
         trap_dir: options.output_dir,
         trap_compression: trap::Compression::from_env("CODEQL_EXTRACTOR_UNIFIED_OPTION_TRAP_COMPRESSION"),
         source_archive_dir: options.source_archive_dir,

unified/extractor/src/languages/mod.rs

@@ -0,0 +1,8 @@
+use codeql_extractor::extractor::simple;
+
+#[path = "swift/swift.rs"]
+mod swift;
+
+pub fn all_language_specs() -> Vec<simple::LanguageSpec> {
+    vec![swift::language_spec()]
+}

unified/extractor/src/main.rs

@@ -3,6 +3,7 @@ use clap::Parser;
 mod autobuilder;
 mod extractor;
 mod generator;
+mod languages;
 
 #[derive(Parser)]
 #[command(author, version, about)]

unified/extractor/tests/corpus/swift/desugar.txt

@@ -0,0 +1,23 @@
+===
+Additive expression is desugared
+===
+
+1 + 2
+
+---
+
+source_file
+  simple_identifier "blah"
+
+
+===
+Another additive expression is desugared
+===
+
+foo + bar
+
+---
+
+source_file
+  simple_identifier "blah"
+

unified/extractor/tests/corpus_tests.rs

@@ -0,0 +1,182 @@
+use std::fs;
+use std::path::Path;
+
+use codeql_extractor::extractor::simple;
+use yeast::{dump::dump_ast, Runner};
+
+#[path = "../src/languages/mod.rs"]
+mod languages;
+
+#[derive(Debug)]
+struct CorpusCase {
+    name: String,
+    input: String,
+    expected: String,
+}
+
+fn update_mode_enabled() -> bool {
+    std::env::var("UNIFIED_UPDATE_CORPUS")
+        .map(|v| matches!(v.to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "on"))
+        .unwrap_or(false)
+}
+
+fn is_header_rule(line: &str) -> bool {
+    let trimmed = line.trim();
+    trimmed.len() >= 3 && trimmed.chars().all(|c| c == '=')
+}
+
+fn parse_corpus(content: &str) -> Vec<CorpusCase> {
+    let lines: Vec<&str> = content.lines().collect();
+    let mut i = 0;
+    let mut cases = Vec::new();
+
+    while i < lines.len() {
+        while i < lines.len() && lines[i].trim().is_empty() {
+            i += 1;
+        }
+        if i >= lines.len() {
+            break;
+        }
+
+        assert!(
+            is_header_rule(lines[i]),
+            "Expected header delimiter at line {}",
+            i + 1
+        );
+        i += 1;
+
+        assert!(i < lines.len(), "Missing test name at line {}", i + 1);
+        let name = lines[i].trim().to_string();
+        i += 1;
+
+        assert!(
+            i < lines.len() && is_header_rule(lines[i]),
+            "Missing closing header delimiter for case {name}"
+        );
+        i += 1;
+
+        let input_start = i;
+        while i < lines.len() && lines[i].trim() != "---" {
+            i += 1;
+        }
+        assert!(i < lines.len(), "Missing --- separator for case {name}");
+        let input = lines[input_start..i].join("\n").trim_end().to_string();
+        i += 1;
+
+        let expected_start = i;
+        while i < lines.len() {
+            if is_header_rule(lines[i])
+                && i + 2 < lines.len()
+                && !lines[i + 1].trim().is_empty()
+                && is_header_rule(lines[i + 2])
+            {
+                break;
+            }
+            i += 1;
+        }
+        let expected = lines[expected_start..i].join("\n").trim().to_string();
+
+        cases.push(CorpusCase {
+            name,
+            input,
+            expected,
+        });
+    }
+
+    cases
+}
+
+fn render_corpus(cases: &[CorpusCase]) -> String {
+    let mut out = String::new();
+
+    for (idx, case) in cases.iter().enumerate() {
+        if idx > 0 {
+            out.push('\n');
+        }
+        out.push_str("===\n");
+        out.push_str(case.name.trim());
+        out.push_str("\n===\n");
+        out.push('\n');
+        out.push_str(case.input.trim());
+        out.push_str("\n\n---\n");
+        out.push('\n');
+        out.push_str(case.expected.trim());
+        out.push_str("\n\n");
+    }
+
+    out
+}
+
+fn run_desugaring(lang: &simple::LanguageSpec, input: &str) -> String {
+    let runner = match lang.desugar.as_ref() {
+        Some(config) => Runner::from_config(lang.ts_language.clone(), config)
+            .expect("Failed to create yeast runner from desugaring config"),
+        None => Runner::new(lang.ts_language.clone(), &[]),
+    };
+    let ast = runner
+        .run(input)
+        .unwrap_or_else(|e| panic!("Failed to parse corpus input: {e}"));
+    dump_ast(&ast, ast.get_root(), input)
+}
+
+#[test]
+fn test_corpus() {
+    let update_mode = update_mode_enabled();
+    let all_languages = languages::all_language_specs();
+    let corpus_dir = Path::new("tests/corpus");
+
+    for lang in all_languages {
+        let lang_corpus_dir = corpus_dir.join(&lang.prefix);
+        if !lang_corpus_dir.exists() {
+            continue;
+        }
+
+        let mut corpus_files: Vec<_> = fs::read_dir(&lang_corpus_dir)
+            .unwrap_or_else(|e| {
+                panic!(
+                    "Failed to read corpus directory {}: {e}",
+                    lang_corpus_dir.display()
+                )
+            })
+            .map(|entry| entry.expect("Failed to read corpus entry").path())
+            .filter(|path| path.extension().is_some_and(|ext| ext == "txt"))
+            .collect();
+        corpus_files.sort();
+
+        for corpus_path in corpus_files {
+            let content = fs::read_to_string(&corpus_path)
+                .unwrap_or_else(|e| panic!("Failed to read {}: {e}", corpus_path.display()));
+            let mut cases = parse_corpus(&content);
+            assert!(
+                !cases.is_empty(),
+                "No corpus cases found in {}",
+                corpus_path.display()
+            );
+
+            for case in &mut cases {
+                let actual = run_desugaring(&lang, &case.input);
+                if update_mode {
+                    case.expected = actual.trim().to_string();
+                } else {
+                    assert_eq!(
+                        case.expected.trim(),
+                        actual.trim(),
+                        "Corpus case failed in {}: {}",
+                        corpus_path.display(),
+                        case.name
+                    );
+                }
+            }
+
+            if update_mode {
+                let updated = render_corpus(&cases);
+                fs::write(&corpus_path, updated).unwrap_or_else(|e| {
+                    panic!(
+                        "Failed to update corpus file {}: {e}",
+                        corpus_path.display()
+                    )
+                });
+            }
+        }
+    }
+}

unified/scripts/update-corpus.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+set -euo pipefail
+IFS=$'\n\t'
+
+cd "$(dirname "$0")/.."
+
+cd extractor
+UNIFIED_UPDATE_CORPUS=1 cargo test