Skip to content

Commit ee3801d

Browse files
committed
Rewrite method matching with TS query
1 parent 3d4b53e commit ee3801d

File tree

2 files changed

+70
-52
lines changed

2 files changed

+70
-52
lines changed

crates/ark/src/lsp/indexer.rs

Lines changed: 30 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -314,61 +314,39 @@ fn index_r6_class(
314314
node: &Node,
315315
entries: &mut Vec<IndexEntry>,
316316
) -> anyhow::Result<()> {
317-
let Some(args_node) = node.child_by_field_name("arguments") else {
318-
return Ok(());
319-
};
320-
321-
let mut cursor = args_node.walk();
322-
for arg in args_node.children(&mut cursor) {
323-
// Only consider `public = ` and `private = ` arguments
324-
let Some(arg_name) = arg.child_by_field_name("name") else {
325-
continue;
326-
};
327-
if !arg_name.is_identifier_or_string() {
328-
continue;
329-
}
330-
let arg_name_str = contents.node_slice(&arg_name)?;
331-
if arg_name_str != "public" && arg_name_str != "private" {
332-
continue;
333-
}
334-
335-
let Some(list_node) = arg.child_by_field_name("value") else {
336-
continue;
337-
};
338-
if !list_node.is_call() {
339-
continue;
340-
}
341-
342-
let Some(list_args) = list_node.child_by_field_name("arguments") else {
343-
return Ok(());
344-
};
345-
346-
let mut cursor = list_args.walk();
347-
for arg in list_args.children(&mut cursor) {
348-
if !arg.is_argument() {
349-
continue;
350-
}
317+
// Tree-sitter query to match individual methods in R6Class public/private lists
318+
let query_str = r#"
319+
(argument
320+
name: (identifier) @access
321+
value: (call
322+
function: (identifier) @_list_fn
323+
arguments: (arguments
324+
(argument
325+
name: (identifier) @method_name
326+
value: (function_definition) @method_fn
327+
)
328+
)
329+
)
330+
(#match? @access "public|private")
331+
(#eq? @_list_fn "list")
332+
)
333+
"#;
334+
let mut ts_query = crate::treesitter::TSQuery::new(query_str)?;
351335

352-
let (Some(mtd_name), Some(mtd_value)) = (
353-
arg.child_by_field_name("name"),
354-
arg.child_by_field_name("value"),
355-
) else {
356-
continue;
357-
};
358-
if !mtd_name.is_identifier_or_string() || !mtd_value.is_function_definition() {
359-
continue;
360-
}
336+
// We'll switch from Rope to String in the near future so let's not
337+
// worry about this conversion now
338+
let contents_str = contents.to_string();
361339

362-
let name = contents.node_slice(&mtd_name)?.to_string();
363-
let start = convert_point_to_position(contents, mtd_name.start_position());
364-
let end = convert_point_to_position(contents, mtd_name.end_position());
340+
for method_node in ts_query.captures_for(*node, "method_name", contents_str.as_bytes()) {
341+
let name = contents.node_slice(&method_node)?.to_string();
342+
let start = convert_point_to_position(contents, method_node.start_position());
343+
let end = convert_point_to_position(contents, method_node.end_position());
365344

366-
entries.push(IndexEntry {
367-
key: name.clone(),
368-
range: Range { start, end },
369-
data: IndexEntryData::Method { name },
370-
});
371-
}
345+
entries.push(IndexEntry {
346+
key: name.clone(),
347+
range: Range { start, end },
348+
data: IndexEntryData::Method { name },
349+
});
372350
}
373351

374352
Ok(())

crates/ark/src/treesitter.rs

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use anyhow::anyhow;
12
use tree_sitter::Node;
23

34
use crate::lsp::traits::node::NodeExt;
@@ -603,3 +604,42 @@ pub(crate) fn node_find_containing_call<'tree>(node: Node<'tree>) -> Option<Node
603604

604605
None
605606
}
607+
608+
pub(crate) struct TSQuery {
609+
query: tree_sitter::Query,
610+
cursor: tree_sitter::QueryCursor,
611+
}
612+
613+
impl TSQuery {
614+
pub(crate) fn new(query_str: &str) -> anyhow::Result<Self> {
615+
let language = &tree_sitter_r::LANGUAGE.into();
616+
let query = tree_sitter::Query::new(language, query_str)
617+
.map_err(|err| anyhow!("Failed to compile query: {err}"))?;
618+
619+
let cursor = tree_sitter::QueryCursor::new();
620+
621+
Ok(Self { query, cursor })
622+
}
623+
624+
/// Match query against `contents` and collect all nodes captured with the
625+
/// given capture name
626+
pub(crate) fn captures_for<'tree>(
627+
&mut self,
628+
node: tree_sitter::Node<'tree>,
629+
capture_name: &str,
630+
contents: &[u8],
631+
) -> Vec<tree_sitter::Node<'tree>> {
632+
let mut result = Vec::new();
633+
634+
for m in self.cursor.matches(&self.query, node, contents) {
635+
for cap in m.captures.iter() {
636+
let cap_name = &self.query.capture_names()[cap.index as usize];
637+
if *cap_name == capture_name {
638+
result.push(cap.node);
639+
}
640+
}
641+
}
642+
643+
result
644+
}
645+
}

0 commit comments

Comments
 (0)