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 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 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 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 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 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 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 #[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}