Tweak QueryCursor to allow iterating either matches or captures

For syntax highlighting, we want to iterate over all of the captures in 
order, and don't care about grouping the captures by pattern.
This commit is contained in:
Max Brunsfeld 2019-09-11 14:44:49 -07:00
parent 33587c924a
commit a1fec71b19
11 changed files with 559 additions and 227 deletions

View file

@ -607,19 +607,23 @@ extern "C" {
#[doc = " Create a new cursor for executing a given query."]
#[doc = ""]
#[doc = " The cursor stores the state that is needed to iteratively search"]
#[doc = " for matches. To use the query cursor:"]
#[doc = " 1. First call `ts_query_cursor_exec` to start running a given query on"]
#[doc = "a given syntax node."]
#[doc = " 2. Then repeatedly call `ts_query_cursor_next` to iterate over the matches."]
#[doc = " This will return `false` when there are no more matches left."]
#[doc = " 3. After each successful call to `ts_query_cursor_next`, you can call"]
#[doc = " `ts_query_cursor_matched_pattern_index` to determine which pattern"]
#[doc = " matched. You can also call `ts_query_cursor_matched_captures` to"]
#[doc = " determine which nodes were captured, and by which capture names."]
#[doc = " for matches. To use the query cursor, first call `ts_query_cursor_exec`"]
#[doc = " to start running a given query on a given syntax node. Then, there are"]
#[doc = " two options for consuming the results of the query:"]
#[doc = " 1. Repeatedly call `ts_query_cursor_next_match` to iterate over all of the"]
#[doc = " the *matches* in the order that they were found. Each match contains the"]
#[doc = " index of the pattern that matched, and an array of captures. Because"]
#[doc = " multiple patterns can match the same set of nodes, one match may contain"]
#[doc = " captures that appear *before* some of the captures from a previous match."]
#[doc = " 2. Repeatedly call `ts_query_cursor_next_capture` to iterate over all of the"]
#[doc = " individual *captures* in the order that they appear. This is useful if"]
#[doc = " don\'t care about which pattern matched, and just want a single ordered"]
#[doc = " sequence of captures."]
#[doc = ""]
#[doc = " If you don\'t care about finding all of the matches, you can stop calling"]
#[doc = " `ts_query_cursor_next` at any point. And you can start executing another"]
#[doc = " query on another node by calling `ts_query_cursor_exec` again."]
#[doc = " If you don\'t care about consuming all of the results, you can stop calling"]
#[doc = " `ts_query_cursor_next_match` or `ts_query_cursor_next_capture` at any point."]
#[doc = " You can then start executing another query on another node by calling"]
#[doc = " `ts_query_cursor_exec` again."]
pub fn ts_query_cursor_new() -> *mut TSQueryCursor;
}
extern "C" {
@ -640,18 +644,26 @@ extern "C" {
}
extern "C" {
#[doc = " Advance to the next match of the currently running query."]
pub fn ts_query_cursor_next(arg1: *mut TSQueryCursor) -> bool;
#[doc = ""]
#[doc = " If there is another match, write its pattern index to `pattern_index`,"]
#[doc = " the number of captures to `capture_count`, and the captures themselves"]
#[doc = " to `*captures`, and return `true`. Otherwise, return `false`."]
pub fn ts_query_cursor_next_match(
self_: *mut TSQueryCursor,
pattern_index: *mut u32,
capture_count: *mut u32,
captures: *mut *const TSQueryCapture,
) -> bool;
}
extern "C" {
#[doc = " Check which pattern matched."]
pub fn ts_query_cursor_matched_pattern_index(arg1: *const TSQueryCursor) -> u32;
}
extern "C" {
#[doc = " Check which pattern matched."]
pub fn ts_query_cursor_matched_captures(
arg1: *const TSQueryCursor,
arg2: *mut u32,
) -> *const TSQueryCapture;
#[doc = " Advance to the next capture of the currently running query."]
#[doc = ""]
#[doc = " If there is another capture, write it to `capture` and return `true`."]
#[doc = " Otherwise, return `false`."]
pub fn ts_query_cursor_next_capture(
arg1: *mut TSQueryCursor,
capture: *mut TSQueryCapture,
) -> bool;
}
extern "C" {
#[doc = " Get the number of distinct node types in the language."]

View file

@ -18,6 +18,7 @@ use std::marker::PhantomData;
use std::os::raw::{c_char, c_void};
use std::sync::atomic::AtomicUsize;
use std::{char, fmt, ptr, slice, str, u16};
use std::mem::MaybeUninit;
pub const LANGUAGE_VERSION: usize = ffi::TREE_SITTER_LANGUAGE_VERSION;
pub const PARSER_HEADER: &'static str = include_str!("../include/tree_sitter/parser.h");
@ -144,7 +145,12 @@ pub struct Query {
pub struct QueryCursor(*mut ffi::TSQueryCursor);
pub struct QueryMatch<'a>(*mut ffi::TSQueryCursor, PhantomData<&'a ()>);
pub struct QueryMatch<'a> {
pattern_index: usize,
capture_count: usize,
captures_ptr: *const ffi::TSQueryCapture,
cursor: PhantomData<&'a ()>,
}
#[derive(Debug, PartialEq, Eq)]
pub enum QueryError<'a> {
@ -996,14 +1002,52 @@ impl QueryCursor {
QueryCursor(unsafe { ffi::ts_query_cursor_new() })
}
pub fn exec<'a>(&'a mut self, query: &'a Query, node: Node<'a>) -> impl Iterator<Item = QueryMatch<'a>> + 'a {
pub fn matches<'a>(
&'a mut self,
query: &'a Query,
node: Node<'a>,
) -> impl Iterator<Item = QueryMatch<'a>> + 'a {
unsafe {
ffi::ts_query_cursor_exec(self.0, query.ptr, node.0);
}
std::iter::from_fn(move || -> Option<QueryMatch<'a>> {
unsafe {
if ffi::ts_query_cursor_next(self.0) {
Some(QueryMatch(self.0, PhantomData))
let mut pattern_index = 0u32;
let mut capture_count = 0u32;
let mut captures = ptr::null();
if ffi::ts_query_cursor_next_match(
self.0,
&mut pattern_index as *mut u32,
&mut capture_count as *mut u32,
&mut captures as *mut *const ffi::TSQueryCapture,
) {
Some(QueryMatch {
pattern_index: pattern_index as usize,
capture_count: capture_count as usize,
captures_ptr: captures,
cursor: PhantomData
})
} else {
None
}
}
})
}
pub fn captures<'a>(
&'a mut self,
query: &'a Query,
node: Node<'a>,
) -> impl Iterator<Item = (usize, Node)> + 'a {
unsafe {
ffi::ts_query_cursor_exec(self.0, query.ptr, node.0);
}
std::iter::from_fn(move || -> Option<(usize, Node<'a>)> {
unsafe {
let mut capture = MaybeUninit::<ffi::TSQueryCapture>::uninit();
if ffi::ts_query_cursor_next_capture(self.0, capture.as_mut_ptr()) {
let capture = capture.assume_init();
Some((capture.index as usize, Node::new(capture.node).unwrap()))
} else {
None
}
@ -1028,19 +1072,14 @@ impl QueryCursor {
impl<'a> QueryMatch<'a> {
pub fn pattern_index(&self) -> usize {
unsafe { ffi::ts_query_cursor_matched_pattern_index(self.0) as usize }
self.pattern_index
}
pub fn captures(&self) -> impl ExactSizeIterator<Item = (usize, Node)> {
unsafe {
let mut capture_count = 0u32;
let captures =
ffi::ts_query_cursor_matched_captures(self.0, &mut capture_count as *mut u32);
let captures = slice::from_raw_parts(captures, capture_count as usize);
captures
.iter()
.map(move |capture| (capture.index as usize, Node::new(capture.node).unwrap()))
}
let captures = unsafe { slice::from_raw_parts(self.captures_ptr, self.capture_count as usize) };
captures
.iter()
.map(|capture| (capture.index as usize, Node::new(capture.node).unwrap()))
}
}