sorted_iter/
sorted_pair_iterator.rs

1//! implementation of the sorted_pair_iterator relational operations
2use 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
14/// marker trait for iterators that are sorted by the key of their Item
15pub 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
174// all this just so I could avoid having this expression twice in the iterator.
175// Sometimes making things DRY in rust is hard...
176impl<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    // type alias does not help much here...
183    #[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
370// not sure how to model this: an iterator that is sorted by key automatically sorted by item as well,
371// since the key comes first in the order of a pair. The reverse is not true, since e.g. (0,1), (0,2)
372// are sorted by item, but not strictly sorted by key!
373
374// mark common std traits
375impl<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    /// just a helper to get good output when a check fails
439    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    /// just a helper to get good output when a check fails
448    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        // creation
553        is_s(iter::empty::<(i64, ())>());
554        is_s(iter::once((0, ())));
555        is_s([1, 2, 3, 4].iter().enumerate());
556        // ranges
557        is_s((0i64..10).pairs());
558        is_s((0i64..=10).pairs());
559        is_s((0i64..).pairs());
560        // wrappers
561        is_s(Box::new((0i64..).pairs()));
562        // skip/take/step/filter
563        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        // identity
570        is_s(s().peekable());
571        is_s(s().fuse());
572        is_s(s().inspect(|_| {}));
573        // relational
574        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        // btreeset
579        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        // one_or_less
585        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}