closed_interval_set/
union_iterator.rs

1//! Computes the union of two normalized iterators.
2use crate::ClosedRange;
3use crate::Endpoint;
4use crate::IntoNormalizedRangeIter;
5use crate::NormalizedRangeIter;
6use crate::RangeVec;
7
8use core::iter::Peekable;
9
10pub struct UnionIterator<T, X, Y>
11where
12    T: Endpoint,
13    X: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
14    Y: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
15{
16    accumulator: Option<(T, T)>,
17    left: Peekable<X>,
18    right: Peekable<Y>,
19}
20
21impl<T, X, Y> UnionIterator<T, X, Y>
22where
23    T: Endpoint,
24    X: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
25    Y: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
26{
27    pub fn new(left: X, right: Y) -> Self {
28        Self {
29            accumulator: None,
30            left: left.peekable(),
31            right: right.peekable(),
32        }
33    }
34}
35
36impl<T, X, Y> Iterator for UnionIterator<T, X, Y>
37where
38    T: Endpoint,
39    X: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
40    Y: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
41{
42    type Item = (T, T);
43
44    fn next(&mut self) -> Option<(T, T)> {
45        use core::cmp::Ordering; // Safe because both iterators are normalized
46
47        loop {
48            let next;
49
50            match (
51                self.left.peek().map(|x| x.get()),
52                self.right.peek().map(|x| x.get()),
53            ) {
54                (Some(left), Some(right)) if T::cmp_range(left, right) <= Ordering::Equal => {
55                    next = left;
56                    self.left.next();
57                }
58                (Some(_), Some(right)) => {
59                    next = right;
60                    self.right.next();
61                }
62                (Some(x), None) => {
63                    next = x;
64                    self.left.next();
65                }
66                (None, Some(x)) => {
67                    next = x;
68                    self.right.next();
69                }
70                (None, None) => return self.accumulator.take(),
71            };
72
73            let (next_start, next_stop) = next;
74            let Some((acc_start, acc_stop)) = self.accumulator.take() else {
75                self.accumulator = Some(next);
76                continue;
77            };
78
79            // Try to join `acc <= next`.
80            if let Some(min_start) = acc_stop.increase_toward(next_start) {
81                // They're disjoint, produce `acc`, buffer `next.
82                if next_start.cmp_end(min_start) > Ordering::Equal {
83                    self.accumulator = Some(next);
84                    return Some((acc_start, acc_stop));
85                }
86            }
87
88            self.accumulator = Some((acc_start, acc_stop.top_end(next_stop)));
89        }
90    }
91
92    #[inline]
93    fn size_hint(&self) -> (usize, Option<usize>) {
94        let (left_min, left_max) = self.left.size_hint();
95        let (right_min, right_max) = self.right.size_hint();
96
97        let max = left_max
98            .unwrap_or(usize::MAX)
99            .checked_add(right_max.unwrap_or(usize::MAX));
100        // If both are empty, the union is empty.  Otherwise,
101        // they could end up as a single big range.
102        let min = ((left_min != 0) | (right_min != 0)) as usize;
103        (min, max)
104    }
105}
106
107impl<T, X, Y> crate::private::Sealed for UnionIterator<T, X, Y>
108where
109    T: Endpoint,
110    X: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
111    Y: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
112{
113}
114
115impl<T, X, Y> NormalizedRangeIter for UnionIterator<T, X, Y>
116where
117    T: Endpoint,
118    X: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
119    Y: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
120{
121}
122
123// We use `Peekable` internally, so we'll reliably return None after None.
124impl<T, X, Y> core::iter::FusedIterator for UnionIterator<T, X, Y>
125where
126    T: Endpoint,
127    X: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
128    Y: Sized + NormalizedRangeIter + Iterator<Item: ClosedRange<EndT = T>>,
129{
130}
131
132impl<T: Endpoint> RangeVec<T> {
133    /// Constructs the union of this [`RangeVec`] and any iterator of
134    /// ranges.
135    ///
136    /// This operation takes time linear in the total input size.  While
137    /// asymptotically better than [`RangeVec::into_union`], the latter's
138    /// constant factors may be better than this method's.
139    ///
140    /// See [`union_vec`] for more general types.
141    ///
142    /// [`union_vec`]: crate::union_vec
143    #[inline(always)]
144    #[must_use]
145    pub fn union(&self, other: impl IntoNormalizedRangeIter<Item: ClosedRange<EndT = T>>) -> Self {
146        #[inline(never)]
147        fn doit<T: Endpoint>(
148            this: &RangeVec<T>,
149            other: impl NormalizedRangeIter<Item: ClosedRange<EndT = T>>,
150        ) -> RangeVec<T> {
151            this.iter().union(other).collect_range_vec()
152        }
153
154        doit(self, other.into_iter())
155    }
156}
157
158#[cfg(test)]
159#[cfg_attr(coverage_nightly, coverage(off))]
160mod test {
161    use super::*;
162    use alloc::vec;
163    use alloc::vec::Vec;
164
165    #[test]
166    fn test_union_iterator_smoke() {
167        use crate::NormalizedRangeIter;
168
169        let acc = vec![(1u8, 4u8), (5u8, 1u8), (5u8, 10u8)];
170        let src = [(0u8, 2u8), (11u8, 15u8)];
171
172        assert_eq!(
173            crate::normalize_vec(acc.clone())
174                .iter()
175                .union(crate::normalize_vec(src.to_vec()).into_iter())
176                .collect_range_vec()
177                .into_vec(),
178            vec![(0u8, 15u8)]
179        );
180
181        assert_eq!(
182            crate::normalize_vec(src.to_vec())
183                .iter()
184                .union(crate::normalize_vec(vec![]))
185                .collect_range_vec()
186                .into_vec(),
187            src.to_vec()
188        );
189
190        {
191            let empty = Vec::<(u8, u8)>::new();
192            let mut union = crate::normalize_vec(empty.clone())
193                .into_iter()
194                .union(crate::normalize_vec(empty).into_iter());
195
196            assert_eq!(union.next(), None);
197            // And still none
198            assert_eq!(union.next(), None);
199        }
200    }
201
202    proptest::proptest! {
203        #[test]
204        fn test_union_iterator(xs: Vec<(u8, u8)>, ys: Vec<(u8, u8)>)
205        {
206            use crate::RangeVec;
207            use crate::ranges_to_bits;
208
209            let bits = {
210                let mut all_ranges = xs.clone();
211                all_ranges.extend(&ys);
212                ranges_to_bits(&all_ranges)
213            };
214
215            let xs = RangeVec::from_vec(xs);
216            let ys = RangeVec::from_vec(ys);
217
218            {
219                let union = xs.iter().union(ys.iter());
220
221                let (min_size, max_size) = union.size_hint();
222                if !xs.is_empty() || !ys.is_empty() {
223                    assert!(min_size > 0);
224                } else {
225                    assert!(min_size == 0);
226                }
227
228                let max_size = max_size.expect("should have limits");
229                assert!(max_size == xs.len() + ys.len());
230
231                let union = union.collect_range_vec();
232                assert_eq!(bits, ranges_to_bits(&union));
233            }
234
235            {
236                let union = ys.iter().union(xs).collect_range_vec();
237                assert_eq!(bits, ranges_to_bits(&union));
238            }
239        }
240
241        #[test]
242        fn test_union_iterator_identities(xs: Vec<(u8, u8)>)
243        {
244            let xs = RangeVec::from_vec(xs);
245
246            // xs | xs = xs
247            {
248                let xxs = xs.iter().union(xs.iter()).collect_range_vec();
249                assert_eq!(&xxs, &xs);
250            }
251
252            // xs | -xs = universe
253            let universe = xs.iter().union(xs.iter().complement()).collect_range_vec();
254            assert_eq!(universe.inner(), &[(0, 255u8)]);
255        }
256    }
257}