Skip to main content

winx_code_agent/utils/
symbols.rs

1//! Symbol/outline extraction via `tree-sitter-tags` and each grammar's embedded
2//! `TAGS_QUERY` — the same engine GitHub uses for code navigation. We surface
3//! only definitions (functions, types, methods, ...), not references.
4//!
5//! 13 languages have a tags query (rust, js/ts, python, elixir, go, c, c++, java,
6//! ruby, c#, php, lua); the others (bash/css/html/json) have no code symbols and
7//! return empty.
8
9use tree_sitter_tags::{TagsConfiguration, TagsContext};
10
11/// A code symbol occurrence (definition or reference) from a source file.
12#[derive(Debug, Clone)]
13pub struct Symbol {
14    /// Symbol name (e.g. `parse_config`).
15    pub name: String,
16    /// Symbol kind from the grammar's tags query (e.g. `function`, `struct`,
17    /// `method`, `class`, `call`).
18    pub kind: String,
19    /// 1-based line where the occurrence starts.
20    pub line: usize,
21    /// True for a definition, false for a reference (call / use site).
22    pub is_definition: bool,
23}
24
25/// Map a file extension to its tree-sitter language + embedded tags query.
26/// Returns `None` for unsupported / non-code extensions.
27fn 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
60/// Whether `ext` has a symbol extractor (used to filter repo-map candidates).
61pub fn supports(ext: &str) -> bool {
62    lang_and_query(ext).is_some()
63}
64
65/// Compile the tags configuration for `ext`, or `None` if unsupported / the
66/// query fails to compile. Compiling is non-trivial, so callers cache per
67/// language across a repo-map pass.
68pub fn config_for(ext: &str) -> Option<TagsConfiguration> {
69    let (language, query) = lang_and_query(ext)?;
70    TagsConfiguration::new(language, query, "").ok()
71}
72
73/// Extract definition symbols from `text` using a pre-built `config`.
74/// `context` is reusable across files.
75pub fn extract(context: &mut TagsContext, config: &TagsConfiguration, text: &str) -> Vec<Symbol> {
76    collect(context, config, text, true)
77}
78
79/// Extract all occurrences — definitions AND references (call/use sites).
80pub fn extract_all(
81    context: &mut TagsContext,
82    config: &TagsConfiguration,
83    text: &str,
84) -> Vec<Symbol> {
85    collect(context, config, text, false)
86}
87
88/// Shared collector. Best-effort: a parse/iteration error yields whatever was
89/// gathered so far (possibly empty). `source[tag.name_range]` is safe because
90/// `text` is valid UTF-8 and the range comes from tree-sitter over that buffer.
91fn 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        // alpha is defined on line 1.
136        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}