async_itertools/features/
zip_longest.rs

1use core::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use futures_lite::{Stream, StreamExt, ready, stream::Fuse};
7use pin_project_lite::pin_project;
8
9use crate::{
10    EitherOrBoth,
11    internal::{
12        check::assert_stream,
13        either_or_none::{EitherOrNone, Entry},
14    },
15};
16
17/// See [`crate::AsyncItertools::zip_longest`].
18pub fn zip_longest<L: Stream, R: Stream>(l: L, r: R) -> crate::ZipLongest<L, R> {
19    assert_stream(ZipLongest {
20        l: l.fuse(),
21        r: r.fuse(),
22        ready: EitherOrNone::None,
23    })
24}
25
26pin_project! {
27    pub struct ZipLongest<L, R, Lt = <L as Stream>::Item, Rt = <R as Stream>::Item> {
28        #[pin]
29        l: Fuse<L>,
30        #[pin]
31        r: Fuse<R>,
32        ready: EitherOrNone<Lt, Rt>,
33    }
34}
35
36impl<L: Stream<Item = Lt>, R: Stream<Item = Rt>, Lt, Rt> Stream for ZipLongest<L, R, Lt, Rt> {
37    type Item = EitherOrBoth<Lt, Rt>;
38
39    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
40        let this = self.project();
41        match this.ready.entry() {
42            Entry::None(entry) => {
43                let l = this.l.poll_next(cx);
44                let r = this.r.poll_next(cx);
45                match (l, r) {
46                    (Poll::Ready(Some(l)), Poll::Ready(Some(r))) => {
47                        Poll::Ready(Some(EitherOrBoth::Both(l, r)))
48                    }
49                    (Poll::Ready(None), Poll::Ready(Some(r))) => {
50                        Poll::Ready(Some(EitherOrBoth::Right(r)))
51                    }
52                    (Poll::Ready(Some(l)), Poll::Ready(None)) => {
53                        Poll::Ready(Some(EitherOrBoth::Left(l)))
54                    }
55                    (Poll::Ready(None), Poll::Ready(None)) => Poll::Ready(None),
56                    (Poll::Ready(Some(l)), Poll::Pending) => {
57                        entry.insert_left(l);
58                        Poll::Pending
59                    }
60                    (Poll::Pending, Poll::Ready(Some(r))) => {
61                        entry.insert_right(r);
62                        Poll::Pending
63                    }
64                    _ => Poll::Pending,
65                }
66            }
67            Entry::Left(entry) => match ready!(this.r.poll_next(cx)) {
68                Some(r) => Poll::Ready(Some(EitherOrBoth::Both(entry.remove(), r))),
69                None => Poll::Ready(Some(EitherOrBoth::Left(entry.remove()))),
70            },
71            Entry::Right(entry) => match ready!(this.l.poll_next(cx)) {
72                Some(l) => Poll::Ready(Some(EitherOrBoth::Both(l, entry.remove()))),
73                None => Poll::Ready(Some(EitherOrBoth::Right(entry.remove()))),
74            },
75        }
76    }
77}
78
79#[cfg(test)]
80mod test {
81    use crate::{AsyncItertools, EitherOrBoth};
82
83    #[test]
84    fn left_longer() {
85        let l = futures_lite::stream::iter([1, 2, 3]);
86        let r = futures_lite::stream::iter([1, 2]);
87        assert!(futures_lite::stream::block_on(l.zip_longest(r)).eq([
88            EitherOrBoth::Both(1, 1),
89            EitherOrBoth::Both(2, 2),
90            EitherOrBoth::Left(3),
91        ]));
92    }
93
94    #[test]
95    fn equal_length() {
96        let l = futures_lite::stream::iter([1, 2, 3]);
97        let r = futures_lite::stream::iter([1, 2, 3]);
98        assert!(futures_lite::stream::block_on(l.zip_longest(r)).eq([
99            EitherOrBoth::Both(1, 1),
100            EitherOrBoth::Both(2, 2),
101            EitherOrBoth::Both(3, 3),
102        ]));
103    }
104
105    #[test]
106    fn right_longer() {
107        let l = futures_lite::stream::iter([1, 2]);
108        let r = futures_lite::stream::iter([1, 2, 3]);
109        assert!(futures_lite::stream::block_on(l.zip_longest(r)).eq([
110            EitherOrBoth::Both(1, 1),
111            EitherOrBoth::Both(2, 2),
112            EitherOrBoth::Right(3),
113        ]));
114    }
115}
async_itertools/features/zip_longest.rs

async_itertools/features/
zip_longest.rs
