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 #[cfg(feature = "languages-extra")]
39 "ex" | "exs" => (tree_sitter_elixir::LANGUAGE.into(), tree_sitter_elixir::TAGS_QUERY),
40 "go" => (tree_sitter_go::LANGUAGE.into(), tree_sitter_go::TAGS_QUERY),
41 "c" | "h" => (tree_sitter_c::LANGUAGE.into(), tree_sitter_c::TAGS_QUERY),
42 #[cfg(feature = "languages-extra")]
43 "cpp" | "cc" | "cxx" | "hpp" | "hxx" | "hh" => {
44 (tree_sitter_cpp::LANGUAGE.into(), tree_sitter_cpp::TAGS_QUERY)
45 }
46 "java" => (tree_sitter_java::LANGUAGE.into(), tree_sitter_java::TAGS_QUERY),
47 #[cfg(feature = "languages-extra")]
48 "rb" => (tree_sitter_ruby::LANGUAGE.into(), tree_sitter_ruby::TAGS_QUERY),
49 #[cfg(feature = "languages-extra")]
50 "cs" => (tree_sitter_c_sharp::LANGUAGE.into(), tree_sitter_c_sharp::TAGS_QUERY),
51 #[cfg(feature = "languages-extra")]
52 "php" => (tree_sitter_php::LANGUAGE_PHP.into(), tree_sitter_php::TAGS_QUERY),
53 #[cfg(feature = "languages-extra")]
54 "lua" => (tree_sitter_lua::LANGUAGE.into(), tree_sitter_lua::TAGS_QUERY),
55 _ => return None,
56 };
57 Some(pair)
58}
59
60pub fn supports(ext: &str) -> bool {
62 lang_and_query(ext).is_some()
63}
64
65pub fn config_for(ext: &str) -> Option<TagsConfiguration> {
69 let (language, query) = lang_and_query(ext)?;
70 TagsConfiguration::new(language, query, "").ok()
71}
72
73pub fn extract(context: &mut TagsContext, config: &TagsConfiguration, text: &str) -> Vec<Symbol> {
76 collect(context, config, text, true)
77}
78
79pub fn extract_all(
81 context: &mut TagsContext,
82 config: &TagsConfiguration,
83 text: &str,
84) -> Vec<Symbol> {
85 collect(context, config, text, false)
86}
87
88fn collect(
92 context: &mut TagsContext,
93 config: &TagsConfiguration,
94 text: &str,
95 definitions_only: bool,
96) -> Vec<Symbol> {
97 let source = text.as_bytes();
98 let mut out = Vec::new();
99 let Ok((tags, _)) = context.generate_tags(config, source, None) else {
100 return out;
101 };
102 for tag in tags.flatten() {
103 if definitions_only && !tag.is_definition {
104 continue;
105 }
106 let name = String::from_utf8_lossy(&source[tag.name_range.clone()]).into_owned();
107 if name.is_empty() {
108 continue;
109 }
110 out.push(Symbol {
111 name,
112 kind: config.syntax_type_name(tag.syntax_type_id).to_string(),
113 line: tag.span.start.row + 1,
114 is_definition: tag.is_definition,
115 });
116 }
117 out
118}
119
120#[cfg(test)]
121mod tests {
122 #![allow(clippy::expect_used)]
123 use super::*;
124
125 #[test]
126 fn extracts_rust_symbols() {
127 let mut ctx = TagsContext::new();
128 let cfg = config_for("rs").expect("rust config");
129 let src = "pub fn alpha() {}\nstruct Beta;\nimpl Beta { fn gamma(&self) {} }\n";
130 let syms = extract(&mut ctx, &cfg, src);
131 let names: Vec<&str> = syms.iter().map(|s| s.name.as_str()).collect();
132 assert!(names.contains(&"alpha"), "got {names:?}");
133 assert!(names.contains(&"Beta"), "got {names:?}");
134 assert!(names.contains(&"gamma"), "got {names:?}");
135 assert_eq!(syms.iter().find(|s| s.name == "alpha").map(|s| s.line), Some(1));
137 }
138
139 #[test]
140 fn unsupported_extension_has_no_config() {
141 assert!(config_for("json").is_none());
142 assert!(!supports("css"));
143 assert!(supports("rs"));
144 assert!(supports("ts"));
145 assert!(supports("py"));
146 assert_eq!(supports("ex"), cfg!(feature = "languages-extra"));
147 }
148
149 #[test]
150 fn extracts_python_symbols() {
151 let mut ctx = TagsContext::new();
152 let cfg = config_for("py").expect("python config");
153 let src = "class Greeter:\n def hello(self, name):\n return name\n\ndef top():\n return Greeter()\n";
154 let syms = extract(&mut ctx, &cfg, src);
155 let names: Vec<&str> = syms.iter().map(|symbol| symbol.name.as_str()).collect();
156 assert!(names.contains(&"Greeter"), "got {names:?}");
157 assert!(names.contains(&"hello"), "got {names:?}");
158 assert!(names.contains(&"top"), "got {names:?}");
159 }
160
161 #[cfg(feature = "languages-extra")]
162 #[test]
163 fn extracts_elixir_symbols() {
164 let mut ctx = TagsContext::new();
165 let cfg = config_for("ex").expect("elixir config");
166 let src = "defmodule Demo do\n def hello(name), do: name\n defmacro traced(expr), do: expr\nend\n";
167 let syms = extract(&mut ctx, &cfg, src);
168 let names: Vec<&str> = syms.iter().map(|symbol| symbol.name.as_str()).collect();
169 assert!(names.contains(&"Demo"), "got {names:?}");
170 assert!(names.contains(&"hello"), "got {names:?}");
171 assert!(names.contains(&"traced"), "got {names:?}");
172 }
173}