tree_house/
text_object.rs

1// TODO: rework using query iter
2
3use std::iter;
4
5use ropey::RopeSlice;
6
7use crate::TREE_SITTER_MATCH_LIMIT;
8use tree_sitter::{InactiveQueryCursor, Node, Query, RopeInput};
9
10#[derive(Debug)]
11pub enum CapturedNode<'a> {
12    Single(Node<'a>),
13    /// Guaranteed to be not empty
14    Grouped(Vec<Node<'a>>),
15}
16
17impl CapturedNode<'_> {
18    pub fn start_byte(&self) -> usize {
19        match self {
20            Self::Single(n) => n.start_byte() as usize,
21            Self::Grouped(ns) => ns[0].start_byte() as usize,
22        }
23    }
24
25    pub fn end_byte(&self) -> usize {
26        match self {
27            Self::Single(n) => n.end_byte() as usize,
28            Self::Grouped(ns) => ns.last().unwrap().end_byte() as usize,
29        }
30    }
31}
32
33#[derive(Debug)]
34pub struct TextObjectQuery {
35    pub query: Query,
36}
37
38impl TextObjectQuery {
39    /// Run the query on the given node and return sub nodes which match given
40    /// capture ("function.inside", "class.around", etc).
41    ///
42    /// Captures may contain multiple nodes by using quantifiers (+, *, etc),
43    /// and support for this is partial and could use improvement.
44    ///
45    /// ```query
46    /// (comment)+ @capture
47    ///
48    /// ; OR
49    /// (
50    ///   (comment)*
51    ///   .
52    ///   (function)
53    /// ) @capture
54    /// ```
55    pub fn capture_nodes<'a>(
56        &'a self,
57        capture_name: &str,
58        node: Node<'a>,
59        slice: RopeSlice<'a>,
60        cursor: InactiveQueryCursor,
61    ) -> Option<impl Iterator<Item = CapturedNode<'a>>> {
62        self.capture_nodes_any(&[capture_name], node, slice, cursor)
63    }
64
65    /// Find the first capture that exists out of all given `capture_names`
66    /// and return sub nodes that match this capture.
67    pub fn capture_nodes_any<'a>(
68        &'a self,
69        capture_names: &[&str],
70        node: Node<'a>,
71        slice: RopeSlice<'a>,
72        mut cursor: InactiveQueryCursor,
73    ) -> Option<impl Iterator<Item = CapturedNode<'a>>> {
74        let capture = capture_names
75            .iter()
76            .find_map(|cap| self.query.get_capture(cap))?;
77
78        cursor.set_match_limit(TREE_SITTER_MATCH_LIMIT);
79        let mut cursor = cursor.execute_query(&self.query, &node, RopeInput::new(slice));
80        let capture_node = iter::from_fn(move || {
81            let (mat, _) = cursor.next_matched_node()?;
82            Some(mat.nodes_for_capture(capture).cloned().collect())
83        })
84        .filter_map(move |nodes: Vec<_>| {
85            if nodes.len() > 1 {
86                Some(CapturedNode::Grouped(nodes))
87            } else {
88                nodes.into_iter().map(CapturedNode::Single).next()
89            }
90        });
91        Some(capture_node)
92    }
93}
tree_house/text_object.rs

tree_house/
text_object.rs
