Skip to main content

range_set_blaze/
sym_diff_iter.rs

1use crate::sym_diff_iter_map::UsizeExtensions;
2use crate::{
3    Integer, Merge, SortedDisjoint, SortedStarts, SymDiffKMerge, SymDiffMerge, merge::KMerge,
4};
5use alloc::collections::BinaryHeap;
6use core::{cmp::Reverse, iter::FusedIterator, ops::RangeInclusive};
7
8/// This `struct` is created by the [`symmetric_difference`] method on [`SortedDisjoint`]. See [`symmetric_difference`]'s
9/// documentation for more.
10///
11/// [`SortedDisjoint`]: crate::SortedDisjoint
12/// [`symmetric_difference`]: crate::SortedDisjointMap::symmetric_difference
13#[derive(Clone, Debug)]
14#[must_use = "iterators are lazy and do nothing unless consumed"]
15pub struct SymDiffIter<T, I> {
16    iter: I,
17    start_or_min_value: T,
18    end_heap: BinaryHeap<Reverse<T>>,
19    next_again: Option<RangeInclusive<T>>,
20    gather: Option<RangeInclusive<T>>,
21}
22
23impl<T, I> FusedIterator for SymDiffIter<T, I>
24where
25    T: Integer,
26    I: SortedStarts<T>,
27{
28}
29
30impl<T, I> Iterator for SymDiffIter<T, I>
31where
32    T: Integer,
33    I: SortedStarts<T>,
34{
35    type Item = RangeInclusive<T>;
36
37    fn next(&mut self) -> Option<RangeInclusive<T>> {
38        loop {
39            let count = self.end_heap.len();
40            let Some(next_range) = self.next_again.take().or_else(|| self.iter.next()) else {
41                // The workspace is empty and next is empty, so return everything gathered.
42                if count == 0 {
43                    return self.gather.take();
44                }
45
46                // The workspace is not empty (but next is empty) is process the next chunk of the workspace.
47                let end = self
48                    .end_heap
49                    .pop()
50                    .expect("Real Assert: the workspace is not empty")
51                    .0;
52                self.remove_same_end(end);
53                let result = self.start_or_min_value..=end;
54                if !self.end_heap.is_empty() {
55                    self.start_or_min_value = end.add_one(); // The 'if' prevents overflow.
56                }
57                if let Some(result) = self.process(count.is_odd(), result) {
58                    return result;
59                }
60                continue;
61            };
62
63            // Next has the same start as the workspace, so add it to the workspace.
64            // (or the workspace is empty, so add it to the workspace.)
65            let (next_start, next_end) = next_range.into_inner();
66            if count == 0 || self.start_or_min_value == next_start {
67                self.start_or_min_value = next_start;
68                self.end_heap.push(Reverse(next_end));
69                continue;
70            }
71
72            // Next start inside the workspace's first chunk, so process up to next_start.
73            let end = self
74                .end_heap
75                .peek()
76                .expect("Real Assert: The workspace has a first chunk.")
77                .0;
78            if next_start <= end {
79                let result = self.start_or_min_value..=next_start.sub_one();
80                self.start_or_min_value = next_start;
81                self.end_heap.push(Reverse(next_end));
82                if let Some(result) = self.process(count.is_odd(), result) {
83                    return result;
84                }
85                continue;
86            }
87
88            // Next start is after the workspaces end, but the workspace contains only one chuck,
89            // so process the workspace and set the workspace to next.
90            self.remove_same_end(end);
91            let result = self.start_or_min_value..=end;
92            if self.end_heap.is_empty() {
93                self.start_or_min_value = next_start;
94                self.end_heap.push(Reverse(next_end));
95                if let Some(result) = self.process(count.is_odd(), result) {
96                    return result;
97                }
98                continue;
99            }
100
101            // Next start is after the workspaces end, and the workspace contains more than one chuck,
102            // so process one chunk and then process next
103            self.start_or_min_value = end.add_one();
104            self.next_again = Some(next_start..=next_end);
105            if let Some(result) = self.process(count.is_odd(), result) {
106                return result;
107            }
108            // continue;
109        }
110    }
111}
112
113impl<T, I> SymDiffIter<T, I>
114where
115    T: Integer,
116    I: SortedStarts<T>,
117{
118    #[inline]
119    fn remove_same_end(&mut self, end: T) {
120        while let Some(end2) = self.end_heap.peek() {
121            if end2.0 == end {
122                self.end_heap.pop();
123            } else {
124                break;
125            }
126        }
127    }
128
129    #[inline]
130    #[allow(clippy::option_option)]
131    fn process(
132        &mut self,
133        keep: bool,
134        next: RangeInclusive<T>,
135    ) -> Option<Option<RangeInclusive<T>>> {
136        if !keep {
137            return None;
138        }
139        let Some(gather) = self.gather.take() else {
140            self.gather = Some(next);
141            return None;
142        };
143        // If there is no "next" then return gather if it exists.
144
145        // Take both next and gather apart.
146        let (next_start, next_end) = next.into_inner();
147        let (gather_start, gather_end) = gather.into_inner();
148
149        // We can assume gather_end < next_start.
150        debug_assert!(gather_end < next_start); // real assert
151
152        // If they touch, set gather to the union and loop.
153        if gather_end.add_one() == next_start {
154            self.gather = Some(gather_start..=next_end);
155            return None;
156        }
157
158        // Next is disjoint from gather, so return gather and set gather to next.
159        self.gather = Some(next_start..=next_end);
160        Some(Some(gather_start..=gather_end))
161    }
162
163    #[inline]
164    pub(crate) fn new(iter: I) -> Self {
165        Self {
166            iter,
167            start_or_min_value: T::min_value(),
168            end_heap: BinaryHeap::with_capacity(10),
169            next_again: None,
170            gather: None,
171        }
172    }
173}
174
175impl<T, L, R> SymDiffMerge<T, L, R>
176where
177    T: Integer,
178    L: SortedDisjoint<T>,
179    R: SortedDisjoint<T>,
180{
181    #[inline]
182    pub(crate) fn new2(left: L, right: R) -> Self {
183        let iter = Merge::new(left, right);
184        Self::new(iter)
185    }
186}
187
188impl<T, J> SymDiffKMerge<T, J>
189where
190    T: Integer,
191    J: SortedDisjoint<T>,
192{
193    #[inline]
194    pub(crate) fn new_k<K>(k: K) -> Self
195    where
196        K: IntoIterator<Item = J>,
197    {
198        let iter = KMerge::new(k);
199        Self::new(iter)
200    }
201}