1use std::cell::RefCell;
8
9use tree_sitter::Parser;
10
11use crate::error::Error;
12use crate::language::LangId;
13
14thread_local! {
15 static PARSERS: RefCell<[Option<Parser>; LangId::COUNT]> = const { RefCell::new([const { None }; LangId::COUNT]) };
16}
17
18fn new_parser(lang: LangId) -> Result<Parser, Error> {
19 let mut parser = Parser::new();
20 let grammar = crate::language::grammar_for(lang);
21 parser
22 .set_language(&grammar)
23 .map_err(|e| Error::Config(format!("failed to set language: {e}")))?;
24 Ok(parser)
25}
26
27fn get_or_init_parser(
28 parsers: &mut [Option<Parser>; LangId::COUNT],
29 lang: LangId,
30) -> Result<&mut Parser, Error> {
31 let idx = lang as usize;
32 if parsers[idx].is_none() {
33 parsers[idx] = Some(new_parser(lang)?);
34 }
35 parsers[idx]
36 .as_mut()
37 .ok_or_else(|| Error::Config("parser slot was not initialized".into()))
38}
39
40pub(crate) fn parse_tree(lang: LangId, source: &[u8]) -> Result<tree_sitter::Tree, Error> {
41 PARSERS.with(|cache| {
42 let Ok(mut pool) = cache.try_borrow_mut() else {
47 let mut parser = new_parser(lang)?;
48 return parse_with(&mut parser, lang, source);
49 };
50 let parser = get_or_init_parser(&mut pool, lang)?;
51 parse_with(parser, lang, source)
52 })
53}
54
55fn parse_with(
56 parser: &mut Parser,
57 lang: LangId,
58 source: &[u8],
59) -> Result<tree_sitter::Tree, Error> {
60 parser.reset();
61 parser.parse(source, None).ok_or_else(|| Error::Parse {
62 path: Default::default(),
63 message: format!("the {lang:?} parser returned no tree"),
64 })
65}
66
67#[derive(Debug, Clone, Copy, PartialEq)]
69pub struct TreeMetrics {
70 pub error_ratio: f64,
71 pub node_count: usize,
72}
73
74pub fn tree_metrics(tree: &tree_sitter::Tree, source: &[u8]) -> TreeMetrics {
76 if source.is_empty() {
77 return TreeMetrics {
78 error_ratio: 0.0,
79 node_count: 0,
80 };
81 }
82 let mut cursor = tree.walk();
83 let mut total = 0u32;
84 let mut errors = 0u32;
85 let mut named = 0u32;
86 let mut reached_root = false;
87
88 while !reached_root {
89 let n = cursor.node();
90 total += 1;
91 if n.is_named() {
92 named += 1;
93 }
94 if n.is_error() || n.is_missing() {
95 errors += 1;
96 }
97 if cursor.goto_first_child() {
98 continue;
99 }
100 if cursor.goto_next_sibling() {
101 continue;
102 }
103 let mut retracing = true;
104 while retracing {
105 if !cursor.goto_parent() {
106 reached_root = true;
107 break;
108 }
109 if cursor.goto_next_sibling() {
110 retracing = false;
111 }
112 }
113 }
114
115 let error_ratio = if total == 0 {
116 0.0
117 } else {
118 errors as f64 / total as f64
119 };
120
121 TreeMetrics {
122 error_ratio,
123 node_count: named as usize,
124 }
125}
126
127#[cfg(test)]
128mod tests {
129 use super::*;
130 use crate::language::LangId;
131
132 #[test]
133 fn parse_tree_valid_python() {
134 let tree = parse_tree(LangId::Python, b"def hello(): pass").unwrap();
135 assert!(!tree.root_node().has_error());
136 assert_eq!(tree.root_node().kind(), "module");
137 }
138
139 #[test]
140 fn parse_tree_switches_languages() {
141 let python = parse_tree(LangId::Python, b"def hello(): pass").unwrap();
142 assert_eq!(python.root_node().kind(), "module");
143
144 let javascript = parse_tree(LangId::JavaScript, b"function hello() {}").unwrap();
145 assert_eq!(javascript.root_node().kind(), "program");
146 }
147
148 #[test]
149 fn tree_metrics_valid_source() {
150 let tree = parse_tree(LangId::Python, b"def hello(): pass").unwrap();
151 let metrics = tree_metrics(&tree, b"def hello(): pass");
152 assert!(metrics.error_ratio < 0.1);
153 assert!(metrics.node_count > 0);
154 }
155
156 #[test]
157 fn tree_metrics_malformed() {
158 let tree = parse_tree(LangId::Python, b"def broken(").unwrap();
159 let metrics = tree_metrics(&tree, b"def broken(");
160 assert!(metrics.error_ratio > 0.0);
161 }
162
163 #[test]
164 fn tree_metrics_empty_source() {
165 let tree = parse_tree(LangId::Python, b"").unwrap();
166 let metrics = tree_metrics(&tree, b"");
167 assert_eq!(metrics.error_ratio, 0.0);
168 assert_eq!(metrics.node_count, 0);
169 }
170 #[test]
171 fn a_busy_pool_still_parses() {
172 PARSERS.with(|cache| {
175 let guard = cache.borrow_mut();
176 let tree = parse_tree(LangId::Python, b"def hello(): pass");
177 assert!(
178 tree.is_ok(),
179 "a busy pool must fall back to a fresh parser: {:?}",
180 tree.err()
181 );
182 drop(guard);
183 });
184 }
185
186 #[test]
187 fn a_parser_is_reused_after_a_parse() {
188 let first = parse_tree(LangId::Python, b"def first(): pass").unwrap();
189 let second = parse_tree(LangId::Python, b"def second(): pass").unwrap();
190 assert!(!first.root_node().has_error());
191 assert!(!second.root_node().has_error());
192 let third = parse_tree(LangId::JavaScript, b"function third() {}").unwrap();
193 assert_eq!(third.root_node().kind(), "program");
194 }
195}