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        "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
54/// Whether `ext` has a symbol extractor (used to filter repo-map candidates).
55pub fn supports(ext: &str) -> bool {
56    lang_and_query(ext).is_some()
57}
58
59/// Compile the tags configuration for `ext`, or `None` if unsupported / the
60/// query fails to compile. Compiling is non-trivial, so callers cache per
61/// language across a repo-map pass.
62pub fn config_for(ext: &str) -> Option<TagsConfiguration> {
63    let (language, query) = lang_and_query(ext)?;
64    TagsConfiguration::new(language, query, "").ok()
65}
66
67/// Extract definition symbols from `text` using a pre-built `config`.
68/// `context` is reusable across files.
69pub fn extract(context: &mut TagsContext, config: &TagsConfiguration, text: &str) -> Vec<Symbol> {
70    collect(context, config, text, true)
71}
72
73/// Extract all occurrences — definitions AND references (call/use sites).
74pub fn extract_all(
75    context: &mut TagsContext,
76    config: &TagsConfiguration,
77    text: &str,
78) -> Vec<Symbol> {
79    collect(context, config, text, false)
80}
81
82/// Shared collector. Best-effort: a parse/iteration error yields whatever was
83/// gathered so far (possibly empty). `source[tag.name_range]` is safe because
84/// `text` is valid UTF-8 and the range comes from tree-sitter over that buffer.
85fn 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        // alpha is defined on line 1.
130        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}