1use std::cell::RefCell;
14use std::collections::HashMap;
15
16use tree_sitter::{Node, Parser, Point};
17
18use crate::lang::Language;
19
20thread_local! {
21 static RESOLVE_PARSERS: RefCell<HashMap<Language, Parser>> = RefCell::new(HashMap::new());
22}
23
24#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct IdentRef {
27 pub name: String,
28 pub is_call: bool,
30 pub is_member: bool,
32 pub is_type: bool,
34 pub is_import: bool,
36}
37
38pub trait LangResolver: Send + Sync {
41 fn language(&self) -> Language;
42
43 fn member_kinds(&self) -> &'static [&'static str];
46
47 fn call_kinds(&self) -> &'static [&'static str];
49
50 fn type_kinds(&self) -> &'static [&'static str] {
52 &["type_identifier", "type_annotation", "generic_type", "type"]
53 }
54
55 fn import_kinds(&self) -> &'static [&'static str];
57}
58
59macro_rules! lang_resolver {
60 ($name:ident, $lang:expr, member = $member:expr, call = $call:expr, import = $import:expr) => {
61 struct $name;
62 impl LangResolver for $name {
63 fn language(&self) -> Language {
64 $lang
65 }
66 fn member_kinds(&self) -> &'static [&'static str] {
67 $member
68 }
69 fn call_kinds(&self) -> &'static [&'static str] {
70 $call
71 }
72 fn import_kinds(&self) -> &'static [&'static str] {
73 $import
74 }
75 }
76 };
77}
78
79lang_resolver!(
80 RustResolver,
81 Language::Rust,
82 member = &["field_expression"],
83 call = &["call_expression", "macro_invocation"],
84 import = &["use_declaration"]
85);
86lang_resolver!(
87 PythonResolver,
88 Language::Python,
89 member = &["attribute"],
90 call = &["call"],
91 import = &["import_statement", "import_from_statement"]
92);
93lang_resolver!(
94 JsResolver,
95 Language::JavaScript,
96 member = &["member_expression"],
97 call = &["call_expression", "new_expression"],
98 import = &["import_statement"]
99);
100lang_resolver!(
101 TsResolver,
102 Language::TypeScript,
103 member = &["member_expression"],
104 call = &["call_expression", "new_expression"],
105 import = &["import_statement"]
106);
107lang_resolver!(
108 TsxResolver,
109 Language::Tsx,
110 member = &["member_expression"],
111 call = &["call_expression", "new_expression"],
112 import = &["import_statement"]
113);
114lang_resolver!(
115 GoResolver,
116 Language::Go,
117 member = &["selector_expression"],
118 call = &["call_expression"],
119 import = &["import_spec"]
120);
121lang_resolver!(
122 DartResolver,
123 Language::Dart,
124 member = &["member_expression"],
125 call = &["call_expression", "constructor_invocation"],
126 import = &["library_import", "import_specification"]
127);
128
129pub fn resolver_for(lang: Language) -> Box<dyn LangResolver> {
131 match lang {
132 Language::Rust => Box::new(RustResolver),
133 Language::Python => Box::new(PythonResolver),
134 Language::JavaScript => Box::new(JsResolver),
135 Language::TypeScript => Box::new(TsResolver),
136 Language::Tsx => Box::new(TsxResolver),
137 Language::Go => Box::new(GoResolver),
138 Language::Dart => Box::new(DartResolver),
139 other => Box::new(GenericResolver(other)),
140 }
141}
142
143struct GenericResolver(Language);
145impl LangResolver for GenericResolver {
146 fn language(&self) -> Language {
147 self.0
148 }
149 fn member_kinds(&self) -> &'static [&'static str] {
150 &[
151 "field_expression",
152 "member_expression",
153 "member_access_expression",
154 "selector_expression",
155 "attribute",
156 "scoped_identifier",
157 ]
158 }
159 fn call_kinds(&self) -> &'static [&'static str] {
160 &[
161 "call_expression",
162 "call",
163 "method_invocation",
164 "invocation_expression",
165 "function_call_expression",
166 "member_call_expression",
167 "object_creation_expression",
168 ]
169 }
170 fn import_kinds(&self) -> &'static [&'static str] {
171 &[
172 "use_declaration",
173 "import_statement",
174 "import_from_statement",
175 "import_declaration",
176 "import_spec",
177 "using_directive",
178 "namespace_use_declaration",
179 ]
180 }
181}
182
183fn is_ident_kind(kind: &str) -> bool {
184 kind.ends_with("identifier") || kind == "name" || kind == "constant" || kind == "property"
185}
186
187pub fn identifier_at(lang: Language, source: &[u8], line: u32, col: u32) -> Option<IdentRef> {
189 let grammar = lang.grammar()?;
190 let res = resolver_for(lang);
191 RESOLVE_PARSERS.with(|cell| {
192 let mut map = cell.borrow_mut();
193 let parser = map.entry(lang).or_insert_with(|| {
194 let mut p = Parser::new();
195 let _ = p.set_language(&grammar);
196 p
197 });
198 let tree = parser.parse(source, None)?;
199 let point = Point {
200 row: line.saturating_sub(1) as usize,
201 column: col.saturating_sub(1) as usize,
202 };
203 let root = tree.root_node();
204 let node = root.named_descendant_for_point_range(point, point)?;
206 let node = if is_ident_kind(node.kind()) {
207 node
208 } else {
209 ident_child_at(node, point)?
211 };
212 let name = node_text(node, source)?;
213
214 let is_import = ancestor_in(node, res.import_kinds());
215 let is_member = parent_in(node, res.member_kinds()) && !is_first_named_child(node);
216 let is_type = ancestor_in(node, res.type_kinds());
217 let is_call = is_callee(node, res.call_kinds());
218
219 Some(IdentRef {
220 name,
221 is_call,
222 is_member,
223 is_type,
224 is_import,
225 })
226 })
227}
228
229fn ident_child_at<'t>(node: Node<'t>, point: Point) -> Option<Node<'t>> {
231 let mut cursor = node.walk();
232 for child in node.named_children(&mut cursor) {
233 let s = child.start_position();
234 let e = child.end_position();
235 let contains = (s.row, s.column) <= (point.row, point.column)
236 && (point.row, point.column) <= (e.row, e.column);
237 if contains {
238 if is_ident_kind(child.kind()) {
239 return Some(child);
240 }
241 if let Some(found) = ident_child_at(child, point) {
242 return Some(found);
243 }
244 }
245 }
246 None
247}
248
249fn node_text(node: Node, source: &[u8]) -> Option<String> {
250 let s = std::str::from_utf8(source.get(node.start_byte()..node.end_byte())?).ok()?;
251 let s = s.trim();
252 if s.is_empty() {
253 None
254 } else {
255 Some(s.to_string())
256 }
257}
258
259fn parent_in(node: Node, kinds: &[&str]) -> bool {
260 node.parent()
261 .map(|p| kinds.contains(&p.kind()))
262 .unwrap_or(false)
263}
264
265fn ancestor_in(node: Node, kinds: &[&str]) -> bool {
266 let mut cur = node.parent();
267 let mut hops = 0;
268 while let Some(p) = cur {
269 if kinds.contains(&p.kind()) {
270 return true;
271 }
272 hops += 1;
273 if hops > 8 {
274 break;
275 }
276 cur = p.parent();
277 }
278 false
279}
280
281fn is_first_named_child(node: Node) -> bool {
282 if let Some(p) = node.parent() {
283 let mut cursor = p.walk();
284 let first = p.named_children(&mut cursor).next();
285 if let Some(first) = first {
286 return first.id() == node.id();
287 }
288 }
289 false
290}
291
292fn is_callee(node: Node, call_kinds: &[&str]) -> bool {
295 let parent = match node.parent() {
296 Some(p) => p,
297 None => return false,
298 };
299 if call_kinds.contains(&parent.kind()) {
300 if let Some(f) = parent.child_by_field_name("function") {
302 return f.id() == node.id();
303 }
304 return is_first_named_child(node);
305 }
306 if let Some(grand) = parent.parent() {
309 if call_kinds.contains(&grand.kind()) {
310 let callee = grand
311 .child_by_field_name("function")
312 .or_else(|| grand.named_child(0));
313 if let Some(c) = callee {
314 return c.id() == parent.id() && !is_first_named_child(node);
315 }
316 }
317 }
318 false
319}
320
321#[cfg(test)]
322mod tests {
323 use super::*;
324
325 #[test]
326 fn finds_call_identifier() {
327 let src = b"fn main() {\n helper();\n}\n";
328 let r = identifier_at(Language::Rust, src, 2, 5).expect("ident");
329 assert_eq!(r.name, "helper");
330 assert!(r.is_call, "expected call: {r:?}");
331 assert!(!r.is_member);
332 }
333
334 #[test]
335 fn finds_member_method() {
336 let src = b"fn main() {\n obj.run();\n}\n";
337 let r = identifier_at(Language::Rust, src, 2, 9).expect("ident");
339 assert_eq!(r.name, "run");
340 assert!(r.is_member, "expected member: {r:?}");
341 }
342
343 #[test]
344 fn finds_import_identifier() {
345 let src = b"from os import path\n";
346 let r = identifier_at(Language::Python, src, 1, 16).expect("ident");
347 assert_eq!(r.name, "path");
348 assert!(r.is_import, "expected import: {r:?}");
349 }
350}