feat(test): allow corpus tests to use cst outputs

This is already exposed for consumers via the CLI, and is a natural way
to express some test expectations over the sexp form.

Also clean up some repeated logic in the internal test code, and narrow
the cst rendering return type to `std::io::Result` rather than
`anyhow::Result`.

crates/cli/src/fuzz.rs

@@ -25,7 +25,7 @@ use crate::{
         random::Rand,
     },
     parse::perform_edit,
-    test::{DiffKey, TestDiff, TestEntry, TestExpectation, parse_tests, strip_sexp_fields},
+    test::{DiffKey, TestDiff, TestEntry, TestExpectation, parse_tests, render_test_output},
 };
 
 pub static LOG_ENABLED: LazyLock<bool> = LazyLock::new(|| env::var("TREE_SITTER_LOG").is_ok());
@@ -169,31 +169,8 @@ pub fn fuzz_language_corpus(
         println!("  {test_index}. {test_name}");
 
         let passed = allocations::record_checked(|| {
-            let mut log_session = None;
-            let mut parser = get_parser(&mut log_session, "log.html");
-            parser.set_language(language).unwrap();
-            set_included_ranges(&mut parser, &test.input, test.template_delimiters);
-
-            let tree = parser.parse(&test.input, None).unwrap();
-
-            if test.error() {
-                return true;
-            }
-
-            let mut actual_output = tree.root_node().to_sexp();
-            if !test.has_fields {
-                actual_output = strip_sexp_fields(&actual_output);
-            }
-
-            if actual_output != test.output {
-                println!("Incorrect initial parse for {test_name}");
-                DiffKey::print();
-                println!("{}", TestDiff::new(&actual_output, &test.output));
-                println!();
-                return false;
-            }
-
-            true
+            let check_output = !test.error();
+            test.check_initial_parse(language, &test_name, check_output)
         })
         .unwrap_or_else(|e| {
             error!("{e}");
@@ -273,10 +250,7 @@ pub fn fuzz_language_corpus(
                 let tree3 = parser.parse(&input, Some(&tree2)).unwrap();
 
                 // Verify that the final tree matches the expectation from the corpus.
-                let mut actual_output = tree3.root_node().to_sexp();
-                if !test.has_fields {
-                    actual_output = strip_sexp_fields(&actual_output);
-                }
+                let actual_output = render_test_output(&input, &tree3, test.cst, test.has_fields).unwrap();
 
                 if actual_output != test.output && !test.error() {
                     println!("Incorrect parse for {test_name} - seed {seed}");
@@ -328,8 +302,10 @@ pub struct FlattenedTest {
     pub languages: Vec<Box<str>>,
     pub expectation: TestExpectation,
     pub has_fields: bool,
+    pub cst: bool,
     pub template_delimiters: Option<(&'static str, &'static str)>,
 }
+
 impl FlattenedTest {
     #[must_use]
     fn skip(&self) -> bool {
@@ -340,6 +316,38 @@ impl FlattenedTest {
     fn error(&self) -> bool {
         self.expectation == TestExpectation::Error
     }
+
+    #[must_use]
+    pub(crate) fn check_initial_parse(
+        &self,
+        language: &Language,
+        display_name: &str,
+        check_output: bool,
+    ) -> bool {
+        let mut log_session = None;
+        let mut parser = get_parser(&mut log_session, "log.html");
+        parser.set_language(language).unwrap();
+        set_included_ranges(&mut parser, &self.input, self.template_delimiters);
+
+        let tree = parser.parse(&self.input, None).unwrap();
+
+        if !check_output {
+            return true;
+        }
+
+        let actual_output =
+            render_test_output(&self.input, &tree, self.cst, self.has_fields).unwrap();
+
+        if actual_output == self.output {
+            true
+        } else {
+            println!("Incorrect initial parse for {display_name}");
+            DiffKey::print();
+            println!("{}", TestDiff::new(&actual_output, &self.output));
+            println!();
+            false
+        }
+    }
 }
 
 #[must_use]
@@ -384,9 +392,10 @@ pub fn flatten_tests(
                     name,
                     input,
                     output,
-                    has_fields,
                     languages: attributes.languages,
                     expectation: attributes.expectation,
+                    has_fields,
+                    cst: attributes.cst,
                     template_delimiters: None,
                 });
             }

crates/cli/src/parse.rs

@@ -774,7 +774,7 @@ pub fn render_cst<'a, 'b: 'a>(
     cursor: &mut TreeCursor<'a>,
     opts: &ParseFileOptions,
     out: &mut impl Write,
-) -> Result<()> {
+) -> io::Result<()> {
     let lossy_source_code = String::from_utf8_lossy(source_code);
     let total_width = lossy_source_code
         .lines()
@@ -850,7 +850,7 @@ fn write_node_text(
     source: &str,
     color: Option<impl Into<Color> + Copy>,
     text_info: (usize, usize),
-) -> Result<()> {
+) -> io::Result<()> {
     let (total_width, indent_level) = text_info;
     let (quote, quote_color) = if is_named {
         ('`', opts.parse_theme.backtick)
@@ -990,7 +990,7 @@ fn cst_render_node(
     total_width: usize,
     indent_level: usize,
     in_error: bool,
-) -> Result<()> {
+) -> io::Result<()> {
     let node = cursor.node();
     let is_named = node.is_named();
     if !opts.no_ranges {

crates/cli/src/test.rs

@@ -965,11 +965,7 @@ fn run_tests(
                                 test_num: test_summary.test_num,
                             },
                         });
-                        let actual = if attributes.cst {
-                            render_test_cst(&input, &tree)?
-                        } else {
-                            tree.root_node().to_sexp()
-                        };
+                        let actual = render_test_output(&input, &tree, attributes.cst, true)?;
                         test_summary.parse_failures.push(TestFailure::new(
                             &name,
                             actual,
@@ -982,14 +978,12 @@ fn run_tests(
                         return Ok(false);
                     }
                 } else {
-                    let mut actual = if attributes.cst {
-                        render_test_cst(&input, &tree)?
-                    } else {
-                        tree.root_node().to_sexp()
-                    };
-                    if !(attributes.cst || opts.show_fields || has_fields) {
-                        actual = strip_sexp_fields(&actual);
-                    }
+                    let actual = render_test_output(
+                        &input,
+                        &tree,
+                        attributes.cst,
+                        opts.show_fields || has_fields,
+                    )?;
 
                     if actual == output {
                         test_summary.parse_results.add_case(TestResult {
@@ -1171,7 +1165,7 @@ fn run_tests(
 }
 
 /// Convenience wrapper to render a CST for a test entry.
-fn render_test_cst(input: &[u8], tree: &Tree) -> Result<String> {
+fn render_test_cst(input: &[u8], tree: &Tree) -> io::Result<String> {
     let mut rendered_cst: Vec<u8> = Vec::new();
     let mut cursor = tree.walk();
     let opts = ParseFileOptions {
@@ -1192,6 +1186,25 @@ fn render_test_cst(input: &[u8], tree: &Tree) -> Result<String> {
     Ok(String::from_utf8_lossy(&rendered_cst).trim().to_string())
 }
 
+/// Render a parsed tree in the output format expected by a corpus test.
+pub(crate) fn render_test_output(
+    input: &[u8],
+    tree: &Tree,
+    cst: bool,
+    include_fields: bool,
+) -> io::Result<String> {
+    if cst {
+        render_test_cst(input, tree)
+    } else {
+        let out = tree.root_node().to_sexp();
+        Ok(if include_fields {
+            out
+        } else {
+            strip_sexp_fields(&out)
+        })
+    }
+}
+
 // Parse time is interpreted in ns before converting to ms to avoid truncation issues
 // Parse rates often have several outliers, leading to a large standard deviation. Taking
 // the log of these rates serves to "flatten" out the distribution, yielding a more

crates/cli/src/tests/corpus_test.rs

@@ -17,7 +17,7 @@ use crate::{
         random::Rand,
     },
     parse::perform_edit,
-    test::{DiffKey, TestDiff, parse_tests, strip_sexp_fields},
+    test::{DiffKey, TestDiff, parse_tests, render_test_output},
     tests::{
         allocations,
         helpers::fixtures::{SCRATCH_BASE_DIR, fixtures_dir, get_language, get_test_language},
@@ -196,28 +196,7 @@ pub fn test_language_corpus(
 
         println!("  {test_index}. {test_name}");
 
-        let passed = allocations::record(|| {
-            let mut log_session = None;
-            let mut parser = get_parser(&mut log_session, "log.html");
-            parser.set_language(&language).unwrap();
-            set_included_ranges(&mut parser, &test.input, test.template_delimiters);
-
-            let tree = parser.parse(&test.input, None).unwrap();
-            let mut actual_output = tree.root_node().to_sexp();
-            if !test.has_fields {
-                actual_output = strip_sexp_fields(&actual_output);
-            }
-
-            if actual_output != test.output {
-                println!("Incorrect initial parse for {test_name}");
-                DiffKey::print();
-                println!("{}", TestDiff::new(&actual_output, &test.output));
-                println!();
-                return false;
-            }
-
-            true
-        });
+        let passed = allocations::record(|| test.check_initial_parse(&language, &test_name, true));
 
         if !passed {
             failure_count += 1;
@@ -293,10 +272,8 @@ pub fn test_language_corpus(
                 let tree3 = parser.parse(&input, Some(&tree2)).unwrap();
 
                 // Verify that the final tree matches the expectation from the corpus.
-                let mut actual_output = tree3.root_node().to_sexp();
-                if !test.has_fields {
-                    actual_output = strip_sexp_fields(&actual_output);
-                }
+                let actual_output =
+                    render_test_output(&input, &tree3, test.cst, test.has_fields).unwrap();
 
                 if actual_output != test.output {
                     println!("Incorrect parse for {test_name} - seed {seed}");
@@ -425,24 +402,8 @@ fn test_feature_corpus_files() {
             for test in tests {
                 eprintln!("  example: {:?}", test.name);
 
-                let passed = allocations::record(|| {
-                    let mut log_session = None;
-                    let mut parser = get_parser(&mut log_session, "log.html");
-                    parser.set_language(&language).unwrap();
-                    let tree = parser.parse(&test.input, None).unwrap();
-                    let mut actual_output = tree.root_node().to_sexp();
-                    if !test.has_fields {
-                        actual_output = strip_sexp_fields(&actual_output);
-                    }
-                    if actual_output == test.output {
-                        true
-                    } else {
-                        DiffKey::print();
-                        print!("{}", TestDiff::new(&actual_output, &test.output));
-                        println!();
-                        false
-                    }
-                });
+                let passed =
+                    allocations::record(|| test.check_initial_parse(&language, &test.name, true));
 
                 if !passed {
                     failure_count += 1;

test/fixtures/test_grammars/basic_cst/corpus.txt

@@ -0,0 +1,12 @@
+==================
+Basic CST rendering
+:cst
+==================
+answer = 42
+---
+
+0:0  - 0:11   document
+0:0  - 0:11     assignment
+0:0  - 0:6        name: identifier `answer`
+0:7  - 0:8        "="
+0:9  - 0:11       value: number `42`

test/fixtures/test_grammars/basic_cst/grammar.js

@@ -0,0 +1,16 @@
+export default grammar({
+  name: 'basic_cst',
+
+  rules: {
+    document: $ => $.assignment,
+
+    assignment: $ => seq(
+      field('name', $.identifier),
+      '=',
+      field('value', $.number),
+    ),
+
+    identifier: _ => /[a-z]+/,
+    number: _ => /[0-9]+/,
+  },
+});