sqruff_lib_core/utils/functional/
segments.rs

1use crate::dialects::syntax::SyntaxSet;
2use crate::parser::segments::ErasedSegment;
3use crate::templaters::TemplatedFile;
4
5#[derive(Debug, Default, Clone)]
6pub struct Segments {
7    pub base: Vec<ErasedSegment>,
8    templated_file: Option<TemplatedFile>,
9}
10
11impl Segments {
12    pub fn into_vec(self) -> Vec<ErasedSegment> {
13        self.base
14    }
15
16    pub fn iter(&self) -> impl Iterator<Item = &ErasedSegment> {
17        self.base.iter()
18    }
19
20    pub fn recursive_crawl(&self, types: &SyntaxSet, recurse_into: bool) -> Segments {
21        let mut segments = Vec::new();
22
23        for segment in &self.base {
24            segments.extend(segment.recursive_crawl(types, recurse_into, &SyntaxSet::EMPTY, true));
25        }
26
27        Segments::from_vec(segments, self.templated_file.clone())
28    }
29
30    pub fn iterate_segments(&self) -> impl Iterator<Item = Segments> + '_ {
31        let mut iter = self.base.iter();
32
33        std::iter::from_fn(move || {
34            let segment = iter.next()?;
35            Segments::new(segment.clone(), self.templated_file.clone()).into()
36        })
37    }
38
39    pub fn from_vec(base: Vec<ErasedSegment>, templated_file: Option<TemplatedFile>) -> Self {
40        Self {
41            base,
42            templated_file,
43        }
44    }
45
46    pub fn reversed(&self) -> Self {
47        let mut base = self.base.clone();
48        base.reverse();
49
50        Self {
51            base,
52            templated_file: self.templated_file.clone(),
53        }
54    }
55
56    pub fn get(&self, index: usize, default: Option<ErasedSegment>) -> Option<ErasedSegment> {
57        self.base.get(index).cloned().or(default)
58    }
59
60    pub fn first(&self) -> Option<&ErasedSegment> {
61        self.base.first()
62    }
63
64    pub fn last(&self) -> Option<&ErasedSegment> {
65        self.base.last()
66    }
67
68    #[track_caller]
69    pub fn pop(&mut self) -> ErasedSegment {
70        self.base.pop().unwrap()
71    }
72
73    pub fn all_match<F: Fn(&ErasedSegment) -> bool>(&self, predicate: F) -> bool {
74        self.base.iter().all(predicate)
75    }
76
77    pub fn any_match<F: Fn(&ErasedSegment) -> bool>(&self, predicate: F) -> bool {
78        self.base.iter().any(predicate)
79    }
80
81    pub fn len(&self) -> usize {
82        self.base.len()
83    }
84
85    pub fn is_empty(&self) -> bool {
86        self.base.is_empty()
87    }
88
89    pub fn new(segment: ErasedSegment, templated_file: Option<TemplatedFile>) -> Self {
90        Self {
91            base: vec![segment],
92            templated_file,
93        }
94    }
95
96    pub fn children_all(&self) -> Segments {
97        let mut child_segments = Vec::with_capacity(self.len());
98        for segment in &self.base {
99            child_segments.extend(segment.segments().iter().cloned());
100        }
101        Segments {
102            base: child_segments,
103            templated_file: self.templated_file.clone(),
104        }
105    }
106
107    pub fn children_where<F: Fn(&ErasedSegment) -> bool>(&self, predicate: F) -> Segments {
108        let mut child_segments = Vec::with_capacity(self.len());
109        for segment in &self.base {
110            for child in segment.segments() {
111                if predicate(child) {
112                    child_segments.push(child.clone());
113                }
114            }
115        }
116        Segments {
117            base: child_segments,
118            templated_file: self.templated_file.clone(),
119        }
120    }
121
122    pub fn children_iter_where<F: Fn(&ErasedSegment) -> bool + 'static>(
123        &self,
124        predicate: F,
125    ) -> impl Iterator<Item = &ErasedSegment> + '_ {
126        self.base
127            .iter()
128            .flat_map(|segment| segment.segments().iter())
129            .filter(move |child| predicate(child))
130    }
131
132    pub fn find_last_where<F: Fn(&ErasedSegment) -> bool>(&self, predicate: F) -> Segments {
133        self.base
134            .iter()
135            .rev()
136            .find_map(|segment| {
137                if predicate(segment) {
138                    Some(Segments {
139                        base: vec![segment.clone()],
140                        templated_file: self.templated_file.clone(),
141                    })
142                } else {
143                    None
144                }
145            })
146            .unwrap_or_else(|| Segments {
147                base: vec![],
148                templated_file: self.templated_file.clone(),
149            })
150    }
151
152    pub fn find(&self, value: &ErasedSegment) -> Option<usize> {
153        self.index(value)
154    }
155
156    pub fn find_first_where<F: Fn(&ErasedSegment) -> bool>(&self, predicate: F) -> Segments {
157        for segment in &self.base {
158            if predicate(segment) {
159                return Segments {
160                    base: vec![segment.clone()],
161                    templated_file: self.templated_file.clone(),
162                };
163            }
164        }
165        Segments {
166            base: vec![],
167            templated_file: self.templated_file.clone(),
168        }
169    }
170
171    pub fn head(&self) -> Segments {
172        if let Some(segment) = self.base.first() {
173            return Segments {
174                base: vec![segment.clone()],
175                templated_file: self.templated_file.clone(),
176            };
177        }
178        Segments {
179            base: vec![],
180            templated_file: self.templated_file.clone(),
181        }
182    }
183
184    pub fn index(&self, value: &ErasedSegment) -> Option<usize> {
185        self.base.iter().position(|it| it == value)
186    }
187
188    pub fn filter<F: Fn(&ErasedSegment) -> bool>(&self, predicate: F) -> Segments {
189        let base = self
190            .base
191            .iter()
192            .filter(|segment| predicate(segment))
193            .cloned()
194            .collect();
195        Segments {
196            base,
197            templated_file: self.templated_file.clone(),
198        }
199    }
200
201    pub fn take_while<F: Fn(&ErasedSegment) -> bool>(&self, predicate: F) -> Segments {
202        let base = self
203            .base
204            .iter()
205            .take_while(|segment| predicate(segment))
206            .cloned()
207            .collect();
208        Segments {
209            base,
210            templated_file: self.templated_file.clone(),
211        }
212    }
213
214    pub fn iter_after<'a>(
215        &'a self,
216        start: &ErasedSegment,
217    ) -> impl Iterator<Item = &'a ErasedSegment> + 'a {
218        let start_idx = self
219            .base
220            .iter()
221            .position(|segment| segment == start)
222            .map(|index| index + 1)
223            .unwrap_or(self.base.len());
224        self.base[start_idx..].iter()
225    }
226
227    pub fn iter_after_while<'a, F: Fn(&ErasedSegment) -> bool + 'a>(
228        &'a self,
229        start: &ErasedSegment,
230        while_cond: F,
231    ) -> impl Iterator<Item = &'a ErasedSegment> + 'a {
232        self.iter_after(start)
233            .take_while(move |segment| while_cond(segment))
234    }
235
236    pub fn after(&self, start: &ErasedSegment) -> Segments {
237        let start_idx = self
238            .base
239            .iter()
240            .position(|segment| segment == start)
241            .expect("start segment not found");
242        let base = self.base[start_idx + 1..].to_vec();
243        Segments {
244            base,
245            templated_file: self.templated_file.clone(),
246        }
247    }
248
249    pub fn before(&self, stop: &ErasedSegment) -> Segments {
250        let stop_idx = self
251            .base
252            .iter()
253            .position(|segment| segment == stop)
254            .expect("stop segment not found");
255        let base = self.base[..stop_idx].to_vec();
256        Segments {
257            base,
258            templated_file: self.templated_file.clone(),
259        }
260    }
261
262    pub fn between_exclusive(&self, start: &ErasedSegment, stop: &ErasedSegment) -> Segments {
263        let len = self.base.len();
264
265        let start_idx = self
266            .base
267            .iter()
268            .position(|s| s == start)
269            .map(|i| i + 1)
270            .unwrap_or(0)
271            .min(len);
272
273        let end_idx = self
274            .base
275            .iter()
276            .position(|s| s == stop)
277            .unwrap_or(len)
278            .min(len);
279
280        let base = self.base.get(start_idx..end_idx).unwrap_or(&[]).to_vec();
281
282        Segments {
283            base,
284            templated_file: self.templated_file.clone(),
285        }
286    }
287}
288
289impl<I: std::slice::SliceIndex<[ErasedSegment]>> std::ops::Index<I> for Segments {
290    type Output = I::Output;
291
292    fn index(&self, index: I) -> &Self::Output {
293        &self.base[index]
294    }
295}
296
297impl IntoIterator for Segments {
298    type Item = ErasedSegment;
299    type IntoIter = std::vec::IntoIter<Self::Item>;
300
301    fn into_iter(self) -> Self::IntoIter {
302        self.base.into_iter()
303    }
304}
sqruff_lib_core/utils/functional/segments.rs

sqruff_lib_core/utils/functional/
segments.rs
