1use std::ops::ControlFlow;
2
3use thiserror::Error;
4use tree_sitter::{Language, ParseOptions, Parser};
5
6use crate::{ExtractionLimitExceeded, ExtractionTracker};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum SourceSyntaxLanguage {
11 JavaScript,
13 TypeScript,
15 Rust,
17 Python,
19 Go,
21 Java,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub struct SourceSyntaxInspection {
28 pub has_error: bool,
30 pub boundary_candidate_count: usize,
32}
33
34#[derive(Debug, Error)]
36pub enum SourceSyntaxError {
37 #[error("failed to configure Tree-sitter grammar: {0}")]
39 Grammar(#[from] tree_sitter::LanguageError),
40 #[error("Tree-sitter did not return a syntax tree")]
42 MissingTree,
43 #[error(transparent)]
45 LimitExceeded(#[from] ExtractionLimitExceeded),
46}
47
48pub fn inspect_source_syntax(
57 language: SourceSyntaxLanguage,
58 source_path: &str,
59 source: &str,
60 tracker: &mut ExtractionTracker,
61) -> Result<SourceSyntaxInspection, SourceSyntaxError> {
62 tracker.check_input_bytes(u64::try_from(source.len()).unwrap_or(u64::MAX))?;
63 tracker.check_tree_sitter_time()?;
64 let grammar = grammar(language, source_path);
65 let mut parser = Parser::new();
66 parser.set_language(&grammar)?;
67 let mut timeout = None;
68 let tree = {
69 let mut progress = |_: &tree_sitter::ParseState| match tracker.check_tree_sitter_time() {
70 Ok(()) => ControlFlow::Continue(()),
71 Err(error) => {
72 timeout = Some(error);
73 ControlFlow::Break(())
74 }
75 };
76 let options = ParseOptions::new().progress_callback(&mut progress);
77 let bytes = source.as_bytes();
78 parser.parse_with_options(
79 &mut |offset, _| bytes.get(offset..).unwrap_or_default(),
80 None,
81 Some(options),
82 )
83 };
84 if let Some(error) = timeout {
85 return Err(error.into());
86 }
87 let tree = tree.ok_or(SourceSyntaxError::MissingTree)?;
88 let root = tree.root_node();
89 let mut boundary_candidate_count = 0;
90 count_candidates(language, root, &mut boundary_candidate_count, tracker)?;
91 tracker.check_tree_sitter_time()?;
92 Ok(SourceSyntaxInspection {
93 has_error: root.has_error(),
94 boundary_candidate_count,
95 })
96}
97
98fn grammar(language: SourceSyntaxLanguage, source_path: &str) -> Language {
99 match language {
100 SourceSyntaxLanguage::JavaScript => tree_sitter_javascript::LANGUAGE.into(),
101 SourceSyntaxLanguage::TypeScript
102 if std::path::Path::new(source_path)
103 .extension()
104 .is_some_and(|extension| extension.eq_ignore_ascii_case("tsx")) =>
105 {
106 tree_sitter_typescript::LANGUAGE_TSX.into()
107 }
108 SourceSyntaxLanguage::TypeScript => tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
109 SourceSyntaxLanguage::Rust => tree_sitter_rust::LANGUAGE.into(),
110 SourceSyntaxLanguage::Python => tree_sitter_python::LANGUAGE.into(),
111 SourceSyntaxLanguage::Go => tree_sitter_go::LANGUAGE.into(),
112 SourceSyntaxLanguage::Java => tree_sitter_java::LANGUAGE.into(),
113 }
114}
115
116fn count_candidates(
117 language: SourceSyntaxLanguage,
118 root: tree_sitter::Node<'_>,
119 count: &mut usize,
120 tracker: &mut ExtractionTracker,
121) -> Result<(), ExtractionLimitExceeded> {
122 tracker.charge_tree_sitter_node(0)?;
123 let mut pending = vec![(root, 0_u64)];
124 while let Some((node, depth)) = pending.pop() {
125 if candidate_kind(language, node.kind()) {
126 *count = count.saturating_add(1);
127 }
128 let child_depth = depth.saturating_add(1);
129 let mut cursor = node.walk();
130 for child in node.children(&mut cursor) {
131 tracker.charge_tree_sitter_node(child_depth)?;
132 pending.push((child, child_depth));
133 }
134 }
135 Ok(())
136}
137
138fn candidate_kind(language: SourceSyntaxLanguage, kind: &str) -> bool {
139 match language {
140 SourceSyntaxLanguage::JavaScript | SourceSyntaxLanguage::TypeScript => {
141 matches!(
142 kind,
143 "call_expression"
144 | "decorator"
145 | "method_definition"
146 | "function_declaration"
147 | "lexical_declaration"
148 )
149 }
150 SourceSyntaxLanguage::Rust => matches!(kind, "call_expression" | "attribute_item"),
151 SourceSyntaxLanguage::Python => {
152 matches!(
153 kind,
154 "call" | "decorator" | "dictionary" | "class_definition"
155 )
156 }
157 SourceSyntaxLanguage::Go => kind == "call_expression",
158 SourceSyntaxLanguage::Java => {
159 matches!(
160 kind,
161 "method_invocation" | "annotation" | "marker_annotation"
162 )
163 }
164 }
165}
166
167#[cfg(test)]
168mod tests {
169 use super::{SourceSyntaxError, SourceSyntaxLanguage, inspect_source_syntax};
170 use crate::{ExtractionBudgets, ExtractionResource, ExtractionTracker};
171
172 #[test]
173 fn mandatory_grammars_should_find_boundary_candidates() {
174 let fixtures = [
175 (
176 SourceSyntaxLanguage::JavaScript,
177 "client.js",
178 "fetch('/orders');",
179 ),
180 (
181 SourceSyntaxLanguage::TypeScript,
182 "router.ts",
183 "app.post('/orders', handler);",
184 ),
185 (
186 SourceSyntaxLanguage::JavaScript,
187 "src/app/api/orders/route.js",
188 "export async function POST() {}",
189 ),
190 (
191 SourceSyntaxLanguage::Rust,
192 "lib.rs",
193 "#[tokio::test]\nasync fn test_order() {}",
194 ),
195 (
196 SourceSyntaxLanguage::Python,
197 "test_api.py",
198 "@app.post('/orders')\ndef create(): pass",
199 ),
200 (
201 SourceSyntaxLanguage::Python,
202 "methods_config.py",
203 "METHODS = {'get_order': ['GET', '/orders/{id}']}",
204 ),
205 (
206 SourceSyntaxLanguage::Python,
207 "factories.py",
208 "class OrderFactory(factory.Factory):\n class Meta:\n model = Order",
209 ),
210 (
211 SourceSyntaxLanguage::Go,
212 "main.go",
213 "package main\nfunc main() { http.Get(\"/orders\") }",
214 ),
215 (
216 SourceSyntaxLanguage::Java,
217 "Client.java",
218 "class Client { void run() { client.get(); } }",
219 ),
220 ];
221
222 for (language, path, source) in fixtures {
223 let mut tracker =
224 ExtractionTracker::new(path, "tree-sitter", &ExtractionBudgets::default());
225 let result = inspect_source_syntax(language, path, source, &mut tracker);
226 assert!(
227 matches!(result, Ok(report) if report.boundary_candidate_count > 0),
228 "{path}: {result:?}"
229 );
230 }
231 }
232
233 #[test]
234 fn javascript_ast_depth_should_accept_256_and_reject_257() {
235 fn nested_arrays(count: usize) -> String {
236 format!("{}0{};", "[".repeat(count), "]".repeat(count))
237 }
238
239 let budgets = ExtractionBudgets {
240 max_ast_depth_per_artifact: 256,
241 ..ExtractionBudgets::default()
242 };
243 let mut exact = ExtractionTracker::new("exact.js", "tree-sitter", &budgets);
244 let mut above = ExtractionTracker::new("above.js", "tree-sitter", &budgets);
245
246 assert!(
247 inspect_source_syntax(
248 SourceSyntaxLanguage::JavaScript,
249 "exact.js",
250 &nested_arrays(254),
251 &mut exact,
252 )
253 .is_ok()
254 );
255 assert!(matches!(
256 inspect_source_syntax(
257 SourceSyntaxLanguage::JavaScript,
258 "above.js",
259 &nested_arrays(255),
260 &mut above,
261 ),
262 Err(SourceSyntaxError::LimitExceeded(error))
263 if error.resource == ExtractionResource::AstDepth
264 && error.observed == 257
265 && error.maximum == 256
266 ));
267 }
268
269 #[test]
270 fn javascript_node_budget_should_accept_500000_and_reject_500001() {
271 let budgets = ExtractionBudgets {
272 max_tree_sitter_nodes_per_artifact: 500_000,
273 ..ExtractionBudgets::default()
274 };
275 let mut exact = ExtractionTracker::new("exact.js", "tree-sitter", &budgets);
276 let mut above = ExtractionTracker::new("above.js", "tree-sitter", &budgets);
277 let exact_source = format!("{}// one", ";".repeat(249_999));
278 let above_source = ";".repeat(250_000);
279
280 assert!(
281 inspect_source_syntax(
282 SourceSyntaxLanguage::JavaScript,
283 "exact.js",
284 &exact_source,
285 &mut exact,
286 )
287 .is_ok()
288 );
289 assert!(matches!(
290 inspect_source_syntax(
291 SourceSyntaxLanguage::JavaScript,
292 "above.js",
293 &above_source,
294 &mut above,
295 ),
296 Err(SourceSyntaxError::LimitExceeded(error))
297 if error.resource == ExtractionResource::TreeSitterNodes
298 && error.observed == 500_001
299 && error.maximum == 500_000
300 ));
301 }
302
303 #[test]
304 fn parser_should_cancel_on_a_real_wall_time_deadline() {
305 let budgets = ExtractionBudgets {
306 max_tree_sitter_wall_time_ms_per_artifact: 1,
307 max_tree_sitter_nodes_per_artifact: 1_000_000,
308 ..ExtractionBudgets::default()
309 };
310 let mut tracker = ExtractionTracker::new("slow.js", "tree-sitter", &budgets);
311 let source = ";".repeat(750_000);
312
313 assert!(matches!(
314 inspect_source_syntax(
315 SourceSyntaxLanguage::JavaScript,
316 "slow.js",
317 &source,
318 &mut tracker,
319 ),
320 Err(SourceSyntaxError::LimitExceeded(error))
321 if error.resource == ExtractionResource::TreeSitterWallTimeMs
322 ));
323 }
324}