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#[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 loop {
44 if let Some(value) = self.ready_to_go.take() {
45 return Some(value);
47 }
48
49 if let Some(next_item) = self.next_item.take() {
51 let (next_start, next_end) = next_item.start_and_end();
52
53 let Some(best) = self.workspace.peek() else {
55 self.workspace.push(next_item);
56 self.next_item = self.iter.next();
57 continue; };
59 if next_start == best.start() {
61 if &next_item > best || next_end > best.end() {
63 self.workspace.push(next_item);
64 }
65 self.next_item = self.iter.next();
66 continue; }
68
69 self.next_item = Some(next_item);
71 }
72
73 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 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 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 gather.0 = *gather.0.start()..=next_end;
95 self.gather = Some(gather);
96 } else {
97 self.ready_to_go = Some(gather);
99 self.gather = Some((best.start()..=next_end, best.value().clone()));
100 }
101 } else {
102 self.gather = Some((best.start()..=next_end, best.value().clone()));
104 }
105
106 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 continue; }
115 item.set_range(next_end.add_one()..=item.end());
116 let Some(new_best) = new_workspace.peek() else {
117 new_workspace.push(item);
119 continue; };
121 if &item < new_best && item.end() <= new_best.end() {
122 continue; }
126
127 new_workspace.push(item);
130 }
131 self.workspace = new_workspace;
132 } }
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
185impl<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 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}