range_set_blaze/
union_iter.rs

1use core::{
2    cmp::max,
3    iter::FusedIterator,
4    ops::{self, RangeInclusive},
5};
6
7use alloc::vec;
8use itertools::Itertools;
9
10use crate::{
11    unsorted_disjoint::{AssumeSortedStarts, UnsortedDisjoint},
12    BitAndMerge, BitOrMerge, BitSubMerge, BitXOrTee, Integer, NotIter, SortedDisjoint,
13    SortedStarts,
14};
15
16/// Turns any number of [`SortedStarts`] iterators into a [`SortedDisjoint`] iterator of their union,
17/// i.e., all the integers in any input iterator, as sorted & disjoint ranges. Uses [`Merge`]
18/// or [`KMerge`].
19///
20/// [`SortedDisjoint`]: crate::SortedDisjoint
21/// [`Merge`]: crate::Merge
22/// [`KMerge`]: crate::KMerge
23///
24/// # Examples
25///
26/// ## For [`SortedDisjoint`]
27///
28/// ```
29/// use itertools::Itertools;
30/// use range_set_blaze::{UnionIter, Merge, SortedDisjoint, CheckSortedDisjoint};
31///
32/// let a = CheckSortedDisjoint::new([1..=2, 5..=100]);
33/// let b = CheckSortedDisjoint::from([2..=6]);
34/// let union = UnionIter::new(Merge::new(a, b));
35/// assert_eq!(union.to_string(), "1..=100");
36///
37/// // Or, equivalently:
38/// let a = CheckSortedDisjoint::new([1..=2, 5..=100]);
39/// let b = CheckSortedDisjoint::from([2..=6]);
40/// let union = a | b;
41/// assert_eq!(union.to_string(), "1..=100")
42/// ```
43///
44/// ## For [`SortedStarts`]
45///
46/// ```
47/// use itertools::Itertools;
48/// use range_set_blaze::{AssumeSortedStarts, SortedDisjoint, UnionIter};
49///
50/// let a = UnionIter::new(AssumeSortedStarts::new([1..=5, 2..=100]));
51/// assert_eq!(a.to_string(), "1..=100");
52/// ```
53#[derive(Clone, Debug)]
54#[must_use = "iterators are lazy and do nothing unless consumed"]
55pub struct UnionIter<T, I>
56where
57    T: Integer,
58    I: SortedStarts<T>,
59{
60    pub(crate) iter: I,
61    pub(crate) option_range: Option<RangeInclusive<T>>,
62}
63
64impl<T, I> UnionIter<T, I>
65where
66    T: Integer,
67    I: SortedStarts<T>,
68{
69    /// Creates a new [`UnionIter`] from zero or more [`SortedDisjoint`] iterators. See [`UnionIter`] for more details and examples.
70    pub fn new(iter: I) -> Self {
71        Self {
72            iter,
73            option_range: None,
74        }
75    }
76}
77
78impl<T: Integer, const N: usize> From<[T; N]> for UnionIter<T, SortedRangeInclusiveVec<T>> {
79    fn from(arr: [T; N]) -> Self {
80        arr.as_slice().into()
81    }
82}
83
84impl<T: Integer> From<&[T]> for UnionIter<T, SortedRangeInclusiveVec<T>> {
85    fn from(slice: &[T]) -> Self {
86        slice.iter().cloned().collect()
87    }
88}
89
90impl<T: Integer, const N: usize> From<[RangeInclusive<T>; N]>
91    for UnionIter<T, SortedRangeInclusiveVec<T>>
92{
93    fn from(arr: [RangeInclusive<T>; N]) -> Self {
94        arr.as_slice().into()
95    }
96}
97
98impl<T: Integer> From<&[RangeInclusive<T>]> for UnionIter<T, SortedRangeInclusiveVec<T>> {
99    fn from(slice: &[RangeInclusive<T>]) -> Self {
100        slice.iter().cloned().collect()
101    }
102}
103
104type SortedRangeInclusiveVec<T> = AssumeSortedStarts<T, vec::IntoIter<RangeInclusive<T>>>;
105
106impl<T: Integer> FromIterator<T> for UnionIter<T, SortedRangeInclusiveVec<T>> {
107    fn from_iter<I>(iter: I) -> Self
108    where
109        I: IntoIterator<Item = T>,
110    {
111        iter.into_iter().map(|x| x..=x).collect()
112    }
113}
114
115impl<T: Integer> FromIterator<RangeInclusive<T>> for UnionIter<T, SortedRangeInclusiveVec<T>> {
116    fn from_iter<I>(iter: I) -> Self
117    where
118        I: IntoIterator<Item = RangeInclusive<T>>,
119    {
120        UnsortedDisjoint::from(iter.into_iter()).into()
121    }
122}
123
124impl<T, I> From<UnsortedDisjoint<T, I>> for UnionIter<T, SortedRangeInclusiveVec<T>>
125where
126    T: Integer,
127    I: Iterator<Item = RangeInclusive<T>>, // Any iterator is OK, because we will sort
128{
129    fn from(unsorted_disjoint: UnsortedDisjoint<T, I>) -> Self {
130        let iter = AssumeSortedStarts {
131            iter: unsorted_disjoint.sorted_by_key(|range| *range.start()),
132        };
133        Self {
134            iter,
135            option_range: None,
136        }
137    }
138}
139
140impl<T: Integer, I> FusedIterator for UnionIter<T, I> where I: SortedStarts<T> + FusedIterator {}
141
142impl<T: Integer, I> Iterator for UnionIter<T, I>
143where
144    I: SortedStarts<T>,
145{
146    type Item = RangeInclusive<T>;
147
148    fn next(&mut self) -> Option<RangeInclusive<T>> {
149        loop {
150            let range = match self.iter.next() {
151                Some(r) => r,
152                None => return self.option_range.take(),
153            };
154
155            let (start, end) = range.into_inner();
156            if end < start {
157                continue;
158            }
159
160            let current_range = match self.option_range.clone() {
161                Some(cr) => cr,
162                None => {
163                    self.option_range = Some(start..=end);
164                    continue;
165                }
166            };
167
168            let (current_start, current_end) = current_range.into_inner();
169            debug_assert!(current_start <= start); // real assert
170            if start <= current_end
171                || (current_end < T::safe_max_value() && start <= current_end + T::one())
172            {
173                self.option_range = Some(current_start..=max(current_end, end));
174                continue;
175            } else {
176                self.option_range = Some(start..=end);
177                return Some(current_start..=current_end);
178            }
179        }
180    }
181
182    // There could be a few as 1 (or 0 if the iter is empty) or as many as the iter.
183    // Plus, possibly one more if we have a range is in progress.
184    fn size_hint(&self) -> (usize, Option<usize>) {
185        let (low, high) = self.iter.size_hint();
186        let low = low.min(1);
187        if self.option_range.is_some() {
188            (low, high.map(|x| x + 1))
189        } else {
190            (low, high)
191        }
192    }
193}
194
195impl<T: Integer, I> ops::Not for UnionIter<T, I>
196where
197    I: SortedStarts<T>,
198{
199    type Output = NotIter<T, Self>;
200
201    fn not(self) -> Self::Output {
202        self.complement()
203    }
204}
205
206impl<T: Integer, R, L> ops::BitOr<R> for UnionIter<T, L>
207where
208    L: SortedStarts<T>,
209    R: SortedDisjoint<T>,
210{
211    type Output = BitOrMerge<T, Self, R>;
212
213    fn bitor(self, rhs: R) -> Self::Output {
214        // It might be fine to optimize to self.iter, but that would require
215        // also considering field 'range'
216        SortedDisjoint::union(self, rhs)
217    }
218}
219
220impl<T: Integer, R, L> ops::Sub<R> for UnionIter<T, L>
221where
222    L: SortedStarts<T>,
223    R: SortedDisjoint<T>,
224{
225    type Output = BitSubMerge<T, Self, R>;
226
227    fn sub(self, rhs: R) -> Self::Output {
228        SortedDisjoint::difference(self, rhs)
229    }
230}
231
232impl<T: Integer, R, L> ops::BitXor<R> for UnionIter<T, L>
233where
234    L: SortedStarts<T>,
235    R: SortedDisjoint<T>,
236{
237    type Output = BitXOrTee<T, Self, R>;
238
239    #[allow(clippy::suspicious_arithmetic_impl)]
240    fn bitxor(self, rhs: R) -> Self::Output {
241        SortedDisjoint::symmetric_difference(self, rhs)
242    }
243}
244
245impl<T: Integer, R, L> ops::BitAnd<R> for UnionIter<T, L>
246where
247    L: SortedStarts<T>,
248    R: SortedDisjoint<T>,
249{
250    type Output = BitAndMerge<T, Self, R>;
251
252    fn bitand(self, other: R) -> Self::Output {
253        SortedDisjoint::intersection(self, other)
254    }
255}
range_set_blaze/union_iter.rs

range_set_blaze/
union_iter.rs
