Skip to main content

tsift_graph/
lang.rs

1use anyhow::Result;
2use tree_sitter::{Language, Parser, Query, QueryCursor, StreamingIterator};
3
4#[derive(Debug, Clone, PartialEq, Eq)]
5pub struct Symbol {
6    pub name: String,
7    pub kind: String,
8    pub line: usize,
9    pub end_line: usize,
10    pub node_kind: String,
11    pub start_byte: usize,
12    pub end_byte: usize,
13    pub body_start_byte: Option<usize>,
14    pub body_end_byte: Option<usize>,
15}
16
17#[allow(dead_code)]
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
19pub enum Lang {
20    #[cfg(feature = "lang-rust")]
21    Rust,
22    #[cfg(feature = "lang-python")]
23    Python,
24    #[cfg(feature = "lang-typescript")]
25    TypeScript,
26    #[cfg(feature = "lang-typescript")]
27    Tsx,
28    #[cfg(feature = "lang-javascript")]
29    JavaScript,
30    #[cfg(feature = "lang-javascript")]
31    Jsx,
32    #[cfg(feature = "lang-kotlin")]
33    Kotlin,
34    #[cfg(feature = "lang-zig")]
35    Zig,
36    #[cfg(feature = "lang-bash")]
37    Bash,
38    #[cfg(feature = "lang-gdscript")]
39    GdScript,
40    #[cfg(feature = "lang-markdown")]
41    Markdown,
42}
43
44#[allow(dead_code)]
45impl Lang {
46    pub fn from_extension(ext: &str) -> Option<Self> {
47        match ext {
48            #[cfg(feature = "lang-rust")]
49            "rs" => Some(Self::Rust),
50            #[cfg(feature = "lang-python")]
51            "py" | "pyi" => Some(Self::Python),
52            #[cfg(feature = "lang-typescript")]
53            "ts" => Some(Self::TypeScript),
54            #[cfg(feature = "lang-typescript")]
55            "tsx" => Some(Self::Tsx),
56            #[cfg(feature = "lang-javascript")]
57            "js" | "mjs" | "cjs" => Some(Self::JavaScript),
58            #[cfg(feature = "lang-javascript")]
59            "jsx" => Some(Self::Jsx),
60            #[cfg(feature = "lang-kotlin")]
61            "kt" | "kts" => Some(Self::Kotlin),
62            #[cfg(feature = "lang-zig")]
63            "zig" => Some(Self::Zig),
64            #[cfg(feature = "lang-bash")]
65            "sh" | "bash" | "zsh" => Some(Self::Bash),
66            #[cfg(feature = "lang-gdscript")]
67            "gd" => Some(Self::GdScript),
68            #[cfg(feature = "lang-markdown")]
69            "md" | "mdx" => Some(Self::Markdown),
70            _ => None,
71        }
72    }
73
74    pub fn tree_sitter_language(&self) -> Language {
75        match self {
76            #[cfg(feature = "lang-rust")]
77            Self::Rust => tree_sitter_rust::LANGUAGE.into(),
78            #[cfg(feature = "lang-python")]
79            Self::Python => tree_sitter_python::LANGUAGE.into(),
80            #[cfg(feature = "lang-typescript")]
81            Self::TypeScript => tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
82            #[cfg(feature = "lang-typescript")]
83            Self::Tsx => tree_sitter_typescript::LANGUAGE_TSX.into(),
84            #[cfg(feature = "lang-javascript")]
85            Self::JavaScript => tree_sitter_javascript::LANGUAGE.into(),
86            #[cfg(feature = "lang-javascript")]
87            Self::Jsx => tree_sitter_javascript::LANGUAGE.into(),
88            #[cfg(feature = "lang-kotlin")]
89            Self::Kotlin => tree_sitter_kotlin_ng::LANGUAGE.into(),
90            #[cfg(feature = "lang-zig")]
91            Self::Zig => tree_sitter_zig::LANGUAGE.into(),
92            #[cfg(feature = "lang-bash")]
93            Self::Bash => tree_sitter_bash::LANGUAGE.into(),
94            #[cfg(feature = "lang-gdscript")]
95            Self::GdScript => tree_sitter_gdscript::LANGUAGE.into(),
96            #[cfg(feature = "lang-markdown")]
97            Self::Markdown => tsift_md_ast::markdown_language(),
98        }
99    }
100
101    pub fn name(&self) -> &'static str {
102        match self {
103            #[cfg(feature = "lang-rust")]
104            Self::Rust => "rust",
105            #[cfg(feature = "lang-python")]
106            Self::Python => "python",
107            #[cfg(feature = "lang-typescript")]
108            Self::TypeScript => "typescript",
109            #[cfg(feature = "lang-typescript")]
110            Self::Tsx => "tsx",
111            #[cfg(feature = "lang-javascript")]
112            Self::JavaScript => "javascript",
113            #[cfg(feature = "lang-javascript")]
114            Self::Jsx => "jsx",
115            #[cfg(feature = "lang-kotlin")]
116            Self::Kotlin => "kotlin",
117            #[cfg(feature = "lang-zig")]
118            Self::Zig => "zig",
119            #[cfg(feature = "lang-bash")]
120            Self::Bash => "bash",
121            #[cfg(feature = "lang-gdscript")]
122            Self::GdScript => "gdscript",
123            #[cfg(feature = "lang-markdown")]
124            Self::Markdown => "markdown",
125        }
126    }
127
128    pub fn symbol_query(&self) -> &'static str {
129        match self {
130            #[cfg(feature = "lang-rust")]
131            Self::Rust => {
132                r#"
133                (function_item name: (identifier) @function.name)
134                (struct_item name: (type_identifier) @struct.name)
135                (enum_item name: (type_identifier) @enum.name)
136                (trait_item name: (type_identifier) @trait.name)
137                (impl_item type: (type_identifier) @impl.name)
138                (mod_item name: (identifier) @mod.name)
139                (type_item name: (type_identifier) @type_alias.name)
140                (const_item name: (identifier) @const.name)
141                (static_item name: (identifier) @static.name)
142            "#
143            }
144            #[cfg(feature = "lang-python")]
145            Self::Python => {
146                r#"
147                (function_definition name: (identifier) @function.name)
148                (class_definition name: (identifier) @class.name)
149            "#
150            }
151            #[cfg(feature = "lang-typescript")]
152            Self::TypeScript | Self::Tsx => {
153                r#"
154                (function_declaration name: (identifier) @function.name)
155                (class_declaration name: (type_identifier) @class.name)
156                (interface_declaration name: (type_identifier) @interface.name)
157                (type_alias_declaration name: (type_identifier) @type_alias.name)
158                (enum_declaration name: (identifier) @enum.name)
159                (variable_declarator name: (identifier) @function.name value: (arrow_function))
160            "#
161            }
162            #[cfg(feature = "lang-javascript")]
163            Self::JavaScript | Self::Jsx => {
164                r#"
165                (function_declaration name: (identifier) @function.name)
166                (class_declaration name: (identifier) @class.name)
167                (variable_declarator name: (identifier) @function.name value: (arrow_function))
168            "#
169            }
170            #[cfg(feature = "lang-kotlin")]
171            Self::Kotlin => {
172                r#"
173                (function_declaration name: (identifier) @function.name)
174                (class_declaration "interface" name: (identifier) @interface.name)
175                (class_declaration (modifiers (class_modifier "data")) name: (identifier) @data_class.name)
176                (class_declaration (modifiers (class_modifier "sealed")) name: (identifier) @sealed_class.name)
177                (class_declaration (modifiers (class_modifier "enum")) name: (identifier) @enum_class.name)
178                (class_declaration "class" name: (identifier) @class.name)
179                (object_declaration name: (identifier) @object.name)
180                (companion_object name: (identifier) @companion_object.name)
181            "#
182            }
183            #[cfg(feature = "lang-zig")]
184            Self::Zig => {
185                r#"
186                (function_declaration (identifier) @function.name)
187                (variable_declaration (identifier) @struct.name (struct_declaration))
188                (variable_declaration (identifier) @enum.name (enum_declaration))
189                (variable_declaration (identifier) @union.name (union_declaration))
190                (variable_declaration (identifier) @const.name)
191            "#
192            }
193            #[cfg(feature = "lang-bash")]
194            Self::Bash => {
195                r#"
196                (function_definition name: (word) @function.name)
197            "#
198            }
199            #[cfg(feature = "lang-gdscript")]
200            Self::GdScript => {
201                // `class_name Foo` declares the script's own type and is the
202                // name every other script refers to it by, so it has to be a
203                // symbol even though it is a statement rather than a block.
204                r#"
205                (function_definition name: (name) @function.name)
206                (class_definition name: (name) @class.name)
207                (class_name_statement name: (name) @class.name)
208                (enum_definition name: (name) @enum.name)
209                (signal_statement name: (name) @signal.name)
210                (const_statement name: (name) @const.name)
211                (variable_statement name: (name) @variable.name)
212                (export_variable_statement name: (name) @variable.name)
213                (onready_variable_statement name: (name) @variable.name)
214            "#
215            }
216            #[cfg(feature = "lang-markdown")]
217            Self::Markdown => {
218                r#"
219                (atx_heading (atx_h1_marker) (inline) @heading.name)
220                (atx_heading (atx_h2_marker) (inline) @heading.name)
221                (atx_heading (atx_h3_marker) (inline) @heading.name)
222                (atx_heading (atx_h4_marker) (inline) @heading.name)
223                (atx_heading (atx_h5_marker) (inline) @heading.name)
224                (atx_heading (atx_h6_marker) (inline) @heading.name)
225                (fenced_code_block (info_string (language) @code_block.name))
226            "#
227            }
228        }
229    }
230
231    pub fn call_query(&self) -> Option<&'static str> {
232        match self {
233            #[cfg(feature = "lang-rust")]
234            Self::Rust => Some(
235                r#"
236                (call_expression function: (identifier) @call.name)
237                (call_expression function: (field_expression field: (field_identifier) @call.name))
238                (call_expression function: (scoped_identifier name: (identifier) @call.name))
239                (macro_invocation macro: (identifier) @call.name)
240            "#,
241            ),
242            #[cfg(feature = "lang-python")]
243            Self::Python => Some(
244                r#"
245                (call function: (identifier) @call.name)
246                (call function: (attribute attribute: (identifier) @call.name))
247            "#,
248            ),
249            #[cfg(feature = "lang-typescript")]
250            Self::TypeScript | Self::Tsx => Some(
251                r#"
252                (call_expression function: (identifier) @call.name)
253                (call_expression function: (member_expression property: (property_identifier) @call.name))
254            "#,
255            ),
256            #[cfg(feature = "lang-javascript")]
257            Self::JavaScript | Self::Jsx => Some(
258                r#"
259                (call_expression function: (identifier) @call.name)
260                (call_expression function: (member_expression property: (property_identifier) @call.name))
261            "#,
262            ),
263            #[cfg(feature = "lang-gdscript")]
264            Self::GdScript => Some(
265                // A bare `foo()` is `(call (identifier) ...)`, while `a.foo()`
266                // puts the callee in an `attribute_call` under `attribute`, and
267                // the Godot-1.x style `.foo()` super call is `base_call`.
268                r#"
269                (call (identifier) @call.name)
270                (attribute_call (identifier) @call.name)
271                (base_call (identifier) @call.name)
272            "#,
273            ),
274            #[cfg(feature = "lang-kotlin")]
275            Self::Kotlin => Some(
276                // tree-sitter-kotlin-ng names these `identifier` /
277                // `navigation_expression`; `simple_identifier` is from the
278                // older tree-sitter-kotlin grammar and does not compile here,
279                // which silently left every Kotlin file with zero call edges.
280                r#"
281(call_expression (identifier) @call.name)
282(call_expression (navigation_expression (identifier) (identifier) @call.name))
283"#,
284            ),
285            #[cfg(feature = "lang-zig")]
286            Self::Zig => Some(
287                r#"
288(call_expression function: (identifier) @call.name)
289(call_expression function: (field_expression member: (identifier) @call.name))
290"#,
291            ),
292            _ => None,
293        }
294    }
295
296    pub fn extract_symbols(&self, source: &[u8]) -> Result<Vec<Symbol>> {
297        let mut parser = Parser::new();
298        let ts_lang = self.tree_sitter_language();
299        parser.set_language(&ts_lang)?;
300        let tree = parser
301            .parse(source, None)
302            .ok_or_else(|| anyhow::anyhow!("parse failed"))?;
303        #[cfg(feature = "lang-markdown")]
304        if *self == Self::Markdown {
305            return Ok(tsift_md_ast::markdown_symbols_from_tree(&tree, source)
306                .into_iter()
307                .map(md_symbol_to_symbol)
308                .collect());
309        }
310        let query = Query::new(&ts_lang, self.symbol_query())?;
311        let mut cursor = QueryCursor::new();
312        let mut symbols = Vec::new();
313        let capture_names: Vec<String> = query
314            .capture_names()
315            .iter()
316            .map(|s| s.to_string())
317            .collect();
318
319        let mut matches = cursor.matches(&query, tree.root_node(), source);
320        while let Some(m) = matches.next() {
321            for capture in m.captures {
322                let capture_name = &capture_names[capture.index as usize];
323                if let Some(kind_str) = capture_name.strip_suffix(".name") {
324                    let name = capture
325                        .node
326                        .utf8_text(source)
327                        .unwrap_or("<invalid utf8>")
328                        .to_string();
329                    let node = symbol_node_for_capture(kind_str, capture.node);
330                    let body_span = symbol_body_span(node);
331                    symbols.push(Symbol {
332                        name,
333                        kind: kind_str.to_string(),
334                        line: node.start_position().row,
335                        end_line: node.end_position().row,
336                        node_kind: node.kind().to_string(),
337                        start_byte: node.start_byte(),
338                        end_byte: node.end_byte(),
339                        body_start_byte: body_span.map(|(start, _)| start),
340                        body_end_byte: body_span.map(|(_, end)| end),
341                    });
342                }
343            }
344        }
345
346        #[cfg(feature = "lang-bash")]
347        if *self == Self::Bash {
348            Self::extract_bash_aliases(&tree, source, &mut symbols);
349        }
350        symbols.sort_by(|a, b| a.line.cmp(&b.line).then(a.name.cmp(&b.name)));
351        symbols.dedup_by(|b, a| {
352            a.name == b.name && a.line == b.line && {
353                let a_generic = matches!(a.kind.as_str(), "variable" | "const");
354                let b_generic = matches!(b.kind.as_str(), "variable" | "const");
355                match (a_generic, b_generic) {
356                    (true, false) => a.kind.clone_from(&b.kind),
357                    (false, true) => {}
358                    _ => {
359                        if b.kind.len() > a.kind.len() {
360                            a.kind.clone_from(&b.kind);
361                        }
362                    }
363                }
364                true
365            }
366        });
367        Ok(symbols)
368    }
369
370    #[cfg(feature = "lang-bash")]
371    fn extract_bash_aliases(tree: &tree_sitter::Tree, source: &[u8], symbols: &mut Vec<Symbol>) {
372        let mut tree_cursor = tree.root_node().walk();
373        if !tree_cursor.goto_first_child() {
374            return;
375        }
376        loop {
377            let node = tree_cursor.node();
378            if node.kind() == "command"
379                && let Some(name_node) = node.child_by_field_name("name")
380            {
381                let cmd = name_node.utf8_text(source).unwrap_or("");
382                if cmd == "alias" {
383                    for i in 0..node.named_child_count() {
384                        if let Some(arg) = node.named_child(i as u32)
385                            && (arg.kind() == "concatenation" || arg.kind() == "word")
386                        {
387                            let text = arg.utf8_text(source).unwrap_or("");
388                            if let Some(alias_name) = text.split('=').next()
389                                && !alias_name.is_empty()
390                                && alias_name != cmd
391                            {
392                                symbols.push(Symbol {
393                                    name: alias_name.to_string(),
394                                    kind: "alias".to_string(),
395                                    line: arg.start_position().row,
396                                    end_line: node.end_position().row,
397                                    node_kind: node.kind().to_string(),
398                                    start_byte: arg.start_byte(),
399                                    end_byte: node.end_byte(),
400                                    body_start_byte: None,
401                                    body_end_byte: None,
402                                });
403                            }
404                        }
405                    }
406                }
407            }
408            if !tree_cursor.goto_next_sibling() {
409                break;
410            }
411        }
412    }
413
414    pub fn all() -> Vec<Self> {
415        vec![
416            #[cfg(feature = "lang-rust")]
417            Self::Rust,
418            #[cfg(feature = "lang-python")]
419            Self::Python,
420            #[cfg(feature = "lang-typescript")]
421            Self::TypeScript,
422            #[cfg(feature = "lang-typescript")]
423            Self::Tsx,
424            #[cfg(feature = "lang-javascript")]
425            Self::JavaScript,
426            #[cfg(feature = "lang-javascript")]
427            Self::Jsx,
428            #[cfg(feature = "lang-kotlin")]
429            Self::Kotlin,
430            #[cfg(feature = "lang-zig")]
431            Self::Zig,
432            #[cfg(feature = "lang-bash")]
433            Self::Bash,
434            #[cfg(feature = "lang-gdscript")]
435            Self::GdScript,
436            #[cfg(feature = "lang-markdown")]
437            Self::Markdown,
438        ]
439    }
440}
441
442fn symbol_node_for_capture<'tree>(
443    kind: &str,
444    name_node: tree_sitter::Node<'tree>,
445) -> tree_sitter::Node<'tree> {
446    let mut node = name_node.parent().unwrap_or(name_node);
447    if kind == "code_block" {
448        while let Some(parent) = node.parent() {
449            node = parent;
450            if node.kind() == "fenced_code_block" {
451                break;
452            }
453        }
454    }
455    node
456}
457
458fn symbol_body_span(node: tree_sitter::Node<'_>) -> Option<(usize, usize)> {
459    if let Some(body) = node.child_by_field_name("body") {
460        return Some((body.start_byte(), body.end_byte()));
461    }
462    for idx in 0..node.named_child_count() {
463        let Some(child) = node.named_child(idx as u32) else {
464            continue;
465        };
466        if matches!(
467            child.kind(),
468            "block"
469                | "declaration_list"
470                | "field_declaration_list"
471                | "enum_variant_list"
472                | "match_block"
473                | "statement_block"
474                | "suite"
475        ) {
476            return Some((child.start_byte(), child.end_byte()));
477        }
478    }
479    None
480}
481
482#[cfg(feature = "lang-markdown")]
483fn md_symbol_to_symbol(md: tsift_md_ast::MdSymbol) -> Symbol {
484    Symbol {
485        name: md.name,
486        kind: md.kind,
487        line: md.line,
488        end_line: md.end_line,
489        node_kind: md.node_kind,
490        start_byte: md.start_byte,
491        end_byte: md.end_byte,
492        body_start_byte: md.body_start_byte,
493        body_end_byte: md.body_end_byte,
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500
501    #[test]
502    fn test_all_grammars_create_parser() {
503        for lang in Lang::all() {
504            let ts_lang = lang.tree_sitter_language();
505            let mut parser = tree_sitter::Parser::new();
506            parser
507                .set_language(&ts_lang)
508                .unwrap_or_else(|e| panic!("failed to set language for {:?}: {}", lang, e));
509        }
510    }
511
512    #[test]
513    fn test_extension_dispatch() {
514        let cases = [
515            ("rs", "rust"),
516            ("py", "python"),
517            ("pyi", "python"),
518            ("ts", "typescript"),
519            ("tsx", "tsx"),
520            ("js", "javascript"),
521            ("mjs", "javascript"),
522            ("cjs", "javascript"),
523            ("jsx", "jsx"),
524            ("kt", "kotlin"),
525            ("kts", "kotlin"),
526            ("zig", "zig"),
527            ("sh", "bash"),
528            ("bash", "bash"),
529            ("zsh", "bash"),
530            ("gd", "gdscript"),
531            ("md", "markdown"),
532            ("mdx", "markdown"),
533        ];
534        for (ext, expected_name) in cases {
535            let lang = Lang::from_extension(ext)
536                .unwrap_or_else(|| panic!("no language for extension: {ext}"));
537            assert_eq!(lang.name(), expected_name, "wrong language for .{ext}");
538        }
539    }
540
541    #[test]
542    fn test_unknown_extension_returns_none() {
543        assert!(Lang::from_extension("xyz").is_none());
544        assert!(Lang::from_extension("").is_none());
545        assert!(Lang::from_extension("txt").is_none());
546    }
547
548    #[cfg(feature = "lang-rust")]
549    #[test]
550    fn test_parse_rust_snippet() {
551        let lang = Lang::Rust;
552        let mut parser = tree_sitter::Parser::new();
553        parser.set_language(&lang.tree_sitter_language()).unwrap();
554        let tree = parser.parse("fn main() {}", None).unwrap();
555        assert_eq!(tree.root_node().kind(), "source_file");
556        assert!(!tree.root_node().has_error());
557    }
558
559    #[cfg(feature = "lang-python")]
560    #[test]
561    fn test_parse_python_snippet() {
562        let lang = Lang::Python;
563        let mut parser = tree_sitter::Parser::new();
564        parser.set_language(&lang.tree_sitter_language()).unwrap();
565        let tree = parser.parse("def hello():\n    pass\n", None).unwrap();
566        assert_eq!(tree.root_node().kind(), "module");
567        assert!(!tree.root_node().has_error());
568    }
569
570    #[cfg(feature = "lang-typescript")]
571    #[test]
572    fn test_parse_typescript_snippet() {
573        let lang = Lang::TypeScript;
574        let mut parser = tree_sitter::Parser::new();
575        parser.set_language(&lang.tree_sitter_language()).unwrap();
576        let tree = parser
577            .parse("function greet(name: string): void {}", None)
578            .unwrap();
579        assert_eq!(tree.root_node().kind(), "program");
580        assert!(!tree.root_node().has_error());
581    }
582
583    #[cfg(feature = "lang-typescript")]
584    #[test]
585    fn test_parse_tsx_snippet() {
586        let lang = Lang::Tsx;
587        let mut parser = tree_sitter::Parser::new();
588        parser.set_language(&lang.tree_sitter_language()).unwrap();
589        let tree = parser
590            .parse("const App = () => <div>hello</div>;", None)
591            .unwrap();
592        assert_eq!(tree.root_node().kind(), "program");
593        assert!(!tree.root_node().has_error());
594    }
595
596    #[cfg(feature = "lang-javascript")]
597    #[test]
598    fn test_parse_javascript_snippet() {
599        let lang = Lang::JavaScript;
600        let mut parser = tree_sitter::Parser::new();
601        parser.set_language(&lang.tree_sitter_language()).unwrap();
602        let tree = parser
603            .parse("function hello() { return 42; }", None)
604            .unwrap();
605        assert_eq!(tree.root_node().kind(), "program");
606        assert!(!tree.root_node().has_error());
607    }
608
609    #[cfg(feature = "lang-kotlin")]
610    #[test]
611    fn test_parse_kotlin_snippet() {
612        let lang = Lang::Kotlin;
613        let mut parser = tree_sitter::Parser::new();
614        parser.set_language(&lang.tree_sitter_language()).unwrap();
615        let tree = parser
616            .parse("fun main() { println(\"hello\") }", None)
617            .unwrap();
618        assert_eq!(tree.root_node().kind(), "source_file");
619        assert!(!tree.root_node().has_error());
620    }
621
622    #[cfg(feature = "lang-zig")]
623    #[test]
624    fn test_parse_zig_snippet() {
625        let lang = Lang::Zig;
626        let mut parser = tree_sitter::Parser::new();
627        parser.set_language(&lang.tree_sitter_language()).unwrap();
628        let tree = parser.parse("pub fn main() !void {}", None).unwrap();
629        assert_eq!(tree.root_node().kind(), "source_file");
630    }
631
632    #[cfg(feature = "lang-bash")]
633    #[test]
634    fn test_parse_bash_snippet() {
635        let lang = Lang::Bash;
636        let mut parser = tree_sitter::Parser::new();
637        parser.set_language(&lang.tree_sitter_language()).unwrap();
638        let tree = parser
639            .parse("#!/bin/bash\nhello() { echo hi; }\n", None)
640            .unwrap();
641        assert_eq!(tree.root_node().kind(), "program");
642        assert!(!tree.root_node().has_error());
643    }
644
645    #[cfg(feature = "lang-gdscript")]
646    #[test]
647    fn test_parse_gdscript_snippet() {
648        let lang = Lang::GdScript;
649        let mut parser = tree_sitter::Parser::new();
650        parser.set_language(&lang.tree_sitter_language()).unwrap();
651        let tree = parser
652            .parse("extends Node\n\nfunc _ready():\n\tprint(\"hi\")\n", None)
653            .unwrap();
654        assert_eq!(tree.root_node().kind(), "source");
655        assert!(!tree.root_node().has_error());
656    }
657
658    #[cfg(feature = "lang-markdown")]
659    #[test]
660    fn test_parse_markdown_snippet() {
661        let lang = Lang::Markdown;
662        let mut parser = tree_sitter::Parser::new();
663        parser.set_language(&lang.tree_sitter_language()).unwrap();
664        let tree = parser.parse("# Hello\n\nSome text.\n", None).unwrap();
665        assert_eq!(tree.root_node().kind(), "document");
666        assert!(!tree.root_node().has_error());
667    }
668
669    #[test]
670    fn test_all_symbol_queries_compile() {
671        for lang in Lang::all() {
672            let ts_lang = lang.tree_sitter_language();
673            tree_sitter::Query::new(&ts_lang, lang.symbol_query())
674                .unwrap_or_else(|e| panic!("query compile failed for {:?}: {}", lang, e));
675        }
676    }
677
678    #[cfg(feature = "lang-rust")]
679    #[test]
680    fn test_extract_rust_symbols() {
681        let source = b"fn main() {}\nstruct Foo;\nenum Bar {}\ntrait Baz {}\nconst X: i32 = 1;\nstatic Y: i32 = 2;\nmod inner {}\ntype Alias = i32;\n";
682        let symbols = Lang::Rust.extract_symbols(source).unwrap();
683        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
684        assert!(names.contains(&"main"), "missing main, got {:?}", names);
685        assert!(names.contains(&"Foo"), "missing Foo, got {:?}", names);
686        assert!(names.contains(&"Bar"), "missing Bar, got {:?}", names);
687        assert!(names.contains(&"Baz"), "missing Baz, got {:?}", names);
688        assert!(names.contains(&"X"), "missing X, got {:?}", names);
689        assert!(names.contains(&"Y"), "missing Y, got {:?}", names);
690        assert!(names.contains(&"inner"), "missing inner, got {:?}", names);
691        assert!(names.contains(&"Alias"), "missing Alias, got {:?}", names);
692        let main_sym = symbols.iter().find(|s| s.name == "main").unwrap();
693        assert_eq!(main_sym.kind, "function");
694        let foo_sym = symbols.iter().find(|s| s.name == "Foo").unwrap();
695        assert_eq!(foo_sym.kind, "struct");
696    }
697
698    #[cfg(feature = "lang-python")]
699    #[test]
700    fn test_extract_python_symbols() {
701        let source =
702            b"def hello():\n    pass\n\nclass MyClass:\n    def method(self):\n        pass\n";
703        let symbols = Lang::Python.extract_symbols(source).unwrap();
704        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
705        assert!(names.contains(&"hello"), "missing hello, got {:?}", names);
706        assert!(
707            names.contains(&"MyClass"),
708            "missing MyClass, got {:?}",
709            names
710        );
711        assert!(names.contains(&"method"), "missing method, got {:?}", names);
712        let cls = symbols.iter().find(|s| s.name == "MyClass").unwrap();
713        assert_eq!(cls.kind, "class");
714    }
715
716    #[cfg(feature = "lang-typescript")]
717    #[test]
718    fn test_extract_typescript_symbols() {
719        let source = b"function greet(name: string): void {}\nclass Foo {}\ninterface Bar {}\ntype Alias = string;\nenum Color { Red, Green }\n";
720        let symbols = Lang::TypeScript.extract_symbols(source).unwrap();
721        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
722        assert!(names.contains(&"greet"), "missing greet, got {:?}", names);
723        assert!(names.contains(&"Foo"), "missing Foo, got {:?}", names);
724        assert!(names.contains(&"Bar"), "missing Bar, got {:?}", names);
725        assert!(names.contains(&"Alias"), "missing Alias, got {:?}", names);
726        assert!(names.contains(&"Color"), "missing Color, got {:?}", names);
727    }
728
729    #[cfg(feature = "lang-javascript")]
730    #[test]
731    fn test_extract_javascript_symbols() {
732        let source = b"function hello() { return 42; }\nclass Widget {}\n";
733        let symbols = Lang::JavaScript.extract_symbols(source).unwrap();
734        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
735        assert!(names.contains(&"hello"), "missing hello, got {:?}", names);
736        assert!(names.contains(&"Widget"), "missing Widget, got {:?}", names);
737    }
738
739    #[cfg(feature = "lang-kotlin")]
740    #[test]
741    fn test_extract_kotlin_symbols() {
742        let source = b"fun main() { println(\"hi\") }\nclass Foo\ninterface Bar\ndata class Baz(val x: Int)\nsealed class Qux\nenum class Color { RED, GREEN }\nobject Singleton\n";
743        let symbols = Lang::Kotlin.extract_symbols(source).unwrap();
744        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
745        assert!(names.contains(&"main"), "missing main, got {:?}", names);
746        assert!(names.contains(&"Foo"), "missing Foo, got {:?}", names);
747        assert!(names.contains(&"Bar"), "missing Bar, got {:?}", names);
748        assert!(names.contains(&"Baz"), "missing Baz, got {:?}", names);
749        assert!(names.contains(&"Qux"), "missing Qux, got {:?}", names);
750        assert!(names.contains(&"Color"), "missing Color, got {:?}", names);
751        assert!(
752            names.contains(&"Singleton"),
753            "missing Singleton, got {:?}",
754            names
755        );
756        let main_sym = symbols.iter().find(|s| s.name == "main").unwrap();
757        assert_eq!(main_sym.kind, "function");
758        let foo_sym = symbols.iter().find(|s| s.name == "Foo").unwrap();
759        assert_eq!(foo_sym.kind, "class");
760        let bar_sym = symbols.iter().find(|s| s.name == "Bar").unwrap();
761        assert_eq!(bar_sym.kind, "interface");
762        let baz_sym = symbols.iter().find(|s| s.name == "Baz").unwrap();
763        assert_eq!(baz_sym.kind, "data_class");
764        let qux_sym = symbols.iter().find(|s| s.name == "Qux").unwrap();
765        assert_eq!(qux_sym.kind, "sealed_class");
766        let color_sym = symbols.iter().find(|s| s.name == "Color").unwrap();
767        assert_eq!(color_sym.kind, "enum_class");
768        let singleton_sym = symbols.iter().find(|s| s.name == "Singleton").unwrap();
769        assert_eq!(singleton_sym.kind, "object");
770        assert_eq!(
771            symbols.len(),
772            7,
773            "expected exactly 7 symbols, got {:?}",
774            symbols
775        );
776    }
777
778    #[cfg(feature = "lang-zig")]
779    #[test]
780    fn test_extract_zig_symbols() {
781        let source = b"const std = @import(\"std\");\npub fn main() !void {}\nconst Point = struct { x: i32, y: i32 };\nconst Color = enum { red, green, blue };\nconst Result = union(enum) { ok: i32, err: []const u8 };\nconst MAX: i32 = 100;\n";
782        let symbols = Lang::Zig.extract_symbols(source).unwrap();
783        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
784        assert!(names.contains(&"main"), "missing main, got {:?}", names);
785        assert!(names.contains(&"Point"), "missing Point, got {:?}", names);
786        assert!(names.contains(&"Color"), "missing Color, got {:?}", names);
787        assert!(names.contains(&"Result"), "missing Result, got {:?}", names);
788        assert!(names.contains(&"std"), "missing std, got {:?}", names);
789        assert!(names.contains(&"MAX"), "missing MAX, got {:?}", names);
790        let main_sym = symbols.iter().find(|s| s.name == "main").unwrap();
791        assert_eq!(main_sym.kind, "function");
792        let point_sym = symbols.iter().find(|s| s.name == "Point").unwrap();
793        assert_eq!(point_sym.kind, "struct");
794        let color_sym = symbols.iter().find(|s| s.name == "Color").unwrap();
795        assert_eq!(color_sym.kind, "enum");
796        let result_sym = symbols.iter().find(|s| s.name == "Result").unwrap();
797        assert_eq!(result_sym.kind, "union");
798        let max_sym = symbols.iter().find(|s| s.name == "MAX").unwrap();
799        assert_eq!(max_sym.kind, "const");
800    }
801
802    #[cfg(feature = "lang-bash")]
803    #[test]
804    fn test_extract_bash_symbols() {
805        let source = b"#!/bin/bash\nhello() { echo hi; }\nfunction world { echo world; }\nalias ll='ls -la'\nalias grep='grep --color=auto'\n";
806        let symbols = Lang::Bash.extract_symbols(source).unwrap();
807        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
808        assert!(names.contains(&"hello"), "missing hello, got {:?}", names);
809        assert!(names.contains(&"world"), "missing world, got {:?}", names);
810        assert!(names.contains(&"ll"), "missing alias ll, got {:?}", names);
811        assert!(
812            names.contains(&"grep"),
813            "missing alias grep, got {:?}",
814            names
815        );
816        let hello_sym = symbols.iter().find(|s| s.name == "hello").unwrap();
817        assert_eq!(hello_sym.kind, "function");
818        let ll_sym = symbols.iter().find(|s| s.name == "ll").unwrap();
819        assert_eq!(ll_sym.kind, "alias");
820    }
821
822    #[cfg(feature = "lang-gdscript")]
823    #[test]
824    fn test_extract_gdscript_symbols() {
825        let source = b"class_name Player\nextends CharacterBody2D\n\nsignal died(cause)\n\nenum State { IDLE, RUNNING }\n\nconst SPEED = 300.0\n\n@export var health := 100\nvar velocity_scale := 1.0\n@onready var sprite = $Sprite2D\n\nclass Inventory:\n\tvar slots = []\n\nfunc _ready():\n\tset_physics_process(true)\n";
826        let symbols = Lang::GdScript.extract_symbols(source).unwrap();
827        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
828        for expected in [
829            "Player",
830            "died",
831            "State",
832            "SPEED",
833            "health",
834            "velocity_scale",
835            "sprite",
836            "Inventory",
837            "_ready",
838        ] {
839            assert!(
840                names.contains(&expected),
841                "missing {expected}, got {names:?}"
842            );
843        }
844        let kind_of = |name: &str| {
845            symbols
846                .iter()
847                .find(|s| s.name == name)
848                .unwrap_or_else(|| panic!("missing {name}"))
849                .kind
850                .clone()
851        };
852        assert_eq!(kind_of("Player"), "class");
853        assert_eq!(kind_of("Inventory"), "class");
854        assert_eq!(kind_of("_ready"), "function");
855        assert_eq!(kind_of("died"), "signal");
856        assert_eq!(kind_of("State"), "enum");
857        assert_eq!(kind_of("SPEED"), "const");
858        assert_eq!(kind_of("health"), "variable");
859        assert_eq!(kind_of("sprite"), "variable");
860    }
861
862    #[cfg(feature = "lang-markdown")]
863    #[test]
864    fn test_extract_markdown_symbols() {
865        let source = b"# Title\n\n## Section One\n\nSome text.\n\n- Run setup\n  - Confirm setup\n\n```rust\nfn main() {}\n```\n\n### Subsection\n\n```python\ndef hello():\n    pass\n```\n\n## Next Section\n\nDone.\n";
866        let symbols = Lang::Markdown.extract_symbols(source).unwrap();
867        let headings: Vec<&Symbol> = symbols.iter().filter(|s| s.kind == "heading").collect();
868        let code_blocks: Vec<&Symbol> = symbols.iter().filter(|s| s.kind == "code_block").collect();
869        let list_items: Vec<&Symbol> = symbols.iter().filter(|s| s.kind == "list_item").collect();
870        assert_eq!(headings.len(), 4, "expected 4 headings, got {:?}", headings);
871        assert_eq!(
872            code_blocks.len(),
873            2,
874            "expected 2 code blocks, got {:?}",
875            code_blocks
876        );
877        assert_eq!(
878            list_items.len(),
879            2,
880            "expected 2 list items, got {:?}",
881            list_items
882        );
883        let title = headings.iter().find(|s| s.name == "Title").unwrap();
884        let section = headings.iter().find(|s| s.name == "Section One").unwrap();
885        let next = headings.iter().find(|s| s.name == "Next Section").unwrap();
886        assert_eq!(title.node_kind, "atx_heading");
887        assert!(title.end_byte > next.start_byte);
888        assert_eq!(section.end_byte, next.start_byte);
889        assert!(
890            section.body_start_byte.unwrap() > section.start_byte,
891            "heading body should begin after the marker line"
892        );
893        assert!(
894            code_blocks.iter().any(|s| s.name == "rust"),
895            "missing rust block, got {:?}",
896            code_blocks
897        );
898        assert!(
899            code_blocks.iter().any(|s| s.name == "python"),
900            "missing python block, got {:?}",
901            code_blocks
902        );
903        assert!(
904            list_items.iter().any(|s| s.name == "Run setup"),
905            "missing top-level list item, got {:?}",
906            list_items
907        );
908    }
909
910    #[cfg(feature = "lang-python")]
911    #[test]
912    fn test_python_async_def() {
913        let source = b"async def fetch_data():\n    await get()\n\ndef sync_fn():\n    pass\n";
914        let symbols = Lang::Python.extract_symbols(source).unwrap();
915        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
916        assert!(
917            names.contains(&"fetch_data"),
918            "missing async function, got {:?}",
919            names
920        );
921        assert!(
922            names.contains(&"sync_fn"),
923            "missing sync function, got {:?}",
924            names
925        );
926    }
927
928    #[cfg(feature = "lang-python")]
929    #[test]
930    fn test_python_decorated_function() {
931        let source = b"@staticmethod\ndef helper():\n    pass\n\n@property\ndef name(self):\n    return self._name\n";
932        let symbols = Lang::Python.extract_symbols(source).unwrap();
933        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
934        assert!(
935            names.contains(&"helper"),
936            "missing decorated function, got {:?}",
937            names
938        );
939        assert!(
940            names.contains(&"name"),
941            "missing property function, got {:?}",
942            names
943        );
944    }
945
946    #[cfg(feature = "lang-typescript")]
947    #[test]
948    fn test_typescript_arrow_exports() {
949        let source = b"export const Foo = () => { return 42; };\nexport const Bar = (x: number): number => x + 1;\nconst local = () => {};\nfunction regular() {}\n";
950        let symbols = Lang::TypeScript.extract_symbols(source).unwrap();
951        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
952        assert!(
953            names.contains(&"Foo"),
954            "missing arrow export Foo, got {:?}",
955            names
956        );
957        assert!(
958            names.contains(&"Bar"),
959            "missing arrow export Bar, got {:?}",
960            names
961        );
962        assert!(
963            names.contains(&"local"),
964            "missing local arrow, got {:?}",
965            names
966        );
967        assert!(
968            names.contains(&"regular"),
969            "missing regular function, got {:?}",
970            names
971        );
972    }
973
974    #[cfg(feature = "lang-typescript")]
975    #[test]
976    fn test_tsx_arrow_component() {
977        let source = b"export const MyComponent = () => <div>hello</div>;\nfunction Other() { return <span/>; }\n";
978        let symbols = Lang::Tsx.extract_symbols(source).unwrap();
979        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
980        assert!(
981            names.contains(&"MyComponent"),
982            "missing arrow component, got {:?}",
983            names
984        );
985        assert!(
986            names.contains(&"Other"),
987            "missing function component, got {:?}",
988            names
989        );
990    }
991
992    #[cfg(feature = "lang-javascript")]
993    #[test]
994    fn test_javascript_arrow_exports() {
995        let source = b"export const handler = () => { return 'ok'; };\nconst helper = (x) => x * 2;\nfunction regular() {}\n";
996        let symbols = Lang::JavaScript.extract_symbols(source).unwrap();
997        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
998        assert!(
999            names.contains(&"handler"),
1000            "missing arrow export, got {:?}",
1001            names
1002        );
1003        assert!(
1004            names.contains(&"helper"),
1005            "missing local arrow, got {:?}",
1006            names
1007        );
1008        assert!(
1009            names.contains(&"regular"),
1010            "missing regular function, got {:?}",
1011            names
1012        );
1013    }
1014
1015    #[cfg(feature = "lang-javascript")]
1016    #[test]
1017    fn test_jsx_arrow_component() {
1018        let source = b"const App = () => <div>hi</div>;\nfunction Page() { return <main/>; }\n";
1019        let symbols = Lang::Jsx.extract_symbols(source).unwrap();
1020        let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
1021        assert!(
1022            names.contains(&"App"),
1023            "missing arrow JSX component, got {:?}",
1024            names
1025        );
1026        assert!(
1027            names.contains(&"Page"),
1028            "missing function component, got {:?}",
1029            names
1030        );
1031    }
1032}