itertools/
zip_longest.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
use super::size_hint;
use std::cmp::Ordering::{Equal, Greater, Less};
use std::iter::{Fuse, FusedIterator};

use crate::either_or_both::EitherOrBoth;

// ZipLongest originally written by SimonSapin,
// and dedicated to itertools https://github.com/rust-lang/rust/pull/19283

/// An iterator which iterates two other iterators simultaneously
/// and wraps the elements in [`EitherOrBoth`].
///
/// This iterator is *fused*.
///
/// See [`.zip_longest()`](crate::Itertools::zip_longest) for more information.
#[derive(Clone, Debug)]
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct ZipLongest<T, U> {
    a: Fuse<T>,
    b: Fuse<U>,
}

/// Create a new `ZipLongest` iterator.
pub fn zip_longest<T, U>(a: T, b: U) -> ZipLongest<T, U>
where
    T: Iterator,
    U: Iterator,
{
    ZipLongest {
        a: a.fuse(),
        b: b.fuse(),
    }
}

impl<T, U> Iterator for ZipLongest<T, U>
where
    T: Iterator,
    U: Iterator,
{
    type Item = EitherOrBoth<T::Item, U::Item>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match (self.a.next(), self.b.next()) {
            (None, None) => None,
            (Some(a), None) => Some(EitherOrBoth::Left(a)),
            (None, Some(b)) => Some(EitherOrBoth::Right(b)),
            (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        size_hint::max(self.a.size_hint(), self.b.size_hint())
    }

    #[inline]
    fn fold<B, F>(self, init: B, mut f: F) -> B
    where
        Self: Sized,
        F: FnMut(B, Self::Item) -> B,
    {
        let Self { mut a, mut b } = self;
        let res = a.try_fold(init, |init, a| match b.next() {
            Some(b) => Ok(f(init, EitherOrBoth::Both(a, b))),
            None => Err(f(init, EitherOrBoth::Left(a))),
        });
        match res {
            Ok(acc) => b.map(EitherOrBoth::Right).fold(acc, f),
            Err(acc) => a.map(EitherOrBoth::Left).fold(acc, f),
        }
    }
}

impl<T, U> DoubleEndedIterator for ZipLongest<T, U>
where
    T: DoubleEndedIterator + ExactSizeIterator,
    U: DoubleEndedIterator + ExactSizeIterator,
{
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        match self.a.len().cmp(&self.b.len()) {
            Equal => match (self.a.next_back(), self.b.next_back()) {
                (None, None) => None,
                (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
                // These can only happen if .len() is inconsistent with .next_back()
                (Some(a), None) => Some(EitherOrBoth::Left(a)),
                (None, Some(b)) => Some(EitherOrBoth::Right(b)),
            },
            Greater => self.a.next_back().map(EitherOrBoth::Left),
            Less => self.b.next_back().map(EitherOrBoth::Right),
        }
    }

    fn rfold<B, F>(self, mut init: B, mut f: F) -> B
    where
        F: FnMut(B, Self::Item) -> B,
    {
        let Self { mut a, mut b } = self;
        let a_len = a.len();
        let b_len = b.len();
        match a_len.cmp(&b_len) {
            Equal => {}
            Greater => {
                init = a
                    .by_ref()
                    .rev()
                    .take(a_len - b_len)
                    .map(EitherOrBoth::Left)
                    .fold(init, &mut f)
            }
            Less => {
                init = b
                    .by_ref()
                    .rev()
                    .take(b_len - a_len)
                    .map(EitherOrBoth::Right)
                    .fold(init, &mut f)
            }
        }
        a.rfold(init, |acc, item_a| {
            f(acc, EitherOrBoth::Both(item_a, b.next_back().unwrap()))
        })
    }
}

impl<T, U> ExactSizeIterator for ZipLongest<T, U>
where
    T: ExactSizeIterator,
    U: ExactSizeIterator,
{
}

impl<T, U> FusedIterator for ZipLongest<T, U>
where
    T: Iterator,
    U: Iterator,
{
}