use tree_sitter::{Language, Node, Parser};
use crate::core::{Kind, Symbol};
pub(crate) mod go;
pub(crate) mod python;
pub(crate) mod ruby;
pub(crate) mod rust;
pub(crate) mod typescript;
pub(crate) struct Ctx<'a> {
src: &'a [u8],
file: &'a str,
language: &'static str,
}
impl Ctx<'_> {
pub(crate) fn field_text(&self, node: Node, field: &str) -> Option<String> {
node.child_by_field_name(field)
.and_then(|n| n.utf8_text(self.src).ok())
.map(str::to_string)
}
pub(crate) fn node_text(&self, node: Node) -> Option<String> {
node.utf8_text(self.src).ok().map(str::to_string)
}
pub(crate) fn symbol(
&self,
name: &str,
kind: Kind,
node: Node,
parent: Option<&str>,
) -> Symbol {
Symbol {
name: name.to_string(),
kind,
language: self.language.to_string(),
file: self.file.to_string(),
line: node.start_position().row as u32 + 1,
end_line: node.end_position().row as u32 + 1,
parent: parent.map(str::to_string),
visibility: None, }
}
}
pub(crate) fn qualify(parent: Option<&str>, name: &str, sep: &str) -> String {
match parent {
Some(p) => format!("{p}{sep}{name}"),
None => name.to_string(),
}
}
thread_local! {
static PARSERS: std::cell::RefCell<std::collections::HashMap<&'static str, Parser>> =
std::cell::RefCell::new(std::collections::HashMap::new());
}
pub(crate) fn extract_with(
language: &'static str,
grammar: Language,
file: &str,
source: &str,
walk: impl FnOnce(&Ctx, Node, &mut Vec<Symbol>),
) -> Vec<Symbol> {
extract_with_key(language, language, grammar, file, source, walk)
}
pub(crate) fn extract_with_key(
key: &'static str,
language: &'static str,
grammar: Language,
file: &str,
source: &str,
walk: impl FnOnce(&Ctx, Node, &mut Vec<Symbol>),
) -> Vec<Symbol> {
PARSERS.with(|cell| {
let mut parsers = cell.borrow_mut();
let parser = match parsers.entry(key) {
std::collections::hash_map::Entry::Occupied(e) => e.into_mut(),
std::collections::hash_map::Entry::Vacant(v) => {
let mut p = Parser::new();
if p.set_language(&grammar).is_err() {
return Vec::new();
}
v.insert(p)
}
};
let Some(tree) = parser.parse(source, None) else {
return Vec::new();
};
let mut out = Vec::new();
let ctx = Ctx {
src: source.as_bytes(),
file,
language,
};
walk(&ctx, tree.root_node(), &mut out);
out
})
}
pub(crate) trait LanguagePlugin {
fn language(&self) -> &'static str;
fn extensions(&self) -> &[&str];
fn extract(&self, file: &str, source: &str) -> Vec<Symbol>;
}
static REGISTRY: [&(dyn LanguagePlugin + Sync); 6] = [
&ruby::Ruby,
&rust::Rust,
&go::Go,
&python::Python,
&typescript::TypeScript,
&typescript::JavaScript,
];
pub(crate) fn languages() -> Vec<&'static str> {
registry().iter().map(|p| p.language()).collect()
}
pub(crate) fn registry() -> &'static [&'static (dyn LanguagePlugin + Sync)] {
®ISTRY
}
pub(crate) fn plugin_for_extension(ext: &str) -> Option<&'static (dyn LanguagePlugin + Sync)> {
REGISTRY
.iter()
.copied()
.find(|p| p.extensions().contains(&ext))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn languages_are_registered_by_extension() {
for ext in ["rb", "rs", "go", "py", "ts", "tsx", "js", "jsx"] {
assert!(plugin_for_extension(ext).is_some(), "{ext} should resolve");
}
assert!(plugin_for_extension("java").is_none());
}
}