type_sitter_lib/query/
captures.rs

1use crate::{raw, Query, UntypedNode};
2use std::fmt::Debug;
3use streaming_iterator::StreamingIterator;
4#[cfg(not(feature = "yak-sitter"))]
5use tree_sitter::Point;
6#[cfg(feature = "yak-sitter")]
7use yak_sitter::PointRange;
8
9/// Iterate a query's captures (see [tree-sitter's `QueryCaptures`](RawQueryCapture))
10#[cfg(feature = "yak-sitter")]
11pub struct QueryCaptures<'query, 'tree: 'query, Query: crate::Query> {
12    typed_query: &'query Query,
13    untyped_captures: raw::QueryCaptures<'query, 'tree>,
14}
15
16/// Iterate a query's captures (see [RawQueryCaptures])
17#[cfg(not(feature = "yak-sitter"))]
18pub struct QueryCaptures<
19    'query,
20    'tree: 'query,
21    Query: crate::Query,
22    Text: raw::TextProvider<I>,
23    I: AsRef<[u8]>,
24> {
25    typed_query: &'query Query,
26    untyped_captures: raw::QueryCaptures<'query, 'tree, Text, I>,
27}
28
29/// A capture from a [`Query`] with [typed nodes](crate::Node)
30pub trait QueryCapture<'query, 'tree: 'query>: Debug + Clone {
31    /// The type of query this capture came from
32    type Query: Query<Capture<'query, 'tree> = Self>;
33
34    /// The query this capture came from
35    fn query(&self) -> &'query Self::Query;
36
37    /// Get the raw tree-sitter query capture
38    #[cfg(feature = "yak-sitter")]
39    fn raw(&self) -> raw::QueryCapture<'query, 'tree>;
40
41    /// Get the raw tree-sitter query captures.
42    #[cfg(not(feature = "yak-sitter"))]
43    fn raw(&self) -> raw::QueryCapture<'tree>;
44
45    /// Get the captured untyped node
46    fn node(&self) -> &UntypedNode<'tree>;
47
48    /// Get a mutable reference to the captured untyped node
49    fn node_mut(&mut self) -> &mut UntypedNode<'tree>;
50
51    /// Get the capture name
52    fn name(&self) -> &'query str;
53
54    /// Get the capture index
55    // The cast is necessary in `tree-sitter` but not `yak-sitter`.
56    #[allow(clippy::unnecessary_cast)]
57    #[inline]
58    fn index(&self) -> usize {
59        self.raw().index as usize
60    }
61}
62
63#[cfg(feature = "yak-sitter")]
64impl<'query, 'tree: 'query, Query: crate::Query> QueryCaptures<'query, 'tree, Query> {
65    /// Construct from [tree-sitter's `QueryCaptures`](raw::QueryCaptures)
66    ///
67    /// # Safety
68    /// The captures must have come from the same query.
69    #[inline]
70    pub(super) unsafe fn new(
71        typed_query: &'query Query,
72        untyped_captures: raw::QueryCaptures<'query, 'tree>,
73    ) -> Self {
74        Self {
75            typed_query,
76            untyped_captures,
77        }
78    }
79
80    /// Limit captures to a byte range
81    #[inline]
82    pub fn set_byte_range(&mut self, range: std::ops::Range<usize>) {
83        self.untyped_captures.set_byte_range(range)
84    }
85
86    /// Limit captures to a point range
87    #[inline]
88    pub fn set_point_range(&mut self, range: PointRange) {
89        self.untyped_captures.set_point_range(range)
90    }
91}
92
93#[cfg(not(feature = "yak-sitter"))]
94impl<'query, 'tree: 'query, Query: crate::Query, Text: raw::TextProvider<I>, I: AsRef<[u8]>>
95    QueryCaptures<'query, 'tree, Query, Text, I>
96{
97    /// Construct from [tree-sitter's `QueryCaptures`](raw::QueryCaptures)
98    ///
99    /// # Safety
100    /// The captures must have come from the same query.
101    #[inline]
102    pub(super) unsafe fn new(
103        typed_query: &'query Query,
104        untyped_captures: raw::QueryCaptures<'query, 'tree, Text, I>,
105    ) -> Self {
106        Self {
107            typed_query,
108            untyped_captures,
109        }
110    }
111
112    /// Limit captures to a byte range
113    #[inline]
114    pub fn set_byte_range(&mut self, range: std::ops::Range<usize>) {
115        self.untyped_captures.set_byte_range(range)
116    }
117
118    /// Limit captures to a point range
119    #[inline]
120    pub fn set_point_range(&mut self, range: std::ops::Range<Point>) {
121        self.untyped_captures.set_point_range(range)
122    }
123}
124
125#[cfg(feature = "yak-sitter")]
126impl<'query, 'tree: 'query, Query: crate::Query> Iterator for QueryCaptures<'query, 'tree, Query> {
127    type Item = Query::Capture<'query, 'tree>;
128
129    //noinspection RsIncorrectFunctionArgumentCount -- IntelliJ inspection bug.
130    #[inline]
131    fn next(&mut self) -> Option<Self::Item> {
132        let query = self.untyped_captures.query();
133        let tree = self.untyped_captures.tree();
134
135        // SAFETY: Captures come from the same query and tree, and node from the same tree
136        unsafe {
137            self.untyped_captures
138                .as_inner_mut()
139                .next()
140                .map(|(m, index)| {
141                    let inner_capture = m.captures[*index];
142                    self.typed_query.wrap_capture(raw::QueryCapture {
143                        node: raw::Node::new(inner_capture.node, tree),
144                        index: inner_capture.index as usize,
145                        name: query.capture_names()[inner_capture.index as usize],
146                    })
147                })
148        }
149    }
150
151    #[inline]
152    fn size_hint(&self) -> (usize, Option<usize>) {
153        self.untyped_captures.size_hint()
154    }
155}
156
157#[cfg(not(feature = "yak-sitter"))]
158impl<'query, 'tree: 'query, Query: crate::Query, Text: raw::TextProvider<I>, I: AsRef<[u8]>>
159    Iterator for QueryCaptures<'query, 'tree, Query, Text, I>
160{
161    type Item = Query::Capture<'query, 'tree>;
162
163    #[inline]
164    fn next(&mut self) -> Option<Self::Item> {
165        // SAFETY: Captures come from the same query
166        unsafe {
167            self.untyped_captures
168                .next()
169                .map(|(m, index)| self.typed_query.wrap_capture(m.captures[*index]))
170        }
171    }
172
173    #[inline]
174    fn size_hint(&self) -> (usize, Option<usize>) {
175        self.untyped_captures.size_hint()
176    }
177}
type_sitter_lib/query/captures.rs

type_sitter_lib/query/
captures.rs
