tree_sitter_cli/
fuzz.rs

1use std::{
2    collections::HashMap,
3    env, fs,
4    path::{Path, PathBuf},
5    sync::LazyLock,
6};
7
8use log::{error, info};
9use rand::RngExt;
10use regex::Regex;
11use tree_sitter::{Language, Parser};
12
13pub mod allocations;
14pub mod corpus_test;
15pub mod edits;
16pub mod random;
17pub mod scope_sequence;
18
19use crate::{
20    fuzz::{
21        corpus_test::{
22            check_changed_ranges, check_consistent_sizes, get_parser, set_included_ranges,
23        },
24        edits::{get_random_edit, invert_edit},
25        random::Rand,
26    },
27    parse::perform_edit,
28    test::{DiffKey, TestDiff, TestEntry, TestExpectation, parse_tests, render_test_output},
29};
30
31pub static LOG_ENABLED: LazyLock<bool> = LazyLock::new(|| env::var("TREE_SITTER_LOG").is_ok());
32
33pub static LOG_GRAPH_ENABLED: LazyLock<bool> =
34    LazyLock::new(|| env::var("TREE_SITTER_LOG_GRAPHS").is_ok());
35
36pub static LANGUAGE_FILTER: LazyLock<Option<String>> =
37    LazyLock::new(|| env::var("TREE_SITTER_LANGUAGE").ok());
38
39pub static EXAMPLE_INCLUDE: LazyLock<Option<Regex>> =
40    LazyLock::new(|| regex_env_var("TREE_SITTER_EXAMPLE_INCLUDE"));
41
42pub static EXAMPLE_EXCLUDE: LazyLock<Option<Regex>> =
43    LazyLock::new(|| regex_env_var("TREE_SITTER_EXAMPLE_EXCLUDE"));
44
45pub static START_SEED: LazyLock<usize> = LazyLock::new(new_seed);
46
47pub const DEFAULT_EDIT_COUNT: usize = 3;
48pub static EDIT_COUNT: LazyLock<usize> =
49    LazyLock::new(|| int_env_var("TREE_SITTER_EDITS").unwrap_or(DEFAULT_EDIT_COUNT));
50
51pub const DEFAULT_ITERATION_COUNT: usize = 10;
52pub static ITERATION_COUNT: LazyLock<usize> =
53    LazyLock::new(|| int_env_var("TREE_SITTER_ITERATIONS").unwrap_or(DEFAULT_ITERATION_COUNT));
54
55fn int_env_var(name: &'static str) -> Option<usize> {
56    env::var(name).ok().and_then(|e| e.parse().ok())
57}
58
59fn regex_env_var(name: &'static str) -> Option<Regex> {
60    env::var(name).ok().and_then(|e| Regex::new(&e).ok())
61}
62
63#[must_use]
64pub fn new_seed() -> usize {
65    int_env_var("TREE_SITTER_SEED").unwrap_or_else(|| {
66        let mut rng = rand::rng();
67        let seed = rng.random_range(0..=usize::MAX);
68        eprintln!("fuzz seed: {seed}");
69        seed
70    })
71}
72
73pub struct FuzzOptions {
74    pub skipped: Option<Vec<String>>,
75    pub subdir: Option<PathBuf>,
76    pub edits: usize,
77    pub iterations: usize,
78    pub include: Option<Regex>,
79    pub exclude: Option<Regex>,
80    pub log_graphs: bool,
81    pub log: bool,
82}
83
84pub fn fuzz_language_corpus(
85    language: &Language,
86    language_name: &str,
87    start_seed: usize,
88    grammar_dir: &Path,
89    options: &mut FuzzOptions,
90) {
91    fn retain(entry: &mut TestEntry, language_name: &str) -> bool {
92        match entry {
93            TestEntry::Example { attributes, .. } => {
94                attributes.languages[0].is_empty()
95                    || attributes
96                        .languages
97                        .iter()
98                        .any(|lang| lang.as_ref() == language_name)
99            }
100            TestEntry::Group { children, .. } => {
101                children.retain_mut(|child| retain(child, language_name));
102                !children.is_empty()
103            }
104        }
105    }
106
107    let subdir = options.subdir.take().unwrap_or_default();
108
109    let corpus_dir = grammar_dir.join(subdir).join("test").join("corpus");
110
111    if !corpus_dir.exists() || !corpus_dir.is_dir() {
112        error!(
113            "No corpus directory found, ensure that you have a `test/corpus` directory in your grammar directory with at least one test file."
114        );
115        return;
116    }
117
118    if std::fs::read_dir(&corpus_dir).unwrap().count() == 0 {
119        error!(
120            "No corpus files found in `test/corpus`, ensure that you have at least one test file in your corpus directory."
121        );
122        return;
123    }
124
125    let mut main_tests = parse_tests(&corpus_dir).unwrap();
126    match main_tests {
127        TestEntry::Group {
128            ref mut children, ..
129        } => {
130            children.retain_mut(|child| retain(child, language_name));
131        }
132        TestEntry::Example { .. } => unreachable!(),
133    }
134    let tests = flatten_tests(
135        main_tests,
136        options.include.as_ref(),
137        options.exclude.as_ref(),
138    );
139
140    let get_test_name = |test: &FlattenedTest| format!("{language_name} - {}", test.name);
141
142    let mut skipped = options
143        .skipped
144        .take()
145        .unwrap_or_default()
146        .into_iter()
147        .chain(tests.iter().filter(|t| t.skip()).map(get_test_name))
148        .map(|x| (x, 0))
149        .collect::<HashMap<String, usize>>();
150
151    let mut failure_count = 0;
152
153    let log_seed = env::var("TREE_SITTER_LOG_SEED").is_ok();
154    let dump_edits = env::var("TREE_SITTER_DUMP_EDITS").is_ok();
155
156    if log_seed {
157        info!("  start seed: {start_seed}");
158    }
159
160    println!();
161    for (test_index, test) in tests.iter().enumerate() {
162        let test_name = get_test_name(test);
163        if let Some(counter) = skipped.get_mut(test_name.as_str()) {
164            println!("  {test_index}. {test_name} - SKIPPED");
165            *counter += 1;
166            continue;
167        }
168
169        println!("  {test_index}. {test_name}");
170
171        let passed = allocations::record_checked(|| {
172            let check_output = !test.error();
173            test.check_initial_parse(language, &test_name, check_output)
174        })
175        .unwrap_or_else(|e| {
176            error!("{e}");
177            false
178        });
179
180        if !passed {
181            failure_count += 1;
182            continue;
183        }
184
185        let mut parser = Parser::new();
186        parser.set_language(language).unwrap();
187        let tree = parser.parse(&test.input, None).unwrap();
188        drop(parser);
189
190        for trial in 0..options.iterations {
191            let seed = start_seed + trial;
192            let passed = allocations::record_checked(|| {
193                let mut rand = Rand::new(seed);
194                let mut log_session = None;
195                let mut parser = get_parser(&mut log_session, "log.html");
196                parser.set_language(language).unwrap();
197                let mut tree = tree.clone();
198                let mut input = test.input.clone();
199
200                if options.log_graphs {
201                    info!("{}\n", String::from_utf8_lossy(&input));
202                }
203
204                // Perform a random series of edits and reparse.
205                let edit_count = rand.unsigned(options.edits);
206                let mut undo_stack = Vec::with_capacity(edit_count);
207                for _ in 0..=edit_count {
208                    let edit = get_random_edit(&mut rand, &input);
209                    undo_stack.push(invert_edit(&input, &edit));
210                    perform_edit(&mut tree, &mut input, &edit).unwrap();
211                }
212
213                if log_seed {
214                    info!("   {test_index}.{trial:<2} seed: {seed}");
215                }
216
217                if dump_edits {
218                    fs::create_dir_all("fuzz").unwrap();
219                    fs::write(
220                        Path::new("fuzz")
221                            .join(format!("edit.{seed}.{test_index}.{trial} {test_name}")),
222                        &input,
223                    )
224                    .unwrap();
225                }
226
227                if options.log_graphs {
228                    info!("{}\n", String::from_utf8_lossy(&input));
229                }
230
231                set_included_ranges(&mut parser, &input, test.template_delimiters);
232                let mut tree2 = parser.parse(&input, Some(&tree)).unwrap();
233
234                // Check that the new tree is consistent.
235                check_consistent_sizes(&tree2, &input);
236                if let Err(message) = check_changed_ranges(&tree, &tree2, &input) {
237                    error!("\nUnexpected scope change in seed {seed} with start seed {start_seed}\n{message}\n\n");
238                    return false;
239                }
240
241                // Undo all of the edits and re-parse again.
242                while let Some(edit) = undo_stack.pop() {
243                    perform_edit(&mut tree2, &mut input, &edit).unwrap();
244                }
245                if options.log_graphs {
246                    info!("{}\n", String::from_utf8_lossy(&input));
247                }
248
249                set_included_ranges(&mut parser, &test.input, test.template_delimiters);
250                let tree3 = parser.parse(&input, Some(&tree2)).unwrap();
251
252                // Verify that the final tree matches the expectation from the corpus.
253                let actual_output = render_test_output(&input, &tree3, test.cst, test.has_fields).unwrap();
254
255                if actual_output != test.output && !test.error() {
256                    println!("Incorrect parse for {test_name} - seed {seed}");
257                    DiffKey::print();
258                    println!("{}", TestDiff::new(&actual_output, &test.output));
259                    println!();
260                    return false;
261                }
262
263                // Check that the edited tree is consistent.
264                check_consistent_sizes(&tree3, &input);
265                if let Err(message) = check_changed_ranges(&tree2, &tree3, &input) {
266                    error!("Unexpected scope change in seed {seed} with start seed {start_seed}\n{message}\n\n");
267                    return false;
268                }
269
270                true
271            }).unwrap_or_else(|e| {
272                error!("{e}");
273                false
274            });
275
276            if !passed {
277                failure_count += 1;
278                break;
279            }
280        }
281    }
282
283    if failure_count != 0 {
284        info!("{failure_count} {language_name} corpus tests failed fuzzing");
285    }
286
287    skipped.retain(|_, v| *v == 0);
288
289    if !skipped.is_empty() {
290        info!("Non matchable skip definitions:");
291        for k in skipped.keys() {
292            info!("  {k}");
293        }
294        panic!("Non matchable skip definitions need to be removed");
295    }
296}
297
298pub struct FlattenedTest {
299    pub name: String,
300    pub input: Vec<u8>,
301    pub output: String,
302    pub languages: Vec<Box<str>>,
303    pub expectation: TestExpectation,
304    pub has_fields: bool,
305    pub cst: bool,
306    pub template_delimiters: Option<(&'static str, &'static str)>,
307}
308
309impl FlattenedTest {
310    #[must_use]
311    fn skip(&self) -> bool {
312        self.expectation == TestExpectation::Skip
313    }
314
315    #[must_use]
316    fn error(&self) -> bool {
317        self.expectation == TestExpectation::Error
318    }
319
320    #[must_use]
321    pub(crate) fn check_initial_parse(
322        &self,
323        language: &Language,
324        display_name: &str,
325        check_output: bool,
326    ) -> bool {
327        let mut log_session = None;
328        let mut parser = get_parser(&mut log_session, "log.html");
329        parser.set_language(language).unwrap();
330        set_included_ranges(&mut parser, &self.input, self.template_delimiters);
331
332        let tree = parser.parse(&self.input, None).unwrap();
333
334        if !check_output {
335            return true;
336        }
337
338        let actual_output =
339            render_test_output(&self.input, &tree, self.cst, self.has_fields).unwrap();
340
341        if actual_output == self.output {
342            true
343        } else {
344            println!("Incorrect initial parse for {display_name}");
345            DiffKey::print();
346            println!("{}", TestDiff::new(&actual_output, &self.output));
347            println!();
348            false
349        }
350    }
351}
352
353#[must_use]
354pub fn flatten_tests(
355    test: TestEntry,
356    include: Option<&Regex>,
357    exclude: Option<&Regex>,
358) -> Vec<FlattenedTest> {
359    fn helper(
360        test: TestEntry,
361        include: Option<&Regex>,
362        exclude: Option<&Regex>,
363        is_root: bool,
364        prefix: &str,
365        result: &mut Vec<FlattenedTest>,
366    ) {
367        match test {
368            TestEntry::Example {
369                mut name,
370                input,
371                output,
372                has_fields,
373                attributes,
374                ..
375            } => {
376                if !prefix.is_empty() {
377                    name.insert_str(0, " - ");
378                    name.insert_str(0, prefix);
379                }
380
381                if let Some(include) = include {
382                    if !include.is_match(&name) {
383                        return;
384                    }
385                } else if let Some(exclude) = exclude
386                    && exclude.is_match(&name)
387                {
388                    return;
389                }
390
391                result.push(FlattenedTest {
392                    name,
393                    input,
394                    output,
395                    languages: attributes.languages,
396                    expectation: attributes.expectation,
397                    has_fields,
398                    cst: attributes.cst,
399                    template_delimiters: None,
400                });
401            }
402            TestEntry::Group {
403                mut name, children, ..
404            } => {
405                if !is_root && !prefix.is_empty() {
406                    name.insert_str(0, " - ");
407                    name.insert_str(0, prefix);
408                }
409                for child in children {
410                    helper(child, include, exclude, false, &name, result);
411                }
412            }
413        }
414    }
415    let mut result = Vec::new();
416    helper(test, include, exclude, true, "", &mut result);
417    result
418}
tree_sitter_cli/fuzz.rs

tree_sitter_cli/
fuzz.rs
