use std::cell::RefCell;
use std::collections::HashMap;
use tree_sitter::{Node, Parser, Point};
use crate::lang::Language;
thread_local! {
static RESOLVE_PARSERS: RefCell<HashMap<Language, Parser>> = RefCell::new(HashMap::new());
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IdentRef {
pub name: String,
pub is_call: bool,
pub is_member: bool,
pub is_type: bool,
pub is_import: bool,
}
pub trait LangResolver: Send + Sync {
fn language(&self) -> Language;
fn member_kinds(&self) -> &'static [&'static str];
fn call_kinds(&self) -> &'static [&'static str];
fn type_kinds(&self) -> &'static [&'static str] {
&["type_identifier", "type_annotation", "generic_type", "type"]
}
fn import_kinds(&self) -> &'static [&'static str];
}
macro_rules! lang_resolver {
($name:ident, $lang:expr, member = $member:expr, call = $call:expr, import = $import:expr) => {
struct $name;
impl LangResolver for $name {
fn language(&self) -> Language {
$lang
}
fn member_kinds(&self) -> &'static [&'static str] {
$member
}
fn call_kinds(&self) -> &'static [&'static str] {
$call
}
fn import_kinds(&self) -> &'static [&'static str] {
$import
}
}
};
}
lang_resolver!(
RustResolver,
Language::Rust,
member = &["field_expression"],
call = &["call_expression", "macro_invocation"],
import = &["use_declaration"]
);
lang_resolver!(
PythonResolver,
Language::Python,
member = &["attribute"],
call = &["call"],
import = &["import_statement", "import_from_statement"]
);
lang_resolver!(
JsResolver,
Language::JavaScript,
member = &["member_expression"],
call = &["call_expression", "new_expression"],
import = &["import_statement"]
);
lang_resolver!(
TsResolver,
Language::TypeScript,
member = &["member_expression"],
call = &["call_expression", "new_expression"],
import = &["import_statement"]
);
lang_resolver!(
TsxResolver,
Language::Tsx,
member = &["member_expression"],
call = &["call_expression", "new_expression"],
import = &["import_statement"]
);
lang_resolver!(
GoResolver,
Language::Go,
member = &["selector_expression"],
call = &["call_expression"],
import = &["import_spec"]
);
lang_resolver!(
DartResolver,
Language::Dart,
member = &["member_expression"],
call = &["call_expression", "constructor_invocation"],
import = &["library_import", "import_specification"]
);
pub fn resolver_for(lang: Language) -> Box<dyn LangResolver> {
match lang {
Language::Rust => Box::new(RustResolver),
Language::Python => Box::new(PythonResolver),
Language::JavaScript => Box::new(JsResolver),
Language::TypeScript => Box::new(TsResolver),
Language::Tsx => Box::new(TsxResolver),
Language::Go => Box::new(GoResolver),
Language::Dart => Box::new(DartResolver),
other => Box::new(GenericResolver(other)),
}
}
struct GenericResolver(Language);
impl LangResolver for GenericResolver {
fn language(&self) -> Language {
self.0
}
fn member_kinds(&self) -> &'static [&'static str] {
&[
"field_expression",
"member_expression",
"member_access_expression",
"selector_expression",
"attribute",
"scoped_identifier",
]
}
fn call_kinds(&self) -> &'static [&'static str] {
&[
"call_expression",
"call",
"method_invocation",
"invocation_expression",
"function_call_expression",
"member_call_expression",
"object_creation_expression",
]
}
fn import_kinds(&self) -> &'static [&'static str] {
&[
"use_declaration",
"import_statement",
"import_from_statement",
"import_declaration",
"import_spec",
"using_directive",
"namespace_use_declaration",
]
}
}
fn is_ident_kind(kind: &str) -> bool {
kind.ends_with("identifier") || kind == "name" || kind == "constant" || kind == "property"
}
pub fn identifier_at(lang: Language, source: &[u8], line: u32, col: u32) -> Option<IdentRef> {
let grammar = lang.grammar()?;
let res = resolver_for(lang);
RESOLVE_PARSERS.with(|cell| {
let mut map = cell.borrow_mut();
let parser = map.entry(lang).or_insert_with(|| {
let mut p = Parser::new();
let _ = p.set_language(&grammar);
p
});
let tree = parser.parse(source, None)?;
let point = Point {
row: line.saturating_sub(1) as usize,
column: col.saturating_sub(1) as usize,
};
let root = tree.root_node();
let node = root.named_descendant_for_point_range(point, point)?;
let node = if is_ident_kind(node.kind()) {
node
} else {
ident_child_at(node, point)?
};
let name = node_text(node, source)?;
let is_import = ancestor_in(node, res.import_kinds());
let is_member = parent_in(node, res.member_kinds()) && !is_first_named_child(node);
let is_type = ancestor_in(node, res.type_kinds());
let is_call = is_callee(node, res.call_kinds());
Some(IdentRef {
name,
is_call,
is_member,
is_type,
is_import,
})
})
}
fn ident_child_at<'t>(node: Node<'t>, point: Point) -> Option<Node<'t>> {
let mut cursor = node.walk();
for child in node.named_children(&mut cursor) {
let s = child.start_position();
let e = child.end_position();
let contains = (s.row, s.column) <= (point.row, point.column)
&& (point.row, point.column) <= (e.row, e.column);
if contains {
if is_ident_kind(child.kind()) {
return Some(child);
}
if let Some(found) = ident_child_at(child, point) {
return Some(found);
}
}
}
None
}
fn node_text(node: Node, source: &[u8]) -> Option<String> {
let s = std::str::from_utf8(source.get(node.start_byte()..node.end_byte())?).ok()?;
let s = s.trim();
if s.is_empty() {
None
} else {
Some(s.to_string())
}
}
fn parent_in(node: Node, kinds: &[&str]) -> bool {
node.parent()
.map(|p| kinds.contains(&p.kind()))
.unwrap_or(false)
}
fn ancestor_in(node: Node, kinds: &[&str]) -> bool {
let mut cur = node.parent();
let mut hops = 0;
while let Some(p) = cur {
if kinds.contains(&p.kind()) {
return true;
}
hops += 1;
if hops > 8 {
break;
}
cur = p.parent();
}
false
}
fn is_first_named_child(node: Node) -> bool {
if let Some(p) = node.parent() {
let mut cursor = p.walk();
let first = p.named_children(&mut cursor).next();
if let Some(first) = first {
return first.id() == node.id();
}
}
false
}
fn is_callee(node: Node, call_kinds: &[&str]) -> bool {
let parent = match node.parent() {
Some(p) => p,
None => return false,
};
if call_kinds.contains(&parent.kind()) {
if let Some(f) = parent.child_by_field_name("function") {
return f.id() == node.id();
}
return is_first_named_child(node);
}
if let Some(grand) = parent.parent() {
if call_kinds.contains(&grand.kind()) {
let callee = grand
.child_by_field_name("function")
.or_else(|| grand.named_child(0));
if let Some(c) = callee {
return c.id() == parent.id() && !is_first_named_child(node);
}
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn finds_call_identifier() {
let src = b"fn main() {\n helper();\n}\n";
let r = identifier_at(Language::Rust, src, 2, 5).expect("ident");
assert_eq!(r.name, "helper");
assert!(r.is_call, "expected call: {r:?}");
assert!(!r.is_member);
}
#[test]
fn finds_member_method() {
let src = b"fn main() {\n obj.run();\n}\n";
let r = identifier_at(Language::Rust, src, 2, 9).expect("ident");
assert_eq!(r.name, "run");
assert!(r.is_member, "expected member: {r:?}");
}
#[test]
fn finds_import_identifier() {
let src = b"from os import path\n";
let r = identifier_at(Language::Python, src, 1, 16).expect("ident");
assert_eq!(r.name, "path");
assert!(r.is_import, "expected import: {r:?}");
}
}