winx_code_agent/utils/
symbols.rs1use tree_sitter_tags::{TagsConfiguration, TagsContext};
10
11#[derive(Debug, Clone)]
13pub struct Symbol {
14 pub name: String,
16 pub kind: String,
19 pub line: usize,
21 pub is_definition: bool,
23}
24
25fn lang_and_query(ext: &str) -> Option<(tree_sitter::Language, &'static str)> {
28 let pair = match ext {
29 "rs" => (tree_sitter_rust::LANGUAGE.into(), tree_sitter_rust::TAGS_QUERY),
30 "js" | "mjs" | "cjs" | "jsx" => {
31 (tree_sitter_javascript::LANGUAGE.into(), tree_sitter_javascript::TAGS_QUERY)
32 }
33 "ts" => {
34 (tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(), tree_sitter_typescript::TAGS_QUERY)
35 }
36 "tsx" => (tree_sitter_typescript::LANGUAGE_TSX.into(), tree_sitter_typescript::TAGS_QUERY),
37 "py" | "pyi" => (tree_sitter_python::LANGUAGE.into(), tree_sitter_python::TAGS_QUERY),
38 "ex" | "exs" => (tree_sitter_elixir::LANGUAGE.into(), tree_sitter_elixir::TAGS_QUERY),
39 "go" => (tree_sitter_go::LANGUAGE.into(), tree_sitter_go::TAGS_QUERY),
40 "c" | "h" => (tree_sitter_c::LANGUAGE.into(), tree_sitter_c::TAGS_QUERY),
41 "cpp" | "cc" | "cxx" | "hpp" | "hxx" | "hh" => {
42 (tree_sitter_cpp::LANGUAGE.into(), tree_sitter_cpp::TAGS_QUERY)
43 }
44 "java" => (tree_sitter_java::LANGUAGE.into(), tree_sitter_java::TAGS_QUERY),
45 "rb" => (tree_sitter_ruby::LANGUAGE.into(), tree_sitter_ruby::TAGS_QUERY),
46 "cs" => (tree_sitter_c_sharp::LANGUAGE.into(), tree_sitter_c_sharp::TAGS_QUERY),
47 "php" => (tree_sitter_php::LANGUAGE_PHP.into(), tree_sitter_php::TAGS_QUERY),
48 "lua" => (tree_sitter_lua::LANGUAGE.into(), tree_sitter_lua::TAGS_QUERY),
49 _ => return None,
50 };
51 Some(pair)
52}
53
54pub fn supports(ext: &str) -> bool {
56 lang_and_query(ext).is_some()
57}
58
59pub fn config_for(ext: &str) -> Option<TagsConfiguration> {
63 let (language, query) = lang_and_query(ext)?;
64 TagsConfiguration::new(language, query, "").ok()
65}
66
67pub fn extract(context: &mut TagsContext, config: &TagsConfiguration, text: &str) -> Vec<Symbol> {
70 collect(context, config, text, true)
71}
72
73pub fn extract_all(
75 context: &mut TagsContext,
76 config: &TagsConfiguration,
77 text: &str,
78) -> Vec<Symbol> {
79 collect(context, config, text, false)
80}
81
82fn collect(
86 context: &mut TagsContext,
87 config: &TagsConfiguration,
88 text: &str,
89 definitions_only: bool,
90) -> Vec<Symbol> {
91 let source = text.as_bytes();
92 let mut out = Vec::new();
93 let Ok((tags, _)) = context.generate_tags(config, source, None) else {
94 return out;
95 };
96 for tag in tags.flatten() {
97 if definitions_only && !tag.is_definition {
98 continue;
99 }
100 let name = String::from_utf8_lossy(&source[tag.name_range.clone()]).into_owned();
101 if name.is_empty() {
102 continue;
103 }
104 out.push(Symbol {
105 name,
106 kind: config.syntax_type_name(tag.syntax_type_id).to_string(),
107 line: tag.span.start.row + 1,
108 is_definition: tag.is_definition,
109 });
110 }
111 out
112}
113
114#[cfg(test)]
115mod tests {
116 #![allow(clippy::expect_used)]
117 use super::*;
118
119 #[test]
120 fn extracts_rust_symbols() {
121 let mut ctx = TagsContext::new();
122 let cfg = config_for("rs").expect("rust config");
123 let src = "pub fn alpha() {}\nstruct Beta;\nimpl Beta { fn gamma(&self) {} }\n";
124 let syms = extract(&mut ctx, &cfg, src);
125 let names: Vec<&str> = syms.iter().map(|s| s.name.as_str()).collect();
126 assert!(names.contains(&"alpha"), "got {names:?}");
127 assert!(names.contains(&"Beta"), "got {names:?}");
128 assert!(names.contains(&"gamma"), "got {names:?}");
129 assert_eq!(syms.iter().find(|s| s.name == "alpha").map(|s| s.line), Some(1));
131 }
132
133 #[test]
134 fn unsupported_extension_has_no_config() {
135 assert!(config_for("json").is_none());
136 assert!(!supports("css"));
137 assert!(supports("rs"));
138 assert!(supports("ts"));
139 assert!(supports("py"));
140 assert!(supports("ex"));
141 }
142
143 #[test]
144 fn extracts_python_symbols() {
145 let mut ctx = TagsContext::new();
146 let cfg = config_for("py").expect("python config");
147 let src = "class Greeter:\n def hello(self, name):\n return name\n\ndef top():\n return Greeter()\n";
148 let syms = extract(&mut ctx, &cfg, src);
149 let names: Vec<&str> = syms.iter().map(|symbol| symbol.name.as_str()).collect();
150 assert!(names.contains(&"Greeter"), "got {names:?}");
151 assert!(names.contains(&"hello"), "got {names:?}");
152 assert!(names.contains(&"top"), "got {names:?}");
153 }
154
155 #[test]
156 fn extracts_elixir_symbols() {
157 let mut ctx = TagsContext::new();
158 let cfg = config_for("ex").expect("elixir config");
159 let src = "defmodule Demo do\n def hello(name), do: name\n defmacro traced(expr), do: expr\nend\n";
160 let syms = extract(&mut ctx, &cfg, src);
161 let names: Vec<&str> = syms.iter().map(|symbol| symbol.name.as_str()).collect();
162 assert!(names.contains(&"Demo"), "got {names:?}");
163 assert!(names.contains(&"hello"), "got {names:?}");
164 assert!(names.contains(&"traced"), "got {names:?}");
165 }
166}