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