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