Skip to main content

range_set_blaze/
union_iter_map.rs

1use crate::map::ValueRef;
2use crate::merge_map::KMergeMap;
3use crate::sorted_disjoint_map::{Priority, PrioritySortedStartsMap};
4use crate::{AssumeSortedStarts, MergeMap, SortedDisjointMap, UnionKMergeMap, UnionMergeMap};
5use alloc::{collections::BinaryHeap, vec};
6use core::cmp::min;
7use core::iter::FusedIterator;
8use core::ops::RangeInclusive;
9use itertools::Itertools;
10
11use crate::Integer;
12use crate::unsorted_priority_map::AssumePrioritySortedStartsMap;
13use crate::unsorted_priority_map::UnsortedPriorityMap;
14
15type SortedStartsInVecMap<T, VR> = AssumePrioritySortedStartsMap<vec::IntoIter<Priority<T, VR>>>;
16#[allow(clippy::redundant_pub_crate)]
17pub(crate) type SortedStartsInVec<T> = AssumeSortedStarts<vec::IntoIter<RangeInclusive<T>>>;
18
19/// This `struct` is created by the [`union`] method. See [`union`]'s
20/// documentation for more.
21///
22/// [`union`]: crate::SortedDisjointMap::union
23#[derive(Clone, Debug)]
24#[must_use = "iterators are lazy and do nothing unless consumed"]
25pub struct UnionIterMap<T, VR, SS> {
26    iter: SS,
27    next_item: Option<Priority<T, VR>>,
28    workspace: BinaryHeap<Priority<T, VR>>,
29    gather: Option<(RangeInclusive<T>, VR)>,
30    ready_to_go: Option<(RangeInclusive<T>, VR)>,
31}
32
33impl<T, VR, I> Iterator for UnionIterMap<T, VR, I>
34where
35    T: Integer,
36    VR: ValueRef,
37    I: PrioritySortedStartsMap<T, VR>,
38{
39    type Item = (RangeInclusive<T>, VR);
40
41    fn next(&mut self) -> Option<(RangeInclusive<T>, VR)> {
42        // Keep doing this until we have something to return.
43        loop {
44            if let Some(value) = self.ready_to_go.take() {
45                // If ready_to_go is Some, return the value immediately.
46                return Some(value);
47            }
48
49            // if self.next_item should go into the workspace, then put it there, get the next, next_item, and loop
50            if let Some(next_item) = self.next_item.take() {
51                let (next_start, next_end) = next_item.start_and_end();
52
53                // If workspace is empty, just push the next item
54                let Some(best) = self.workspace.peek() else {
55                    self.workspace.push(next_item);
56                    self.next_item = self.iter.next();
57                    continue; // return to top of the main processing loop
58                };
59                // LATER: Could add this special case: If next value is the same as best value and the ending is later, and the start overlaps/touches, then just extend the best value.
60                if next_start == best.start() {
61                    // Only push if the priority is better or the end is greater
62                    if &next_item > best || next_end > best.end() {
63                        self.workspace.push(next_item);
64                    }
65                    self.next_item = self.iter.next();
66                    continue; // return to top of the main processing loop
67                }
68
69                // It does not go into the workspace, so just hold it and keep processing.
70                self.next_item = Some(next_item);
71            }
72
73            // If the workspace is empty, we are done.
74            let Some(best) = self.workspace.peek() else {
75                debug_assert!(self.next_item.is_none());
76                debug_assert!(self.ready_to_go.is_none());
77                return self.gather.take();
78            };
79
80            // We buffer for output the best item up to the start of the next item (if any).
81
82            // Find the start of the next item, if any.
83            let next_end = self.next_item.as_ref().map_or_else(
84                || best.end(),
85                |next_item| min(next_item.start().sub_one(), best.end()),
86            );
87
88            // Add the front of best to the gather buffer.
89            if let Some(mut gather) = self.gather.take() {
90                if gather.1.borrow() == best.value().borrow()
91                    && (*gather.0.end()).add_one() == best.start()
92                {
93                    // if the gather is contiguous with the best, then merge them
94                    gather.0 = *gather.0.start()..=next_end;
95                    self.gather = Some(gather);
96                } else {
97                    // if the gather is not contiguous with the best, then output the gather and set the gather to the best
98                    self.ready_to_go = Some(gather);
99                    self.gather = Some((best.start()..=next_end, best.value().clone()));
100                }
101            } else {
102                // if there is no gather, then set the gather to the best
103                self.gather = Some((best.start()..=next_end, best.value().clone()));
104            }
105
106            // We also update the workspace to removing any items that are completely covered by the new_start.
107            // We also don't need to keep any items that have a lower priority and are shorter than the new best.
108            let mut new_workspace = BinaryHeap::new();
109            while let Some(item) = self.workspace.pop() {
110                let mut item = item;
111                if item.end() <= next_end {
112                    // too short, don't keep
113                    continue; // while loop
114                }
115                item.set_range(next_end.add_one()..=item.end());
116                let Some(new_best) = new_workspace.peek() else {
117                    // new_workspace is empty, so keep
118                    new_workspace.push(item);
119                    continue; // while loop
120                };
121                if &item < new_best && item.end() <= new_best.end() {
122                    // item.priority, item.0, new_best.priority, new_best.0);
123                    // not as good as new_best, and shorter, so don't keep
124                    continue; // while loop
125                }
126
127                // higher priority or longer, so keep
128                // item.priority, item.0, new_best.priority, new_best.0);
129                new_workspace.push(item);
130            }
131            self.workspace = new_workspace;
132        } // end of main loop
133    }
134}
135
136impl<T, VR, I> UnionIterMap<T, VR, I>
137where
138    T: Integer,
139    VR: ValueRef,
140    I: PrioritySortedStartsMap<T, VR>,
141{
142    #[inline]
143    pub(crate) fn new(mut iter: I) -> Self {
144        let item = iter.next();
145        Self {
146            iter,
147            next_item: item,
148            workspace: BinaryHeap::new(),
149            gather: None,
150            ready_to_go: None,
151        }
152    }
153}
154
155impl<T, VR, L, R> UnionMergeMap<T, VR, L, R>
156where
157    T: Integer,
158    VR: ValueRef,
159    L: SortedDisjointMap<T, VR>,
160    R: SortedDisjointMap<T, VR>,
161{
162    #[inline]
163    pub(crate) fn new2(left: L, right: R) -> Self {
164        let iter = MergeMap::new(left, right);
165        Self::new(iter)
166    }
167}
168
169impl<T, VR, J> UnionKMergeMap<T, VR, J>
170where
171    T: Integer,
172    VR: ValueRef,
173    J: SortedDisjointMap<T, VR>,
174{
175    #[inline]
176    pub(crate) fn new_k<K>(k: K) -> Self
177    where
178        K: IntoIterator<Item = J>,
179    {
180        let iter = KMergeMap::new(k);
181        Self::new(iter)
182    }
183}
184
185// UnionIterMap from iter (T, VR)
186impl<T, VR> FromIterator<(RangeInclusive<T>, VR)>
187    for UnionIterMap<T, VR, SortedStartsInVecMap<T, VR>>
188where
189    T: Integer,
190    VR: ValueRef,
191{
192    fn from_iter<I>(iter: I) -> Self
193    where
194        I: IntoIterator<Item = (RangeInclusive<T>, VR)>,
195    {
196        let iter = iter.into_iter();
197        let iter = UnsortedPriorityMap::new(iter);
198        // We sort only by start -- priority is not used until later.
199        let iter = iter.sorted_by(|a, b| a.start().cmp(&b.start()));
200        let iter = AssumePrioritySortedStartsMap::new(iter);
201        Self::new(iter)
202    }
203}
204
205impl<T, VR, I> FusedIterator for UnionIterMap<T, VR, I>
206where
207    T: Integer,
208    VR: ValueRef,
209    I: PrioritySortedStartsMap<T, VR> + FusedIterator,
210{
211}