Skip to main content

astmap_languages/
language.rs

1// Per-language module consolidation.
2// Concrete implementations of traits defined in astmap-core.
3// Provides dispatch (parser_for, resolver_for) and DefaultLanguageRegistry.
4
5pub mod go;
6pub mod java;
7pub mod kotlin;
8pub mod python;
9pub mod rust;
10pub mod typescript;
11
12use std::path::{Path, PathBuf};
13use std::sync::OnceLock;
14
15use tracing::trace;
16
17// Re-export types and traits from core so language implementations
18// can keep using `super::*`.
19pub use astmap_core::{
20    ExtractedCall, ExtractedImport, ExtractedSymbol, ExtractedTypeRef, LanguageParser,
21    LanguageRegistry, LanguageResolver, ParseResult,
22};
23
24// ── Shared utilities used by language modules ──
25
26/// Find the symbol that contains the given line number (innermost/smallest range).
27pub fn find_enclosing_symbol(symbols: &[ExtractedSymbol], line: usize) -> Option<&str> {
28    symbols
29        .iter()
30        .filter(|sym| line >= sym.start_line && line <= sym.end_line)
31        .min_by_key(|sym| sym.end_line - sym.start_line)
32        .map(|s| s.name.as_str())
33}
34
35// ── Shared utilities used by language modules (depend on tree-sitter) ──
36
37pub fn node_text(node: &tree_sitter::Node, source: &str) -> String {
38    source[node.start_byte()..node.end_byte()].to_string()
39}
40
41pub fn first_line(node: &tree_sitter::Node, source: &str) -> String {
42    let text = &source[node.start_byte()..node.end_byte()];
43    text.lines().next().unwrap_or("").to_string()
44}
45
46/// Build an ExtractedSymbol from a tree-sitter node.
47///
48/// Every language extracts name/kind from different node shapes, but the
49/// conversion from a tree-sitter Node to position fields is always the same.
50pub fn symbol_from_node(
51    name: String,
52    kind: astmap_core::SymbolKind,
53    node: &tree_sitter::Node,
54    source: &str,
55    parent_index: Option<usize>,
56) -> ExtractedSymbol {
57    ExtractedSymbol {
58        name,
59        kind,
60        signature: Some(first_line(node, source)),
61        start_line: node.start_position().row + 1,
62        end_line: node.end_position().row + 1,
63        start_byte: node.start_byte(),
64        end_byte: node.end_byte(),
65        parent_index,
66    }
67}
68
69/// Parse source code with a tree-sitter language and language-specific extractors.
70///
71/// The parse-tree construction and `ParseResult` assembly is identical for every
72/// language; only the extraction callbacks differ.
73pub fn parse_with_tree_sitter(
74    source: &str,
75    language: &tree_sitter::Language,
76    extract_symbols: impl FnOnce(
77        tree_sitter::Node,
78        &str,
79        &mut Vec<ExtractedSymbol>,
80        &mut Vec<ExtractedImport>,
81    ),
82    extract_calls: impl FnOnce(tree_sitter::Node, &str, &[ExtractedSymbol]) -> Vec<ExtractedCall>,
83    extract_type_refs: impl FnOnce(tree_sitter::Node, &str, &[ExtractedSymbol]) -> Vec<ExtractedTypeRef>,
84) -> ParseResult {
85    let mut parser = tree_sitter::Parser::new();
86    parser
87        .set_language(language)
88        .expect("failed to set language");
89
90    let tree = match parser.parse(source, None) {
91        Some(t) => t,
92        None => {
93            return ParseResult {
94                symbols: vec![],
95                imports: vec![],
96                calls: vec![],
97                type_refs: vec![],
98                has_errors: true,
99            };
100        }
101    };
102
103    let has_errors = tree.root_node().has_error();
104    let mut symbols = Vec::new();
105    let mut imports = Vec::new();
106
107    extract_symbols(tree.root_node(), source, &mut symbols, &mut imports);
108
109    let calls = extract_calls(tree.root_node(), source, &symbols);
110    let type_refs = extract_type_refs(tree.root_node(), source, &symbols);
111
112    ParseResult {
113        symbols,
114        calls,
115        type_refs,
116        imports,
117        has_errors,
118    }
119}
120
121/// Collect type references by walking the tree for nodes of a given kind
122/// and filtering out language-specific builtins.
123///
124/// Rust and TypeScript both look for `"type_identifier"` nodes; Python uses
125/// a different strategy so it does not use this helper.
126pub fn collect_type_refs_by_kind(
127    root: tree_sitter::Node,
128    source: &str,
129    symbols: &[ExtractedSymbol],
130    node_kind: &str,
131    is_builtin: &dyn Fn(&str) -> bool,
132) -> Vec<ExtractedTypeRef> {
133    let mut refs = Vec::new();
134    walk_type_refs(root, source, symbols, node_kind, is_builtin, &mut refs);
135    refs
136}
137
138fn walk_type_refs(
139    node: tree_sitter::Node,
140    source: &str,
141    symbols: &[ExtractedSymbol],
142    node_kind: &str,
143    is_builtin: &dyn Fn(&str) -> bool,
144    refs: &mut Vec<ExtractedTypeRef>,
145) {
146    if node.kind() == node_kind {
147        let type_name = node_text(&node, source);
148        let line = node.start_position().row + 1;
149
150        if !is_builtin(&type_name) {
151            if let Some(from_sym) = find_enclosing_symbol(symbols, line) {
152                refs.push(ExtractedTypeRef {
153                    from_symbol: from_sym.to_string(),
154                    to_type: type_name,
155                    line,
156                });
157            }
158        }
159    }
160
161    for i in 0..node.child_count() {
162        if let Some(child) = node.child(i as u32) {
163            walk_type_refs(child, source, symbols, node_kind, is_builtin, refs);
164        }
165    }
166}
167
168// ── Shared JVM helper ──
169
170/// Detect JVM source roots by checking for conventional directories.
171/// Used by both Java and Kotlin resolvers to find where source files live.
172pub fn detect_jvm_source_roots(root: &Path) -> Vec<PathBuf> {
173    let candidates = [
174        "src/main/java",
175        "src/main/kotlin",
176        "src/test/java",
177        "src/test/kotlin",
178        "src/commonMain/kotlin",
179        "src/jvmMain/kotlin",
180        "src/androidMain/kotlin",
181        "src",
182    ];
183
184    let mut roots = Vec::new();
185    for candidate in &candidates {
186        let path = root.join(candidate);
187        if path.is_dir() {
188            roots.push(path);
189        }
190    }
191
192    if roots.is_empty() {
193        roots.push(root.to_path_buf());
194    }
195    roots
196}
197
198// ── Language registration and dispatch ──
199
200pub(crate) struct LanguageSupport {
201    pub name: &'static str,
202    pub extensions: &'static [&'static str],
203    pub parser: &'static dyn LanguageParser,
204    pub resolver_factory: fn(&Path) -> Box<dyn LanguageResolver>,
205    /// Config files that trigger a full re-scan when changed (e.g., "go.mod").
206    pub config_files: &'static [&'static str],
207    /// Given a file's relative path, return sibling expansion rules (or None).
208    pub sibling_fn: fn(&str) -> Option<astmap_core::SiblingExpansion>,
209}
210
211/// All registered languages. Add new languages here.
212fn all() -> &'static [LanguageSupport] {
213    static LANGS: OnceLock<Vec<LanguageSupport>> = OnceLock::new();
214    LANGS.get_or_init(|| {
215        vec![
216            rust::lang(),
217            typescript::lang(),
218            python::lang(),
219            go::lang(),
220            java::lang(),
221            kotlin::lang(),
222        ]
223    })
224}
225
226pub fn parser_for(ext: &str) -> Option<&'static dyn LanguageParser> {
227    all().iter().find(|l| l.extensions.contains(&ext)).map(|l| {
228        trace!("selected parser: {} for extension: {}", l.name, ext);
229        l.parser
230    })
231}
232
233pub fn resolver_for(ext: &str, root: &Path) -> Option<Box<dyn LanguageResolver>> {
234    all()
235        .iter()
236        .find(|l| l.extensions.contains(&ext))
237        .map(|l| (l.resolver_factory)(root))
238}
239
240pub fn supported_extensions() -> &'static [&'static str] {
241    static EXTENSIONS: OnceLock<Vec<&str>> = OnceLock::new();
242    EXTENSIONS.get_or_init(|| {
243        all()
244            .iter()
245            .flat_map(|l| l.extensions.iter().copied())
246            .collect()
247    })
248}
249
250pub fn config_files() -> &'static [&'static str] {
251    static CONFIGS: OnceLock<Vec<&str>> = OnceLock::new();
252    CONFIGS.get_or_init(|| {
253        let mut files: Vec<&str> = all()
254            .iter()
255            .flat_map(|l| l.config_files.iter().copied())
256            .collect();
257        files.dedup();
258        files
259    })
260}
261
262pub fn sibling_expansion(rel_path: &str) -> Option<astmap_core::SiblingExpansion> {
263    let ext = Path::new(rel_path).extension()?.to_str()?;
264    let lang = all().iter().find(|l| l.extensions.contains(&ext))?;
265    (lang.sibling_fn)(rel_path)
266}
267
268// ── DefaultLanguageRegistry ──
269
270/// Concrete implementation of LanguageRegistry backed by the static dispatch tables.
271pub struct DefaultLanguageRegistry;
272
273impl LanguageRegistry for DefaultLanguageRegistry {
274    fn parser_for(&self, ext: &str) -> Option<&dyn LanguageParser> {
275        parser_for(ext)
276    }
277
278    fn resolver_for(&self, ext: &str, root: &Path) -> Option<Box<dyn LanguageResolver>> {
279        resolver_for(ext, root)
280    }
281
282    fn supported_extensions(&self) -> &[&str] {
283        supported_extensions()
284    }
285
286    fn config_files(&self) -> &[&str] {
287        config_files()
288    }
289
290    fn sibling_expansion(&self, rel_path: &str) -> Option<astmap_core::SiblingExpansion> {
291        sibling_expansion(rel_path)
292    }
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298    use astmap_core::SymbolKind;
299
300    #[test]
301    fn test_parser_for_known_extensions() {
302        assert!(parser_for("rs").is_some());
303        assert!(parser_for("ts").is_some());
304        assert!(parser_for("tsx").is_some());
305        assert!(parser_for("go").is_some());
306        assert_eq!(parser_for("rs").unwrap().language_name(), "rust");
307        assert_eq!(parser_for("ts").unwrap().language_name(), "typescript");
308        assert_eq!(parser_for("tsx").unwrap().language_name(), "typescript");
309        assert_eq!(parser_for("go").unwrap().language_name(), "go");
310    }
311
312    #[test]
313    fn test_parser_for_python() {
314        assert!(parser_for("py").is_some());
315        assert_eq!(parser_for("py").unwrap().language_name(), "python");
316    }
317
318    #[test]
319    fn test_parser_for_unknown_extension() {
320        assert!(parser_for("rb").is_none());
321        assert!(parser_for("js").is_none());
322        assert!(parser_for("").is_none());
323    }
324
325    #[test]
326    fn test_supported_extensions() {
327        let extensions = supported_extensions();
328        assert!(extensions.contains(&"rs"));
329        assert!(extensions.contains(&"ts"));
330        assert!(extensions.contains(&"tsx"));
331        assert!(extensions.contains(&"py"));
332        assert!(extensions.contains(&"go"));
333        assert!(!extensions.contains(&"js"));
334    }
335
336    #[test]
337    fn test_resolver_for() {
338        let tmp = tempfile::tempdir().unwrap();
339        assert!(resolver_for("rs", tmp.path()).is_some());
340        assert!(resolver_for("ts", tmp.path()).is_some());
341        assert!(resolver_for("tsx", tmp.path()).is_some());
342        assert!(resolver_for("py", tmp.path()).is_some());
343        assert!(resolver_for("go", tmp.path()).is_some());
344    }
345
346    #[test]
347    fn test_parser_produces_results() {
348        let parser = parser_for("ts").unwrap();
349        let result = parser.parse("function foo() {}", &std::path::PathBuf::from("test.ts"));
350        assert!(!result.has_errors);
351        assert_eq!(result.symbols.len(), 1);
352    }
353
354    #[test]
355    fn test_node_text_basic() {
356        let mut parser = tree_sitter::Parser::new();
357        parser
358            .set_language(&tree_sitter_rust::LANGUAGE.into())
359            .unwrap();
360        let source = "fn hello() {}";
361        let tree = parser.parse(source, None).unwrap();
362        let root = tree.root_node();
363        let text = node_text(&root, source);
364        assert_eq!(text, source);
365    }
366
367    #[test]
368    fn test_node_text_child_node() {
369        let mut parser = tree_sitter::Parser::new();
370        parser
371            .set_language(&tree_sitter_rust::LANGUAGE.into())
372            .unwrap();
373        let source = "fn hello() {}";
374        let tree = parser.parse(source, None).unwrap();
375        // Get the function_item node, then the name (identifier) child
376        let func = tree.root_node().child(0).unwrap();
377        let name_node = func.child_by_field_name("name").unwrap();
378        let text = node_text(&name_node, source);
379        assert_eq!(text, "hello");
380    }
381
382    #[test]
383    fn test_first_line_single_line() {
384        let mut parser = tree_sitter::Parser::new();
385        parser
386            .set_language(&tree_sitter_rust::LANGUAGE.into())
387            .unwrap();
388        let source = "fn hello() {}";
389        let tree = parser.parse(source, None).unwrap();
390        let root = tree.root_node();
391        let line = first_line(&root, source);
392        assert_eq!(line, "fn hello() {}");
393    }
394
395    #[test]
396    fn test_first_line_multi_line() {
397        let mut parser = tree_sitter::Parser::new();
398        parser
399            .set_language(&tree_sitter_rust::LANGUAGE.into())
400            .unwrap();
401        let source = "fn hello() {\n    42\n}";
402        let tree = parser.parse(source, None).unwrap();
403        let func = tree.root_node().child(0).unwrap();
404        let line = first_line(&func, source);
405        assert_eq!(line, "fn hello() {");
406    }
407
408    #[test]
409    fn test_all_returns_known_languages() {
410        let langs = all();
411        assert!(langs.len() >= 4);
412        let names: Vec<&str> = langs.iter().map(|l| l.name).collect();
413        assert!(names.contains(&"rust"));
414        assert!(names.contains(&"typescript"));
415        assert!(names.contains(&"python"));
416        assert!(names.contains(&"go"));
417    }
418
419    #[test]
420    fn test_default_language_registry_parser_for() {
421        let registry = DefaultLanguageRegistry;
422        assert!(registry.parser_for("rs").is_some());
423        assert!(registry.parser_for("ts").is_some());
424        assert!(registry.parser_for("py").is_some());
425        assert!(registry.parser_for("js").is_none());
426    }
427
428    #[test]
429    fn test_default_language_registry_resolver_for() {
430        let registry = DefaultLanguageRegistry;
431        let tmp = tempfile::tempdir().unwrap();
432        assert!(registry.resolver_for("rs", tmp.path()).is_some());
433        assert!(registry.resolver_for("tsx", tmp.path()).is_some());
434        assert!(registry.resolver_for("py", tmp.path()).is_some());
435        assert!(registry.resolver_for("rb", tmp.path()).is_none());
436    }
437
438    #[test]
439    fn test_default_language_registry_supported_extensions() {
440        let registry = DefaultLanguageRegistry;
441        let exts = registry.supported_extensions();
442        assert!(exts.contains(&"rs"));
443        assert!(exts.contains(&"ts"));
444        assert!(exts.contains(&"tsx"));
445        assert!(exts.contains(&"py"));
446        assert!(!exts.contains(&"js"));
447    }
448
449    #[test]
450    fn test_resolver_for_unknown_extension() {
451        let tmp = tempfile::tempdir().unwrap();
452        assert!(resolver_for("rb", tmp.path()).is_none());
453        assert!(resolver_for("", tmp.path()).is_none());
454    }
455
456    #[test]
457    fn test_find_enclosing_symbol_innermost() {
458        let symbols = vec![
459            ExtractedSymbol {
460                name: "Outer".to_string(),
461                kind: SymbolKind::Struct,
462                signature: None,
463                start_line: 1,
464                end_line: 20,
465                start_byte: 0,
466                end_byte: 200,
467                parent_index: None,
468            },
469            ExtractedSymbol {
470                name: "inner_fn".to_string(),
471                kind: SymbolKind::Method,
472                signature: None,
473                start_line: 5,
474                end_line: 10,
475                start_byte: 50,
476                end_byte: 100,
477                parent_index: Some(0),
478            },
479        ];
480        assert_eq!(find_enclosing_symbol(&symbols, 7), Some("inner_fn"));
481        assert_eq!(find_enclosing_symbol(&symbols, 15), Some("Outer"));
482        assert_eq!(find_enclosing_symbol(&symbols, 25), None);
483    }
484}