Skip to main content

reference_query/lang/
mod.rs

1//! Language plugins — the only seam languages plug into.
2//!
3//! A plugin maps source text to the common [`Symbol`](crate::core::Symbol)
4//! model. The core stays language-agnostic; adding a language is a new plugin,
5//! not a core change.
6
7use tree_sitter::{Language, Node, Parser};
8
9use crate::core::{Kind, Symbol};
10
11pub(crate) mod go;
12pub(crate) mod python;
13pub(crate) mod ruby;
14pub(crate) mod rust;
15pub(crate) mod typescript;
16
17/// Per-file extraction context shared by every plugin: the source bytes, the
18/// repo-relative path, and the language tag stamped on each emitted symbol.
19pub(crate) struct Ctx<'a> {
20    src: &'a [u8],
21    file: &'a str,
22    language: &'static str,
23}
24
25impl Ctx<'_> {
26    /// The text of `node`'s named field, if present.
27    pub(crate) fn field_text(&self, node: Node, field: &str) -> Option<String> {
28        node.child_by_field_name(field)
29            .and_then(|n| n.utf8_text(self.src).ok())
30            .map(str::to_string)
31    }
32
33    /// The text of `node` itself.
34    pub(crate) fn node_text(&self, node: Node) -> Option<String> {
35        node.utf8_text(self.src).ok().map(str::to_string)
36    }
37
38    /// Build a [`Symbol`] for `node` (1-based line span).
39    pub(crate) fn symbol(
40        &self,
41        name: &str,
42        kind: Kind,
43        node: Node,
44        parent: Option<&str>,
45    ) -> Symbol {
46        Symbol {
47            name: name.to_string(),
48            kind,
49            language: self.language.to_string(),
50            file: self.file.to_string(),
51            line: node.start_position().row as u32 + 1,
52            end_line: node.end_position().row as u32 + 1,
53            parent: parent.map(str::to_string),
54            visibility: None, // plugins that know it set it on the result
55        }
56    }
57}
58
59/// Join a name onto its enclosing qualified name with the language's separator.
60pub(crate) fn qualify(parent: Option<&str>, name: &str, sep: &str) -> String {
61    match parent {
62        Some(p) => format!("{p}{sep}{name}"),
63        None => name.to_string(),
64    }
65}
66
67thread_local! {
68    /// One parser per language per thread. `set_language` (grammar table
69    /// loading) is the expensive step of parser setup, and the indexer calls
70    /// `extract` once per file — reuse makes that a one-time cost per worker.
71    static PARSERS: std::cell::RefCell<std::collections::HashMap<&'static str, Parser>> =
72        std::cell::RefCell::new(std::collections::HashMap::new());
73}
74
75/// Parse `source` with `grammar` and hand the tree's root (plus a [`Ctx`]) to
76/// the plugin's `walk`. All the per-file plumbing lives here; a plugin is just
77/// its walk. (The parser cache is borrowed across the walk, so a walk must
78/// never recurse into another `extract` — none does.)
79pub(crate) fn extract_with(
80    language: &'static str,
81    grammar: Language,
82    file: &str,
83    source: &str,
84    walk: impl FnOnce(&Ctx, Node, &mut Vec<Symbol>),
85) -> Vec<Symbol> {
86    extract_with_key(language, language, grammar, file, source, walk)
87}
88
89/// [`extract_with`] with the parser-cache key named separately from the language
90/// tag — for a plugin that spans more than one grammar (TypeScript's `.ts` vs
91/// `.tsx`) or a grammar shared by two tags. The key identifies the *grammar*, so
92/// it must be distinct per grammar and identical wherever that grammar is used.
93pub(crate) fn extract_with_key(
94    key: &'static str,
95    language: &'static str,
96    grammar: Language,
97    file: &str,
98    source: &str,
99    walk: impl FnOnce(&Ctx, Node, &mut Vec<Symbol>),
100) -> Vec<Symbol> {
101    PARSERS.with(|cell| {
102        let mut parsers = cell.borrow_mut();
103        let parser = match parsers.entry(key) {
104            std::collections::hash_map::Entry::Occupied(e) => e.into_mut(),
105            std::collections::hash_map::Entry::Vacant(v) => {
106                let mut p = Parser::new();
107                if p.set_language(&grammar).is_err() {
108                    return Vec::new();
109                }
110                v.insert(p)
111            }
112        };
113        let Some(tree) = parser.parse(source, None) else {
114            return Vec::new();
115        };
116        let mut out = Vec::new();
117        let ctx = Ctx {
118            src: source.as_bytes(),
119            file,
120            language,
121        };
122        walk(&ctx, tree.root_node(), &mut out);
123        out
124    })
125}
126
127/// Extracts definitions from a single source file.
128pub(crate) trait LanguagePlugin {
129    /// The language tag emitted on every [`Symbol`] (e.g. `"ruby"`). Also the
130    /// canonical name `--lang` matches against.
131    fn language(&self) -> &'static str;
132
133    /// File extensions this plugin handles, without the dot (e.g. `["rb"]`).
134    fn extensions(&self) -> &[&str];
135
136    /// Extract definitions from `source`. `file` is the repo-relative path,
137    /// recorded on each emitted [`Symbol`].
138    fn extract(&self, file: &str, source: &str) -> Vec<Symbol>;
139}
140
141/// The registered language plugins. Adding a language is one line here.
142static REGISTRY: [&(dyn LanguagePlugin + Sync); 6] = [
143    &ruby::Ruby,
144    &rust::Rust,
145    &go::Go,
146    &python::Python,
147    &typescript::TypeScript,
148    &typescript::JavaScript,
149];
150
151/// The tags of all registered languages — the set `--lang` matches against, so
152/// it can't drift from the registry.
153pub(crate) fn languages() -> Vec<&'static str> {
154    registry().iter().map(|p| p.language()).collect()
155}
156
157/// The registered language plugins.
158pub(crate) fn registry() -> &'static [&'static (dyn LanguagePlugin + Sync)] {
159    &REGISTRY
160}
161
162/// The plugin handling files with the given extension (without the dot), if any.
163pub(crate) fn plugin_for_extension(ext: &str) -> Option<&'static (dyn LanguagePlugin + Sync)> {
164    REGISTRY
165        .iter()
166        .copied()
167        .find(|p| p.extensions().contains(&ext))
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn languages_are_registered_by_extension() {
176        for ext in ["rb", "rs", "go", "py", "ts", "tsx", "js", "jsx"] {
177            assert!(plugin_for_extension(ext).is_some(), "{ext} should resolve");
178        }
179        assert!(plugin_for_extension("java").is_none());
180    }
181}