1use super::*;
3use crate::sorted_iterator::SortedByItem;
4use core::{
5 cmp::Ordering::*,
6 cmp::{max, min},
7 fmt::Debug,
8 iter,
9 iter::Peekable,
10 option, result,
11};
12use std::{collections, iter::FusedIterator};
13
14pub trait SortedByKey {}
16
17pub struct Join<I: Iterator, J: Iterator> {
18 pub(crate) a: Peekable<I>,
19 pub(crate) b: Peekable<J>,
20}
21
22impl<K, A, B, I, J> Iterator for Join<I, J>
23where
24 K: Ord,
25 I: Iterator<Item = (K, A)> + SortedByKey,
26 J: Iterator<Item = (K, B)> + SortedByKey,
27{
28 type Item = (K, (A, B));
29
30 fn next(&mut self) -> Option<Self::Item> {
31 while let (Some((ak, _)), Some((bk, _))) = (self.a.peek(), self.b.peek()) {
32 match ak.cmp(&bk) {
33 Less => {
34 self.a.next();
35 }
36 Greater => {
37 self.b.next();
38 }
39 Equal => {
40 if let (Some((ak, av)), Some((_, bv))) = (self.a.next(), self.b.next()) {
41 return Some((ak, (av, bv)));
42 } else {
43 unreachable!();
44 }
45 }
46 }
47 }
48 None
49 }
50
51 fn size_hint(&self) -> (usize, Option<usize>) {
52 let (_, amax) = self.a.size_hint();
53 let (_, bmax) = self.b.size_hint();
54 let rmin = 0;
55 let rmax = amax.and_then(|amax| bmax.map(|bmax| min(amax, bmax)));
56 (rmin, rmax)
57 }
58}
59
60impl<I: Iterator + Clone, J: Iterator + Clone> Clone for Join<I, J>
61where
62 I::Item: Clone,
63 J::Item: Clone,
64{
65 fn clone(&self) -> Self {
66 Self {
67 a: self.a.clone(),
68 b: self.b.clone(),
69 }
70 }
71}
72
73pub struct LeftJoin<I: Iterator, J: Iterator> {
74 pub(crate) a: Peekable<I>,
75 pub(crate) b: Peekable<J>,
76}
77
78impl<K, A, B, I, J> Iterator for LeftJoin<I, J>
79where
80 K: Ord,
81 I: Iterator<Item = (K, A)>,
82 J: Iterator<Item = (K, B)>,
83{
84 type Item = (K, (A, Option<B>));
85
86 fn next(&mut self) -> Option<Self::Item> {
87 let (ak, av) = self.a.next()?;
88 while let Some((bk, _)) = self.b.peek() {
89 match ak.cmp(bk) {
90 Less => break,
91 Greater => {
92 self.b.next();
93 }
94 Equal => {
95 let (_, bv) = self.b.next()?;
96 return Some((ak, (av, Some(bv))));
97 }
98 }
99 }
100 Some((ak, (av, None)))
101 }
102
103 fn size_hint(&self) -> (usize, Option<usize>) {
104 self.a.size_hint()
105 }
106}
107
108impl<I: Iterator + Clone, J: Iterator + Clone> Clone for LeftJoin<I, J>
109where
110 I::Item: Clone,
111 J::Item: Clone,
112{
113 fn clone(&self) -> Self {
114 Self {
115 a: self.a.clone(),
116 b: self.b.clone(),
117 }
118 }
119}
120
121pub struct RightJoin<I: Iterator, J: Iterator> {
122 pub(crate) a: Peekable<I>,
123 pub(crate) b: Peekable<J>,
124}
125
126impl<K, A, B, I, J> Iterator for RightJoin<I, J>
127where
128 K: Ord,
129 I: Iterator<Item = (K, A)>,
130 J: Iterator<Item = (K, B)>,
131{
132 type Item = (K, (Option<A>, B));
133
134 fn next(&mut self) -> Option<Self::Item> {
135 let (bk, bv) = self.b.next()?;
136 while let Some((ak, _)) = self.a.peek() {
137 match bk.cmp(ak) {
138 Less => break,
139 Greater => {
140 self.a.next();
141 }
142 Equal => {
143 let (_, av) = self.a.next()?;
144 return Some((bk, (Some(av), bv)));
145 }
146 }
147 }
148 Some((bk, (None, bv)))
149 }
150
151 fn size_hint(&self) -> (usize, Option<usize>) {
152 self.b.size_hint()
153 }
154}
155
156impl<I: Iterator + Clone, J: Iterator + Clone> Clone for RightJoin<I, J>
157where
158 I::Item: Clone,
159 J::Item: Clone,
160{
161 fn clone(&self) -> Self {
162 Self {
163 a: self.a.clone(),
164 b: self.b.clone(),
165 }
166 }
167}
168
169pub struct OuterJoin<I: Iterator, J: Iterator> {
170 pub(crate) a: Peekable<I>,
171 pub(crate) b: Peekable<J>,
172}
173
174impl<K, A, B, I, J> OuterJoin<I, J>
177where
178 K: Ord,
179 I: Iterator<Item = (K, A)>,
180 J: Iterator<Item = (K, B)>,
181{
182 #[allow(clippy::type_complexity)]
184 fn next_a(&mut self) -> Option<(K, (Option<A>, Option<B>))> {
185 self.a.next().map(|(ak, av)| (ak, (Some(av), None)))
186 }
187
188 #[allow(clippy::type_complexity)]
189 fn next_b(&mut self) -> Option<(K, (Option<A>, Option<B>))> {
190 self.b.next().map(|(bk, bv)| (bk, (None, Some(bv))))
191 }
192}
193
194impl<K, A, B, I, J> Iterator for OuterJoin<I, J>
195where
196 K: Ord,
197 I: Iterator<Item = (K, A)>,
198 J: Iterator<Item = (K, B)>,
199{
200 type Item = (K, (Option<A>, Option<B>));
201
202 fn next(&mut self) -> Option<Self::Item> {
203 if let (Some((ak, _)), Some((bk, _))) = (self.a.peek(), self.b.peek()) {
204 match ak.cmp(&bk) {
205 Less => self.next_a(),
206 Greater => self.next_b(),
207 Equal => self
208 .a
209 .next()
210 .and_then(|(ak, av)| self.b.next().map(|(_, bv)| (ak, (Some(av), Some(bv))))),
211 }
212 } else {
213 self.next_a().or_else(|| self.next_b())
214 }
215 }
216
217 fn size_hint(&self) -> (usize, Option<usize>) {
218 let (amin, amax) = self.a.size_hint();
219 let (bmin, bmax) = self.b.size_hint();
220 let rmin = max(amin, bmin);
221 let rmax = amax.and_then(|amax| bmax.map(|bmax| amax + bmax));
222 (rmin, rmax)
223 }
224}
225
226impl<I: Iterator + Clone, J: Iterator + Clone> Clone for OuterJoin<I, J>
227where
228 I::Item: Clone,
229 J::Item: Clone,
230{
231 fn clone(&self) -> Self {
232 Self {
233 a: self.a.clone(),
234 b: self.b.clone(),
235 }
236 }
237}
238
239#[derive(Clone, Debug)]
240pub struct Keys<I: Iterator> {
241 pub(crate) i: I,
242}
243
244impl<K, V, I> Iterator for Keys<I>
245where
246 K: Ord,
247 I: Iterator<Item = (K, V)>,
248{
249 type Item = K;
250
251 fn next(&mut self) -> Option<Self::Item> {
252 self.i.next().map(|(k, _)| k)
253 }
254
255 fn size_hint(&self) -> (usize, Option<usize>) {
256 self.i.size_hint()
257 }
258}
259impl<K: Ord, V, I: ExactSizeIterator + Iterator<Item = (K, V)>> ExactSizeIterator for Keys<I> {
260 fn len(&self) -> usize {
261 self.i.len()
262 }
263}
264impl<K: Ord, V, I: DoubleEndedIterator + Iterator<Item = (K, V)>> DoubleEndedIterator for Keys<I> {
265 fn next_back(&mut self) -> Option<Self::Item> {
266 self.i.next_back().map(|(k, _)| k)
267 }
268}
269
270#[derive(Clone, Debug)]
271pub struct MapValues<I: Iterator, F> {
272 pub(crate) i: I,
273 pub(crate) f: F,
274}
275
276impl<K, V, W, I, F> Iterator for MapValues<I, F>
277where
278 K: Ord,
279 I: Iterator<Item = (K, V)>,
280 F: FnMut(V) -> W,
281{
282 type Item = (K, W);
283
284 fn next(&mut self) -> Option<Self::Item> {
285 self.i.next().map(|(k, v)| (k, (self.f)(v)))
286 }
287
288 fn size_hint(&self) -> (usize, Option<usize>) {
289 self.i.size_hint()
290 }
291}
292impl<K: Ord, V, W, I, F> ExactSizeIterator for MapValues<I, F>
293where
294 I: ExactSizeIterator + Iterator<Item = (K, V)>,
295 F: FnMut(V) -> W,
296{
297 fn len(&self) -> usize {
298 self.i.len()
299 }
300}
301impl<K: Ord, V, W, I, F> DoubleEndedIterator for MapValues<I, F>
302where
303 I: DoubleEndedIterator + Iterator<Item = (K, V)>,
304 F: FnMut(V) -> W,
305{
306 fn next_back(&mut self) -> Option<Self::Item> {
307 self.i.next_back().map(|(k, v)| (k, (self.f)(v)))
308 }
309}
310
311#[derive(Clone, Debug)]
312pub struct FilterMapValues<I: Iterator, F> {
313 pub(crate) i: I,
314 pub(crate) f: F,
315}
316
317impl<K, V, W, I, F> Iterator for FilterMapValues<I, F>
318where
319 K: Ord,
320 I: Iterator<Item = (K, V)>,
321 F: FnMut(V) -> Option<W>,
322{
323 type Item = (K, W);
324
325 fn next(&mut self) -> Option<Self::Item> {
326 while let Some((k, v)) = self.i.next() {
327 if let Some(w) = (self.f)(v) {
328 return Some((k, w));
329 }
330 }
331 None
332 }
333
334 fn size_hint(&self) -> (usize, Option<usize>) {
335 let (_, imax) = self.i.size_hint();
336 (0, imax)
337 }
338}
339
340#[derive(Clone, Debug)]
341pub struct AssumeSortedByKey<I: Iterator> {
342 pub(crate) i: I,
343}
344
345impl<I: Iterator> Iterator for AssumeSortedByKey<I> {
346 type Item = I::Item;
347
348 fn next(&mut self) -> Option<Self::Item> {
349 self.i.next()
350 }
351
352 fn size_hint(&self) -> (usize, Option<usize>) {
353 self.i.size_hint()
354 }
355}
356
357impl<I: Iterator> ExactSizeIterator for AssumeSortedByKey<I> where I: ExactSizeIterator {}
358
359impl<I: Iterator> FusedIterator for AssumeSortedByKey<I> where I: FusedIterator {}
360
361impl<I: Iterator> DoubleEndedIterator for AssumeSortedByKey<I>
362where
363 I: DoubleEndedIterator,
364{
365 fn next_back(&mut self) -> Option<Self::Item> {
366 self.i.next_back()
367 }
368}
369
370impl<I> SortedByKey for iter::Empty<I> {}
376impl<I> SortedByKey for iter::Once<I> {}
377impl<'a, T> SortedByKey for option::Iter<'a, T> {}
378impl<'a, T> SortedByKey for result::Iter<'a, T> {}
379impl<T> SortedByKey for option::IntoIter<T> {}
380impl<T> SortedByKey for result::IntoIter<T> {}
381impl<I> SortedByKey for iter::Enumerate<I> {}
382
383impl<I: Iterator + SortedByItem> SortedByKey for Pairs<I> {}
384impl<I: Iterator, F> SortedByKey for MapValues<I, F> {}
385impl<I: Iterator, F> SortedByKey for FilterMapValues<I, F> {}
386
387impl<I: SortedByKey> SortedByKey for iter::Take<I> {}
388impl<I: SortedByKey> SortedByKey for iter::Skip<I> {}
389impl<I: SortedByKey> SortedByKey for iter::StepBy<I> {}
390impl<I: SortedByKey> SortedByKey for iter::Cloned<I> {}
391impl<I: SortedByKey> SortedByKey for iter::Copied<I> {}
392impl<I: SortedByKey> SortedByKey for iter::Fuse<I> {}
393impl<I: SortedByKey, F> SortedByKey for iter::Inspect<I, F> {}
394impl<I: SortedByKey, P> SortedByKey for iter::TakeWhile<I, P> {}
395impl<I: SortedByKey, P> SortedByKey for iter::SkipWhile<I, P> {}
396impl<I: SortedByKey, P> SortedByKey for iter::Filter<I, P> {}
397impl<I: SortedByKey + Iterator> SortedByKey for iter::Peekable<I> {}
398impl<I: SortedByItem, J> SortedByKey for iter::Zip<I, J> {}
399
400impl<I: Iterator, J: Iterator> SortedByKey for Join<I, J> {}
401impl<I: Iterator, J: Iterator> SortedByKey for LeftJoin<I, J> {}
402impl<I: Iterator, J: Iterator> SortedByKey for RightJoin<I, J> {}
403impl<I: Iterator, J: Iterator> SortedByKey for OuterJoin<I, J> {}
404impl<I: Iterator> SortedByKey for AssumeSortedByKey<I> {}
405
406impl<K, V> SortedByKey for collections::btree_map::IntoIter<K, V> {}
407impl<'a, K, V> SortedByKey for collections::btree_map::Iter<'a, K, V> {}
408impl<'a, K, V> SortedByKey for collections::btree_map::IterMut<'a, K, V> {}
409impl<'a, K, V> SortedByKey for collections::btree_map::Range<'a, K, V> {}
410impl<'a, K, V> SortedByKey for collections::btree_map::RangeMut<'a, K, V> {}
411
412impl<I: SortedByKey> SortedByKey for Box<I> {}
413
414impl<I: OneOrLess, F> SortedByKey for iter::Map<I, F> {}
415impl<Iin, J, Iout, F> SortedByKey for iter::FlatMap<Iin, J, F>
416where
417 Iin: OneOrLess,
418 J: IntoIterator<IntoIter = Iout>,
419 Iout: SortedByKey,
420{
421}
422impl<Iin, J, Iout> SortedByKey for iter::Flatten<Iin>
423where
424 Iin: OneOrLess + Iterator<Item = J>,
425 J: IntoIterator<IntoIter = Iout>,
426 Iout: SortedByKey,
427{
428}
429
430#[cfg(test)]
431mod tests {
432 extern crate maplit;
433 use super::*;
434 use collections::BTreeMap;
435 use core::fmt::Debug;
436 use maplit::*;
437
438 fn unary_op<E: Debug, R: Eq + Debug>(x: E, expected: R, actual: R) -> bool {
440 let res = expected == actual;
441 if !res {
442 println!("x:{:?} expected:{:?} actual:{:?}", x, expected, actual);
443 }
444 res
445 }
446
447 fn binary_op<E: Debug, R: Eq + Debug>(a: E, b: E, expected: R, actual: R) -> bool {
449 let res = expected == actual;
450 if !res {
451 println!(
452 "a:{:?} b:{:?} expected:{:?} actual:{:?}",
453 a, b, expected, actual
454 );
455 }
456 res
457 }
458
459 type Element = i64;
460 type Reference = BTreeMap<Element, Element>;
461
462 #[quickcheck]
463 fn join(a: Reference, b: Reference) -> bool {
464 type Result = BTreeMap<Element, (Element, Element)>;
465 let expected: Result = a
466 .keys()
467 .intersection(b.keys())
468 .map(|k| (k.clone(), (a[k], b[k])))
469 .collect();
470 let actual: Result = a.clone().into_iter().join(b.clone().into_iter()).collect();
471 binary_op(a, b, expected, actual)
472 }
473
474 #[quickcheck]
475 fn left_join(a: Reference, b: Reference) -> bool {
476 type Result = BTreeMap<Element, (Element, Option<Element>)>;
477 let expected: Result = a
478 .keys()
479 .map(|k| (k.clone(), (a[k], b.get(k).cloned())))
480 .collect();
481 let actual: Result = a
482 .clone()
483 .into_iter()
484 .left_join(b.clone().into_iter())
485 .collect();
486 binary_op(a, b, expected, actual)
487 }
488
489 #[quickcheck]
490 fn right_join(a: Reference, b: Reference) -> bool {
491 type Result = BTreeMap<Element, (Option<Element>, Element)>;
492 let expected: Result = b
493 .keys()
494 .map(|k| (k.clone(), (a.get(k).cloned(), b[k])))
495 .collect();
496 let actual: Result = a
497 .clone()
498 .into_iter()
499 .right_join(b.clone().into_iter())
500 .collect();
501 binary_op(a, b, expected, actual)
502 }
503
504 #[quickcheck]
505 fn outer_join(a: Reference, b: Reference) -> bool {
506 type Result = BTreeMap<Element, (Option<Element>, Option<Element>)>;
507 let expected: Result = a
508 .keys()
509 .union(b.keys())
510 .map(|k| (k.clone(), (a.get(k).cloned(), b.get(k).cloned())))
511 .collect();
512 let actual: Result = a
513 .clone()
514 .into_iter()
515 .outer_join(b.clone().into_iter())
516 .collect();
517 binary_op(a, b, expected, actual)
518 }
519
520 #[quickcheck]
521 fn map_values(x: Reference) -> bool {
522 type Result = BTreeMap<Element, Element>;
523 let expected: Result = x.clone().into_iter().map(|(k, v)| (k, v * 2)).collect();
524 let actual: Result = x.clone().into_iter().map_values(|v| v * 2).collect();
525 unary_op(x, expected, actual)
526 }
527
528 #[quickcheck]
529 fn filter_map_values(x: Reference) -> bool {
530 type Result = BTreeMap<Element, Element>;
531 let expected: Result = x
532 .clone()
533 .into_iter()
534 .filter_map(|(k, v)| if v % 2 != 0 { Some((k, v * 2)) } else { None })
535 .collect();
536 let actual: Result = x
537 .clone()
538 .into_iter()
539 .filter_map_values(|v| if v % 2 != 0 { Some(v * 2) } else { None })
540 .collect();
541 unary_op(x, expected, actual)
542 }
543
544 fn is_s<K, V, I: Iterator<Item = (K, V)> + SortedByKey>(_v: I) {}
545
546 fn s() -> impl Iterator<Item = (i64, ())> + SortedByKey {
547 (0i64..10).pairs()
548 }
549
550 #[test]
551 fn instances() {
552 is_s(iter::empty::<(i64, ())>());
554 is_s(iter::once((0, ())));
555 is_s([1, 2, 3, 4].iter().enumerate());
556 is_s((0i64..10).pairs());
558 is_s((0i64..=10).pairs());
559 is_s((0i64..).pairs());
560 is_s(Box::new((0i64..).pairs()));
562 is_s(s().step_by(1));
564 is_s(s().take(1));
565 is_s(s().skip(1));
566 is_s(s().take_while(|_| true));
567 is_s(s().skip_while(|_| true));
568 is_s(s().filter(|_| true));
569 is_s(s().peekable());
571 is_s(s().fuse());
572 is_s(s().inspect(|_| {}));
573 is_s(s().join(s()));
575 is_s(s().left_join(s()));
576 is_s(s().right_join(s()));
577 is_s(s().outer_join(s()));
578 is_s(btreemap! { 0i64 => "" }.iter());
580 is_s(btreemap! { 0i64 => "" }.into_iter());
581 is_s(btreemap! { 0i64 => "" }.iter_mut());
582 is_s(btreemap! { 0i64 => "" }.range(..));
583 is_s(btreemap! { 0i64 => "" }.range_mut(..));
584 let a_btree = BTreeMap::<i64, f32>::new();
586 is_s(
587 Some(())
588 .iter()
589 .map(|_| &a_btree)
590 .filter(|b| !b.is_empty())
591 .flatten(),
592 );
593 is_s(
594 iter::once(Result::<i64, f32>::Ok(12))
595 .flatten()
596 .map(|_| &a_btree)
597 .filter(|b| !b.is_empty())
598 .flat_map(|m| m.iter()),
599 );
600 }
601}