avltriee/
iter.rs

1use std::{cmp::Ordering, num::NonZeroU32};
2
3use crate::{search::AvltrieeSearch, AvltrieeAllocator};
4
5use super::Avltriee;
6
7#[derive(PartialEq)]
8enum Order {
9    Asc,
10    Desc,
11}
12
13pub struct AvltrieeIter<'a, T, I: ?Sized, A> {
14    now: Option<NonZeroU32>,
15    end_row: Option<NonZeroU32>,
16    same_branch: Option<NonZeroU32>,
17    triee: &'a Avltriee<T, I, A>,
18    next_func: fn(
19        &Avltriee<T, I, A>,
20        NonZeroU32,
21        Option<NonZeroU32>,
22    ) -> Option<(NonZeroU32, Option<NonZeroU32>)>,
23}
24
25impl<'a, T, I: ?Sized, A: AvltrieeAllocator<T>> AvltrieeIter<'a, T, I, A> {
26    fn new(
27        triee: &'a Avltriee<T, I, A>,
28        now: Option<NonZeroU32>,
29        end_row: Option<NonZeroU32>,
30        order: Order,
31    ) -> AvltrieeIter<'a, T, I, A> {
32        match order {
33            Order::Asc => AvltrieeIter {
34                now,
35                end_row,
36                same_branch: None,
37                triee,
38                next_func: Avltriee::<T, I, A>::next,
39            },
40            Order::Desc => AvltrieeIter {
41                now: end_row,
42                end_row: now,
43                same_branch: None,
44                triee,
45                next_func: Avltriee::<T, I, A>::next_desc,
46            },
47        }
48    }
49
50    /// Generates an iterator of nodes with the same value as the specified value.
51    pub fn by<S: AvltrieeSearch<T, I, A>>(s: &'a S, value: &I) -> AvltrieeIter<'a, T, I, A> {
52        let triee = s.as_ref();
53        let edge = s.edge(value);
54        let row = if edge.1 == Ordering::Equal {
55            edge.0
56        } else {
57            None
58        };
59        AvltrieeIter::new(triee, row, row, Order::Asc)
60    }
61
62    /// Generates an iterator that is greater than or equal to the specified value.
63    pub fn from_asc<S: AvltrieeSearch<T, I, A>>(s: &'a S, value: &I) -> AvltrieeIter<'a, T, I, A> {
64        Self::from_inner(s, value, Order::Asc)
65    }
66
67    /// Generates a descending iterator with greater than or equal to the specified value.
68    pub fn from_desc<S: AvltrieeSearch<T, I, A>>(s: &'a S, value: &I) -> AvltrieeIter<'a, T, I, A> {
69        Self::from_inner(s, value, Order::Desc)
70    }
71
72    fn from_inner<S: AvltrieeSearch<T, I, A>>(
73        s: &'a S,
74        value: &I,
75        order: Order,
76    ) -> AvltrieeIter<'a, T, I, A> {
77        let triee = s.as_ref();
78        let now = s.ge(value);
79        AvltrieeIter::new(triee, now, now.and_then(|_| triee.max(triee.root())), order)
80    }
81
82    /// Generates an iterator of nodes with values less than or equal to the specified value.
83    pub fn to_asc<S: AvltrieeSearch<T, I, A>>(s: &'a S, value: &I) -> AvltrieeIter<'a, T, I, A> {
84        Self::to_inner(s, value, Order::Asc)
85    }
86
87    /// Generates an iterator of nodes with values less than or equal to the specified value. Iterates in descending order.
88    pub fn to_desc<S: AvltrieeSearch<T, I, A>>(s: &'a S, value: &I) -> AvltrieeIter<'a, T, I, A> {
89        Self::to_inner(s, value, Order::Desc)
90    }
91
92    fn to_inner<S: AvltrieeSearch<T, I, A>>(
93        s: &'a S,
94        value: &I,
95        order: Order,
96    ) -> AvltrieeIter<'a, T, I, A> {
97        let triee = s.as_ref();
98        let end_row = s.le(value);
99        AvltrieeIter::new(
100            triee,
101            end_row.and_then(|_| triee.min(triee.root())),
102            end_row,
103            order,
104        )
105    }
106
107    /// Generates an iterator of nodes with values greater than the specified value.
108    pub fn over_asc<S: AvltrieeSearch<T, I, A>>(s: &'a S, value: &I) -> AvltrieeIter<'a, T, I, A> {
109        Self::over_inner(s, value, Order::Asc)
110    }
111
112    /// Generates an iterator of nodes with values greater than the specified value. Iterates in descending order.
113    pub fn over_desc<S: AvltrieeSearch<T, I, A>>(s: &'a S, value: &I) -> AvltrieeIter<'a, T, I, A> {
114        Self::over_inner(s, value, Order::Desc)
115    }
116
117    fn over_inner<S: AvltrieeSearch<T, I, A>>(
118        s: &'a S,
119        value: &I,
120        order: Order,
121    ) -> AvltrieeIter<'a, T, I, A> {
122        let triee = s.as_ref();
123        let now = s.gt(value);
124        AvltrieeIter::new(triee, now, now.and_then(|_| triee.max(triee.root())), order)
125    }
126
127    /// Generates an iterator of nodes with values less than the specified value.
128    pub fn under_asc<S: AvltrieeSearch<T, I, A>>(s: &'a S, value: &I) -> AvltrieeIter<'a, T, I, A> {
129        Self::under_inner(s, value, Order::Asc)
130    }
131
132    /// Generates an iterator of nodes with values less than the specified value. Iterates in descending order.
133    pub fn under_desc<S: AvltrieeSearch<T, I, A>>(
134        s: &'a S,
135        value: &I,
136    ) -> AvltrieeIter<'a, T, I, A> {
137        Self::under_inner(s, value, Order::Desc)
138    }
139
140    fn under_inner<S: AvltrieeSearch<T, I, A>>(
141        s: &'a S,
142        value: &I,
143        order: Order,
144    ) -> AvltrieeIter<'a, T, I, A> {
145        let triee = s.as_ref();
146        let end_row = s.lt(value);
147        AvltrieeIter::new(
148            triee,
149            end_row.and_then(|_| triee.min(triee.root())),
150            end_row,
151            order,
152        )
153    }
154
155    /// Generates an iterator of nodes with the specified range of values.
156    pub fn range_asc<S: AvltrieeSearch<T, I, A>>(
157        s: &'a S,
158        start: &I,
159        end: &I,
160    ) -> AvltrieeIter<'a, T, I, A> {
161        Self::range_inner(s, start, end, Order::Asc)
162    }
163
164    /// Generates an iterator of nodes with the specified range of values. Iterates in descending order.
165    pub fn range_desc<S: AvltrieeSearch<T, I, A>>(
166        s: &'a S,
167        start: &I,
168        end: &I,
169    ) -> AvltrieeIter<'a, T, I, A> {
170        Self::range_inner(s, start, end, Order::Desc)
171    }
172
173    fn range_inner<S: AvltrieeSearch<T, I, A>>(
174        s: &'a S,
175        start: &I,
176        end: &I,
177        order: Order,
178    ) -> AvltrieeIter<'a, T, I, A> {
179        let triee = s.as_ref();
180        if let Some(range) = s.range(start, end) {
181            AvltrieeIter::new(triee, Some(range.start), Some(range.end), order)
182        } else {
183            AvltrieeIter::new(triee, None, None, order)
184        }
185    }
186}
187
188impl<'a, T, I: ?Sized, A: AvltrieeAllocator<T>> Iterator for AvltrieeIter<'a, T, I, A> {
189    type Item = NonZeroU32;
190
191    fn next(&mut self) -> Option<Self::Item> {
192        self.now.map(|c| {
193            self.now = if Some(c) == self.end_row {
194                let same = unsafe { self.triee.node_unchecked(c) }.same;
195                if same.is_some() {
196                    self.end_row = same;
197                }
198                same
199            } else {
200                let next_func = self.next_func;
201                next_func(self.triee, c, self.same_branch).map(|(i, b)| {
202                    self.same_branch = b;
203                    i
204                })
205            };
206            c
207        })
208    }
209}
210
211impl<T, I: ?Sized, A: AvltrieeAllocator<T>> Avltriee<T, I, A> {
212    /// Generate an iterator.
213    pub fn iter(&self) -> AvltrieeIter<T, I, A> {
214        AvltrieeIter::new(
215            &self,
216            self.min(self.root()),
217            self.max(self.root()),
218            Order::Asc,
219        )
220    }
221
222    /// Generate an iterator. Iterates in descending order.
223    pub fn desc_iter(&self) -> AvltrieeIter<T, I, A> {
224        AvltrieeIter::new(
225            &self,
226            self.min(self.root()),
227            self.max(self.root()),
228            Order::Desc,
229        )
230    }
231
232    /// Generates an iterator of nodes with the same value as the specified value.
233    pub fn iter_by<'a>(&'a self, value: &I) -> AvltrieeIter<T, I, A>
234    where
235        Self: AvltrieeSearch<T, I, A>,
236    {
237        AvltrieeIter::by(self, value)
238    }
239
240    /// Generates an iterator with values starting from the specified value.
241    pub fn iter_from<'a>(&'a self, value: &I) -> AvltrieeIter<T, I, A>
242    where
243        Self: AvltrieeSearch<T, I, A>,
244    {
245        AvltrieeIter::from_asc(self, value)
246    }
247
248    /// Generates an iterator with values starting from the specified value. Iterates in descending order.
249    pub fn desc_iter_from<'a>(&'a self, value: &I) -> AvltrieeIter<T, I, A>
250    where
251        Self: AvltrieeSearch<T, I, A>,
252    {
253        AvltrieeIter::from_desc(self, value)
254    }
255
256    /// Generates an iterator of nodes with values less than or equal to the specified value.
257    pub fn iter_to<'a>(&'a self, value: &I) -> AvltrieeIter<T, I, A>
258    where
259        Self: AvltrieeSearch<T, I, A>,
260    {
261        AvltrieeIter::to_asc(self, value)
262    }
263
264    /// Generates an iterator of nodes with values less than or equal to the specified value. Iterates in descending order.
265    pub fn desc_iter_to<'a>(&'a self, value: &I) -> AvltrieeIter<T, I, A>
266    where
267        Self: AvltrieeSearch<T, I, A>,
268    {
269        AvltrieeIter::to_desc(self, value)
270    }
271
272    /// Generates an iterator of nodes with values greater than the specified value.
273    pub fn iter_over<'a>(&'a self, value: &I) -> AvltrieeIter<T, I, A>
274    where
275        Self: AvltrieeSearch<T, I, A>,
276    {
277        AvltrieeIter::over_asc(self, value)
278    }
279
280    /// Generates an iterator of nodes with values greater than the specified value. Iterates in descending order.
281    pub fn desc_iter_over<'a>(&'a self, value: &I) -> AvltrieeIter<T, I, A>
282    where
283        Self: AvltrieeSearch<T, I, A>,
284    {
285        AvltrieeIter::over_desc(self, value)
286    }
287
288    /// Generates an iterator of nodes with values less than the specified value.
289    pub fn iter_under<'a>(&'a self, value: &I) -> AvltrieeIter<T, I, A>
290    where
291        Self: AvltrieeSearch<T, I, A>,
292    {
293        AvltrieeIter::under_asc(self, value)
294    }
295
296    /// Generates an iterator of nodes with values less than the specified value. Iterates in descending order.
297    pub fn desc_iter_under<'a>(&'a self, value: &I) -> AvltrieeIter<T, I, A>
298    where
299        Self: AvltrieeSearch<T, I, A>,
300    {
301        AvltrieeIter::under_desc(self, value)
302    }
303
304    /// Generates an iterator of nodes with the specified range of values.
305    pub fn iter_range<'a>(&'a self, start: &I, end: &I) -> AvltrieeIter<T, I, A>
306    where
307        Self: AvltrieeSearch<T, I, A>,
308    {
309        AvltrieeIter::range_asc(self, start, end)
310    }
311
312    /// Generates an iterator of nodes with the specified range of values. Iterates in descending order.
313    pub fn desc_iter_range<'a>(&'a self, start: &I, end: &I) -> AvltrieeIter<T, I, A>
314    where
315        Self: AvltrieeSearch<T, I, A>,
316    {
317        AvltrieeIter::range_desc(self, start, end)
318    }
319
320    fn next(
321        &self,
322        c: NonZeroU32,
323        same_branch: Option<NonZeroU32>,
324    ) -> Option<(NonZeroU32, Option<NonZeroU32>)> {
325        let mut node = unsafe { self.node_unchecked(c) };
326        if let Some(same) = node.same {
327            Some((
328                same,
329                if same_branch.is_some() {
330                    same_branch
331                } else {
332                    Some(c)
333                },
334            ))
335        } else {
336            let current = if let Some(same_branch) = same_branch {
337                node = unsafe { self.node_unchecked(same_branch) };
338                same_branch
339            } else {
340                c
341            };
342            if node.right.is_some() {
343                Some((self.min(node.right).unwrap(), None))
344            } else {
345                node.parent.and_then(|parent| {
346                    if unsafe { self.node_unchecked(parent) }.left == Some(current) {
347                        Some((parent, None))
348                    } else {
349                        self.retroactive(parent).map(|i| (i, None))
350                    }
351                })
352            }
353        }
354    }
355
356    fn retroactive(&self, c: NonZeroU32) -> Option<NonZeroU32> {
357        unsafe { self.node_unchecked(c) }.parent.and_then(|parent| {
358            if unsafe { self.node_unchecked(parent) }.right == Some(c) {
359                self.retroactive(parent).filter(|p| p.get() != c.get())
360            } else {
361                Some(parent)
362            }
363        })
364    }
365
366    fn next_desc(
367        &self,
368        c: NonZeroU32,
369        same_branch: Option<NonZeroU32>,
370    ) -> Option<(NonZeroU32, Option<NonZeroU32>)> {
371        let mut node = unsafe { self.node_unchecked(c) };
372        if let Some(same) = node.same {
373            Some((
374                same,
375                Some(if let Some(same_branch) = same_branch {
376                    same_branch
377                } else {
378                    c
379                }),
380            ))
381        } else {
382            let mut current = c;
383            if let Some(same_branch) = same_branch {
384                current = same_branch;
385                node = unsafe { self.node_unchecked(current) };
386            }
387            if node.left.is_some() {
388                Some((self.max(node.left).unwrap(), None))
389            } else {
390                node.parent.and_then(|parent| {
391                    if unsafe { self.node_unchecked(parent) }.right == Some(current) {
392                        Some((parent, None))
393                    } else {
394                        self.retroactive_desc(parent).map(|i| (i, None))
395                    }
396                })
397            }
398        }
399    }
400
401    fn retroactive_desc(&self, c: NonZeroU32) -> Option<NonZeroU32> {
402        unsafe { self.node_unchecked(c) }.parent.and_then(|parent| {
403            if unsafe { self.node_unchecked(parent) }.left == Some(c) {
404                self.retroactive_desc(parent).filter(|p| *p != c)
405            } else {
406                Some(parent)
407            }
408        })
409    }
410}
avltriee/iter.rs

avltriee/
iter.rs
