Skip to main content

code_system_graph_core/
source_syntax.rs

1use std::ops::ControlFlow;
2
3use thiserror::Error;
4use tree_sitter::{Language, ParseOptions, Parser};
5
6use crate::{ExtractionLimitExceeded, ExtractionTracker};
7
8/// Source grammar selected for focused boundary inspection.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum SourceSyntaxLanguage {
11    /// JavaScript or JSX.
12    JavaScript,
13    /// TypeScript or TSX.
14    TypeScript,
15    /// Rust.
16    Rust,
17    /// Python.
18    Python,
19    /// Go.
20    Go,
21    /// Java.
22    Java,
23}
24
25/// Bounded syntax inspection used to corroborate focused textual facts.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub struct SourceSyntaxInspection {
28    /// Whether Tree-sitter recovered from at least one syntax error.
29    pub has_error: bool,
30    /// Number of syntax nodes that can contain a supported boundary.
31    pub boundary_candidate_count: usize,
32}
33
34/// Error produced when a mandatory Tree-sitter grammar cannot inspect an input.
35#[derive(Debug, Error)]
36pub enum SourceSyntaxError {
37    /// The selected grammar could not be configured.
38    #[error("failed to configure Tree-sitter grammar: {0}")]
39    Grammar(#[from] tree_sitter::LanguageError),
40    /// Tree-sitter did not return a syntax tree.
41    #[error("Tree-sitter did not return a syntax tree")]
42    MissingTree,
43    /// Parsing or traversal exceeded one configured invocation resource.
44    #[error(transparent)]
45    LimitExceeded(#[from] ExtractionLimitExceeded),
46}
47
48/// Inspects a focused source file with its mandatory Tree-sitter grammar.
49///
50/// The resulting candidate count is used as a structural guard for the narrower
51/// framework recognizers; it does not create or persist a repository-local AST.
52///
53/// # Errors
54///
55/// Returns [`SourceSyntaxError`] if the grammar cannot be selected or parsing fails.
56pub 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}