iset/
iter.rs

1//! Module with various iterators over `IntervalMap` and `IntervalSet`.
2
3use alloc::vec::Vec;
4use core::ops::{Range, RangeBounds, Bound};
5use core::iter::FusedIterator;
6use core::mem;
7
8use super::{Interval, IntervalMap, Node, IndexType, check_ordered, BitVec};
9
10fn should_go_left<T, V, Ix>(nodes: &[Node<T, V, Ix>], index: Ix, start_bound: Bound<&T>) -> bool
11where T: PartialOrd + Copy,
12      Ix: IndexType,
13{
14    if !nodes[index.get()].left.defined() {
15        return false;
16    }
17    let left_end = nodes[nodes[index.get()].left.get()].subtree_interval.end;
18    match start_bound {
19        Bound::Included(&value) | Bound::Excluded(&value) => left_end >= value,
20        Bound::Unbounded => true,
21    }
22}
23
24fn should_go_right<T, V, Ix>(nodes: &[Node<T, V, Ix>], index: Ix, end_bound: Bound<&T>) -> bool
25where T: PartialOrd + Copy,
26      Ix: IndexType,
27{
28    if !nodes[index.get()].right.defined() {
29        return false;
30    }
31    let right_start = nodes[nodes[index.get()].right.get()].subtree_interval.start;
32    match end_bound {
33        Bound::Included(&value) => right_start <= value,
34        Bound::Excluded(&value) => right_start < value,
35        Bound::Unbounded => true,
36    }
37}
38
39#[derive(Debug, Clone)]
40struct ActionStack(BitVec);
41
42impl ActionStack {
43    fn new() -> Self {
44        Self(BitVec::from_elem(2, false))
45    }
46
47    #[inline]
48    fn push(& mut self) {
49        self.0.push(false);
50        self.0.push(false);
51    }
52
53    // 00 - just entered
54    // 01 - was to the left
55    // 10 - returned
56    // 11 - was to the right
57
58    #[inline]
59    fn can_go_left(&self) -> bool {
60        !self.0.get_tail(1) && !self.0.get_tail(0)
61    }
62
63    #[inline]
64    fn go_left(&mut self) {
65        self.0.set1(self.0.len() - 1);
66    }
67
68    #[inline]
69    fn can_match(&self) -> bool {
70        !self.0.get_tail(1)
71    }
72
73    #[inline]
74    fn make_match(&mut self) {
75        self.0.set1(self.0.len() - 2);
76        self.0.set0(self.0.len() - 1);
77    }
78
79    #[inline]
80    fn can_go_right(&self) -> bool {
81        !self.0.get_tail(0)
82    }
83
84    #[inline]
85    fn go_right(&mut self) {
86        self.0.set1(self.0.len() - 2);
87        self.0.set1(self.0.len() - 1);
88    }
89
90    #[inline]
91    fn pop(&mut self) {
92        self.0.pop();
93        self.0.pop();
94    }
95
96    #[inline]
97    fn is_empty(&self) -> bool {
98        self.0.is_empty()
99    }
100}
101
102fn move_to_next<T, V, R, Ix>(nodes: &[Node<T, V, Ix>], mut index: Ix, range: &R, stack: &mut ActionStack) -> Ix
103where T: PartialOrd + Copy,
104      R: RangeBounds<T>,
105      Ix: IndexType,
106{
107    while index.defined() {
108        if stack.can_go_left() {
109            while should_go_left(nodes, index, range.start_bound()) {
110                stack.go_left();
111                stack.push();
112                index = nodes[index.get()].left;
113            }
114            stack.go_left();
115        }
116
117        if stack.can_match() {
118            stack.make_match();
119            if nodes[index.get()].interval.intersects_range(range) {
120                return index;
121            }
122        }
123
124        if stack.can_go_right() && should_go_right(nodes, index, range.end_bound()) {
125            stack.go_right();
126            stack.push();
127            index = nodes[index.get()].right;
128        } else {
129            stack.pop();
130            index = nodes[index.get()].parent;
131        }
132    }
133    index
134}
135
136/// Macro that generates Iterator over IntervalMap.
137macro_rules! iterator {
138    (
139        $(#[$outer:meta])*
140        struct $name:ident -> $elem:ty,
141        $node:ident -> $out:expr, {$( $mut_:tt )?}
142    ) => {
143        $(#[$outer])*
144        pub struct $name<'a, T, V, R, Ix>
145        where T: PartialOrd + Copy,
146              R: RangeBounds<T>,
147              Ix: IndexType,
148        {
149            pub(crate) index: Ix,
150            range: R,
151            nodes: &'a $( $mut_ )? [Node<T, V, Ix>],
152            stack: ActionStack,
153        }
154
155        impl<'a, T: PartialOrd + Copy, V, R: RangeBounds<T>, Ix: IndexType> $name<'a, T, V, R, Ix> {
156            pub(crate) fn new(tree: &'a $( $mut_ )? IntervalMap<T, V, Ix>, range: R) -> Self {
157                check_ordered(&range);
158                Self {
159                    index: tree.root,
160                    range,
161                    nodes: & $( $mut_ )? tree.nodes,
162                    stack: ActionStack::new(),
163                }
164            }
165        }
166
167        impl<'a, T: PartialOrd + Copy, V, R: RangeBounds<T>, Ix: IndexType> Iterator for $name<'a, T, V, R, Ix> {
168            type Item = $elem;
169
170            fn next(&mut self) -> Option<Self::Item> {
171                self.index = move_to_next(self.nodes, self.index, &self.range, &mut self.stack);
172                if self.index.defined() {
173                    let $node = & $( $mut_ )? self.nodes[self.index.get()];
174                    Some($out)
175                } else {
176                    None
177                }
178            }
179
180            fn size_hint(& self) -> (usize, Option<usize>) {
181                // Not optimal implementation, basically, always returns lower bound = 0, upper bound = map.len().
182                (0, Some(self.nodes.len()))
183            }
184        }
185
186        impl<'a, T: PartialOrd + Copy, V, R: RangeBounds<T>, Ix: IndexType> FusedIterator for $name<'a, T, V, R, Ix> { }
187    };
188}
189
190iterator! {
191    #[doc="Iterator over pairs `(x..y, &value)`."]
192    #[derive(Clone, Debug)]
193    struct Iter -> (Range<T>, &'a V),
194    node -> (node.interval.to_range(), &node.value), { /* no mut */ }
195}
196
197iterator! {
198    #[doc="Iterator over intervals `x..y`."]
199    #[derive(Clone, Debug)]
200    struct Intervals -> Range<T>,
201    node -> node.interval.to_range(), { /* no mut */ }
202}
203
204iterator! {
205    #[doc="Iterator over values."]
206    #[derive(Clone, Debug)]
207    struct Values -> &'a V,
208    node -> &node.value, { /* no mut */ }
209}
210
211iterator! {
212    #[doc="Iterator over pairs `(x..y, &mut value)`."]
213    #[derive(Debug)]
214    struct IterMut -> (Range<T>, &'a mut V),
215    node -> (node.interval.to_range(), unsafe { &mut *(&mut node.value as *mut V) }), { mut }
216}
217
218iterator! {
219    #[doc="Iterator over mutable values."]
220    #[derive(Debug)]
221    struct ValuesMut -> &'a mut V,
222    node -> unsafe { &mut *(&mut node.value as *mut V) }, { mut }
223}
224
225/// Macro that generates IntoIterator over IntervalMap.
226macro_rules! into_iterator {
227    (
228        $(#[$outer:meta])*
229        struct $name:ident -> $elem:ty,
230        $node:ident -> $out:expr
231    ) => {
232        $(#[$outer])*
233        pub struct $name<T, V, R, Ix>
234        where T: PartialOrd + Copy,
235              R: RangeBounds<T>,
236              Ix: IndexType,
237        {
238            index: Ix,
239            range: R,
240            nodes: Vec<Node<T, V, Ix>>,
241            stack: ActionStack,
242        }
243
244        impl<T: PartialOrd + Copy, V, R: RangeBounds<T>, Ix: IndexType> $name<T, V, R, Ix> {
245            pub(crate) fn new(tree: IntervalMap<T, V, Ix>, range: R) -> Self {
246                check_ordered(&range);
247                Self {
248                    index: tree.root,
249                    range,
250                    nodes: tree.nodes,
251                    stack: ActionStack::new(),
252                }
253            }
254        }
255
256        impl<T: PartialOrd + Copy, V, R: RangeBounds<T>, Ix: IndexType> Iterator for $name<T, V, R, Ix> {
257            type Item = $elem;
258
259            fn next(&mut self) -> Option<Self::Item> {
260                self.index = move_to_next(&self.nodes, self.index, &self.range, &mut self.stack);
261                if self.index.defined() {
262                    let $node = &mut self.nodes[self.index.get()];
263                    Some($out)
264                } else {
265                    None
266                }
267            }
268
269            fn size_hint(& self) -> (usize, Option<usize>) {
270                // Not optimal implementation, basically, always returns lower bound = 0, upper bound = map.len().
271                (0, Some(self.nodes.len()))
272            }
273        }
274
275        impl<T: PartialOrd + Copy, V, R: RangeBounds<T>, Ix: IndexType> FusedIterator for $name<T, V, R, Ix> { }
276    };
277}
278
279into_iterator! {
280    #[doc="Iterator over pairs `(x..y, value)`. Takes ownership of the interval map/set."]
281    #[derive(Debug)]
282    struct IntoIter -> (Range<T>, V),
283    node -> (node.interval.to_range(), mem::replace(&mut node.value, unsafe { mem::zeroed() }))
284}
285
286into_iterator! {
287    #[doc="Iterator over intervals `x..y`. Takes ownership of the interval map/set."]
288    #[derive(Debug)]
289    struct IntoIntervals -> Range<T>,
290    node -> node.interval.to_range()
291}
292
293into_iterator! {
294    #[doc="Iterator over values. Takes ownership of the interval map/set."]
295    #[derive(Debug)]
296    struct IntoValues -> V,
297    node -> mem::replace(&mut node.value, unsafe { mem::zeroed() })
298}
299
300/// Macro that generates unsorted iterator over IntervalMap.
301macro_rules! unsorted_iterator {
302    (
303        $(#[$outer:meta])*
304        struct $name:ident -> $elem:ty,
305        ($get_iter:ident -> $iter_type:ty),
306        $node:ident -> $out:expr, {$( $mut_:tt )?}
307    ) => {
308        $(#[$outer])*
309        pub struct $name<'a, T, V, Ix>($iter_type)
310        where T: PartialOrd + Copy,
311              Ix: IndexType;
312
313        impl<'a, T: PartialOrd + Copy, V, Ix: IndexType> $name<'a, T, V, Ix> {
314            pub(crate) fn new(tree: &'a $( $mut_ )? IntervalMap<T, V, Ix>) -> Self {
315                Self(tree.nodes.$get_iter())
316            }
317        }
318
319        impl<'a, T: PartialOrd + Copy, V, Ix: IndexType> Iterator for $name<'a, T, V, Ix> {
320            type Item = $elem;
321
322            fn next(&mut self) -> Option<Self::Item> {
323                match self.0.next() {
324                    Some($node) => Some($out),
325                    None => None,
326                }
327            }
328
329            #[inline]
330            fn size_hint(&self) -> (usize, Option<usize>) {
331                self.0.size_hint()
332            }
333        }
334
335        impl<'a, T: PartialOrd + Copy, V, Ix: IndexType> DoubleEndedIterator for $name<'a, T, V, Ix> {
336            fn next_back(&mut self) -> Option<Self::Item> {
337                match self.0.next_back() {
338                    Some($node) => Some($out),
339                    None => None,
340                }
341            }
342        }
343
344        impl<'a, T: PartialOrd + Copy, V, Ix: IndexType> FusedIterator for $name<'a, T, V, Ix> { }
345
346        impl<'a, T: PartialOrd + Copy, V, Ix: IndexType> ExactSizeIterator for $name<'a, T, V, Ix> {
347            #[inline]
348            fn len(&self) -> usize {
349                self.0.len()
350            }
351        }
352    };
353}
354
355unsorted_iterator! {
356    #[doc="Unsorted iterator over pairs `(x..y, &value)`."]
357    #[derive(Clone, Debug)]
358    struct UnsIter -> (Range<T>, &'a V),
359    (iter -> alloc::slice::Iter<'a, Node<T, V, Ix>>),
360    node -> (node.interval.to_range(), &node.value), { /* no mut */ }
361}
362
363unsorted_iterator! {
364    #[doc="Unsorted iterator over intervals `x..y`."]
365    #[derive(Clone, Debug)]
366    struct UnsIntervals -> Range<T>,
367    (iter -> alloc::slice::Iter<'a, Node<T, V, Ix>>),
368    node -> node.interval.to_range(), { /* no mut */ }
369}
370
371unsorted_iterator! {
372    #[doc="Unsorted iterator over values `&V`."]
373    #[derive(Clone, Debug)]
374    struct UnsValues -> &'a V,
375    (iter -> alloc::slice::Iter<'a, Node<T, V, Ix>>),
376    node -> &node.value, { /* no mut */ }
377}
378
379unsorted_iterator! {
380    #[doc="Unsorted iterator over pairs `(x..y, &mut V)`."]
381    #[derive(Debug)]
382    struct UnsIterMut -> (Range<T>, &'a mut V),
383    (iter_mut -> alloc::slice::IterMut<'a, Node<T, V, Ix>>),
384    node -> (node.interval.to_range(), &mut node.value), { mut }
385}
386
387unsorted_iterator! {
388    #[doc="Unsorted iterator over mutable values `&mut V`."]
389    #[derive(Debug)]
390    struct UnsValuesMut -> &'a mut V,
391    (iter_mut -> alloc::slice::IterMut<'a, Node<T, V, Ix>>),
392    node -> &mut node.value, { mut }
393}
394
395/// Macro that generates unsorted IntoIterator over IntervalMap.
396macro_rules! unsorted_into_iterator {
397    (
398        $(#[$outer:meta])*
399        struct $name:ident -> $elem:ty,
400        $node:ident -> $out:expr
401    ) => {
402        $(#[$outer])*
403        pub struct $name<T, V, Ix>(alloc::vec::IntoIter<Node<T, V, Ix>>)
404        where T: PartialOrd + Copy,
405              Ix: IndexType;
406
407        impl<T: PartialOrd + Copy, V, Ix: IndexType> $name<T, V, Ix> {
408            pub(crate) fn new(tree: IntervalMap<T, V, Ix>) -> Self {
409                Self(tree.nodes.into_iter())
410            }
411        }
412
413        impl<T: PartialOrd + Copy, V, Ix: IndexType> Iterator for $name<T, V, Ix> {
414            type Item = $elem;
415
416            fn next(&mut self) -> Option<Self::Item> {
417                match self.0.next() {
418                    Some($node) => Some($out),
419                    None => None,
420                }
421            }
422
423            #[inline]
424            fn size_hint(&self) -> (usize, Option<usize>) {
425                self.0.size_hint()
426            }
427        }
428
429        impl<T: PartialOrd + Copy, V, Ix: IndexType> FusedIterator for $name<T, V, Ix> { }
430
431        impl<T: PartialOrd + Copy, V, Ix: IndexType> ExactSizeIterator for $name<T, V, Ix> {
432            #[inline]
433            fn len(&self) -> usize {
434                self.0.len()
435            }
436        }
437    };
438}
439
440unsorted_into_iterator! {
441    #[doc="Unsorted IntoIterator over pairs `(x..y, V)`. Takes ownership of the interval map/set."]
442    #[derive(Debug)]
443    struct UnsIntoIter -> (Range<T>, V),
444    node -> (node.interval.to_range(), node.value)
445}
446
447unsorted_into_iterator! {
448    #[doc="Unsorted IntoIterator over intervals `x..y`. Takes ownership of the interval map/set."]
449    #[derive(Debug)]
450    struct UnsIntoIntervals -> Range<T>,
451    node -> node.interval.to_range()
452}
453
454unsorted_into_iterator! {
455    #[doc="Unsorted IntoIterator over intervals `x..y`. Takes ownership of the interval map/set."]
456    #[derive(Debug)]
457    struct UnsIntoValues -> V,
458    node -> node.value
459}
460
461fn should_go_left_exact<T, V, Ix>(nodes: &[Node<T, V, Ix>], index: Ix, query: &Interval<T>) -> bool
462where T: PartialOrd + Copy,
463      Ix: IndexType,
464{
465    let node = &nodes[index.get()];
466    let left_index = nodes[index.get()].left;
467    left_index.defined() && query <= &node.interval && nodes[left_index.get()].subtree_interval.contains(query)
468}
469
470fn should_go_right_exact<T, V, Ix>(nodes: &[Node<T, V, Ix>], index: Ix, query: &Interval<T>) -> bool
471where T: PartialOrd + Copy,
472      Ix: IndexType,
473{
474    let node = &nodes[index.get()];
475    let right_index = nodes[index.get()].right;
476    right_index.defined() && query >= &node.interval && nodes[right_index.get()].subtree_interval.contains(query)
477}
478
479fn move_to_next_exact<T, V, Ix>(
480    nodes: &[Node<T, V, Ix>],
481    mut index: Ix,
482    query: &Interval<T>,
483    stack: &mut ActionStack
484) -> Ix
485where T: PartialOrd + Copy,
486      Ix: IndexType,
487{
488    while !stack.is_empty() && index.defined() {
489        if stack.can_go_left() {
490            while should_go_left_exact(nodes, index, query) {
491                stack.go_left();
492                stack.push();
493                index = nodes[index.get()].left;
494            }
495            stack.go_left();
496        }
497
498        if stack.can_match() {
499            stack.make_match();
500            if query == &nodes[index.get()].interval {
501                return index;
502            }
503        }
504
505        if stack.can_go_right() && should_go_right_exact(nodes, index, query) {
506            stack.go_right();
507            stack.push();
508            index = nodes[index.get()].right;
509        } else {
510            stack.pop();
511            index = nodes[index.get()].parent;
512        }
513    }
514    Ix::MAX
515}
516
517/// Macro that generates iterator over exactly matching intervals in the IntervalMap.
518macro_rules! iterator_exact {
519    (
520        $(#[$outer:meta])*
521        struct $name:ident -> $elem:ty,
522        $node:ident -> $out:expr, {$( $mut_:tt )?}
523    ) => {
524        $(#[$outer])*
525        pub struct $name<'a, T, V, Ix>
526        where T: PartialOrd + Copy,
527              Ix: IndexType,
528        {
529            pub(crate) index: Ix,
530            query: Interval<T>,
531            nodes: &'a $( $mut_ )? [Node<T, V, Ix>],
532            stack: ActionStack,
533        }
534
535        impl<'a, T: PartialOrd + Copy, V, Ix: IndexType> $name<'a, T, V, Ix> {
536            pub(crate) fn new(tree: &'a $( $mut_ )? IntervalMap<T, V, Ix>, query: Interval<T>) -> Self {
537                Self {
538                    index: tree.find_index(&query),
539                    query,
540                    nodes: & $( $mut_ )? tree.nodes,
541                    stack: ActionStack::new(),
542                }
543            }
544        }
545
546        impl<'a, T: PartialOrd + Copy, V, Ix: IndexType> Iterator for $name<'a, T, V, Ix> {
547            type Item = $elem;
548
549            fn next(&mut self) -> Option<Self::Item> {
550                self.index = move_to_next_exact(self.nodes, self.index, &self.query, &mut self.stack);
551                if self.index.defined() {
552                    let $node = & $( $mut_ )? self.nodes[self.index.get()];
553                    Some($out)
554                } else {
555                    None
556                }
557            }
558
559            fn size_hint(& self) -> (usize, Option<usize>) {
560                // Not optimal implementation, basically, always returns lower bound = 0, upper bound = map.len().
561                (0, Some(self.nodes.len()))
562            }
563        }
564
565        impl<'a, T: PartialOrd + Copy, V, Ix: IndexType> FusedIterator for $name<'a, T, V, Ix> { }
566    };
567}
568
569iterator_exact! {
570    #[doc="Iterator over values `&V` for exact matches with the query."]
571    #[derive(Clone, Debug)]
572    struct ValuesExact -> &'a V,
573    node -> &node.value, { /* no mut */ }
574}
575
576iterator_exact! {
577    #[doc="Iterator over mutable values `&mut V` for exact matches with the query."]
578    #[derive(Debug)]
579    struct ValuesExactMut -> &'a mut V,
580    node -> unsafe { &mut *(&mut node.value as *mut V) }, { mut }
581}
iset/iter.rs

iset/
iter.rs
