context_builder/tree_sitter/
mod.rs

1//! Tree-sitter integration for intelligent code parsing.
2//!
3//! This module provides:
4//! - Signature extraction (function/class signatures without bodies)
5//! - Smart truncation (truncate at AST boundaries)
6//! - Structure extraction (imports, exports, symbol counts)
7//!
8//! Feature-gated: Only compiled when one of the tree-sitter-* features is enabled.
9
10#[cfg(feature = "tree-sitter-base")]
11pub mod language_support;
12
13#[cfg(feature = "tree-sitter-base")]
14pub mod signatures;
15
16#[cfg(feature = "tree-sitter-base")]
17pub mod structure;
18
19#[cfg(feature = "tree-sitter-base")]
20pub mod truncation;
21
22#[cfg(feature = "tree-sitter-base")]
23pub mod languages;
24
25#[cfg(feature = "tree-sitter-base")]
26use std::path::Path;
27
28#[cfg(feature = "tree-sitter-base")]
29pub use language_support::{CodeStructure, LanguageSupport, Signature, SignatureKind, Visibility};
30
31#[cfg(feature = "tree-sitter-base")]
32pub use signatures::extract_signatures;
33
34#[cfg(feature = "tree-sitter-base")]
35pub use structure::extract_structure;
36
37#[cfg(feature = "tree-sitter-base")]
38pub use truncation::find_truncation_point;
39
40/// Check if tree-sitter is available for a given file extension.
41#[cfg(feature = "tree-sitter-base")]
42pub fn is_supported_extension(ext: &str) -> bool {
43    languages::get_language_support(ext).is_some()
44}
45
46#[cfg(not(feature = "tree-sitter-base"))]
47pub fn is_supported_extension(_ext: &str) -> bool {
48    false
49}
50
51/// Extract file extension from a path.
52#[cfg(feature = "tree-sitter-base")]
53fn get_extension(path: &Path) -> Option<String> {
54    path.extension()
55        .and_then(|e| e.to_str())
56        .map(|s| s.to_lowercase())
57}
58
59/// Get language support for a file path.
60#[cfg(feature = "tree-sitter-base")]
61pub fn get_language_for_path(path: &Path) -> Option<&'static dyn LanguageSupport> {
62    let ext = get_extension(path)?;
63    languages::get_language_support(&ext)
64}
65
66/// Extract signatures from source code for a given file extension.
67#[cfg(feature = "tree-sitter-base")]
68pub fn extract_signatures_for_file(
69    source: &str,
70    ext: &str,
71    visibility_filter: Visibility,
72) -> Option<Vec<Signature>> {
73    let support = languages::get_language_support(ext)?;
74    Some(extract_signatures(source, support, visibility_filter))
75}
76
77/// Extract structure from source code for a given file extension.
78#[cfg(feature = "tree-sitter-base")]
79pub fn extract_structure_for_file(source: &str, ext: &str) -> Option<CodeStructure> {
80    let support = languages::get_language_support(ext)?;
81    Some(extract_structure(source, support))
82}
83
84/// Find a smart truncation point for a given file extension.
85#[cfg(feature = "tree-sitter-base")]
86pub fn find_smart_truncation_point(source: &str, max_bytes: usize, ext: &str) -> Option<usize> {
87    let support = languages::get_language_support(ext)?;
88    Some(find_truncation_point(source, max_bytes, support))
89}
90
91#[cfg(not(feature = "tree-sitter-base"))]
92pub fn extract_signatures_for_file(
93    _source: &str,
94    _ext: &str,
95    _visibility_filter: (),
96) -> Option<()> {
97    None
98}
99
100#[cfg(not(feature = "tree-sitter-base"))]
101pub fn extract_structure_for_file(_source: &str, _ext: &str) -> Option<()> {
102    None
103}
104
105#[cfg(not(feature = "tree-sitter-base"))]
106pub fn find_smart_truncation_point(_source: &str, _max_bytes: usize, _ext: &str) -> Option<usize> {
107    None
108}
109
110#[cfg(not(feature = "tree-sitter-base"))]
111pub fn get_language_for_path(_path: &std::path::Path) -> Option<()> {
112    None
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[test]
120    #[cfg(feature = "tree-sitter-base")]
121    fn test_is_supported_extension() {
122        #[cfg(feature = "tree-sitter-rust")]
123        assert!(is_supported_extension("rs"));
124        #[cfg(feature = "tree-sitter-python")]
125        assert!(is_supported_extension("py"));
126        #[cfg(feature = "tree-sitter-js")]
127        assert!(is_supported_extension("js"));
128        assert!(!is_supported_extension("xyz"));
129    }
130
131    #[test]
132    #[cfg(not(feature = "tree-sitter-base"))]
133    fn test_no_tree_sitter_support() {
134        assert!(!is_supported_extension("rs"));
135        assert!(!is_supported_extension("py"));
136    }
137
138    #[test]
139    #[cfg(feature = "tree-sitter-base")]
140    fn test_get_extension() {
141        assert_eq!(get_extension(Path::new("foo.rs")), Some("rs".to_string()));
142        assert_eq!(get_extension(Path::new("foo.RS")), Some("rs".to_string()));
143        assert_eq!(get_extension(Path::new("foo.PY")), Some("py".to_string()));
144        assert_eq!(get_extension(Path::new("foo")), None);
145        assert_eq!(get_extension(Path::new(".gitignore")), None);
146    }
147
148    #[test]
149    #[cfg(feature = "tree-sitter-rust")]
150    fn test_extract_signatures_for_file_rust() {
151        let source = "pub fn hello() { }\nfn world() { }";
152        let sigs = extract_signatures_for_file(source, "rs", Visibility::All);
153        assert!(sigs.is_some());
154        let sigs = sigs.unwrap();
155        assert!(sigs.len() >= 2);
156    }
157
158    #[test]
159    #[cfg(feature = "tree-sitter-base")]
160    fn test_extract_signatures_for_file_unsupported() {
161        let sigs = extract_signatures_for_file("anything", "xyz", Visibility::All);
162        assert!(sigs.is_none());
163    }
164
165    #[test]
166    #[cfg(feature = "tree-sitter-rust")]
167    fn test_extract_structure_for_file_rust() {
168        let source = "use std::io;\nfn foo() { }\nstruct Bar { }\nenum Baz { A, B }";
169        let structure = extract_structure_for_file(source, "rs");
170        assert!(structure.is_some());
171        let s = structure.unwrap();
172        assert!(s.functions >= 1);
173        assert!(s.structs >= 1);
174        assert!(s.enums >= 1);
175    }
176
177    #[test]
178    #[cfg(feature = "tree-sitter-base")]
179    fn test_extract_structure_for_file_unsupported() {
180        let structure = extract_structure_for_file("anything", "xyz");
181        assert!(structure.is_none());
182    }
183
184    #[test]
185    #[cfg(feature = "tree-sitter-rust")]
186    fn test_find_smart_truncation_point_within_bounds() {
187        let source = "fn foo() { }\nfn bar() { }\nfn baz() { }";
188        // Max bytes larger than source — should return source length
189        let point = find_smart_truncation_point(source, 1000, "rs");
190        assert!(point.is_some());
191        assert_eq!(point.unwrap(), source.len());
192    }
193
194    #[test]
195    #[cfg(feature = "tree-sitter-rust")]
196    fn test_find_smart_truncation_point_truncated() {
197        let source = "fn foo() {\n    let x = 1;\n}\nfn bar() {\n    let y = 2;\n}";
198        let point = find_smart_truncation_point(source, 15, "rs");
199        assert!(point.is_some());
200        // Should truncate at an AST boundary, not mid-token
201        assert!(point.unwrap() <= source.len());
202    }
203
204    #[test]
205    #[cfg(feature = "tree-sitter-base")]
206    fn test_find_smart_truncation_point_unsupported() {
207        let point = find_smart_truncation_point("anything", 100, "xyz");
208        assert!(point.is_none());
209    }
210
211    #[test]
212    #[cfg(feature = "tree-sitter-rust")]
213    fn test_get_language_for_path_known() {
214        let support = get_language_for_path(Path::new("src/main.rs"));
215        assert!(support.is_some());
216    }
217
218    #[test]
219    #[cfg(feature = "tree-sitter-base")]
220    fn test_get_language_for_path_unknown() {
221        let support = get_language_for_path(Path::new("README.md"));
222        assert!(support.is_none());
223    }
224
225    #[test]
226    #[cfg(feature = "tree-sitter-base")]
227    fn test_get_language_for_path_no_extension() {
228        let support = get_language_for_path(Path::new("Makefile"));
229        assert!(support.is_none());
230    }
231}
context_builder/tree_sitter/mod.rs

context_builder/tree_sitter/
mod.rs
