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 option_entry::{Entry, OptionEntry};
8use pin_project_lite::pin_project;
9
10use crate::{EitherOrBoth, internal::check::assert_stream};
11
12/// See [`crate::AsyncItertools::zip_longest`].
13pub fn zip_longest<L: Stream, R: Stream>(l: L, r: R) -> crate::ZipLongest<L, R> {
14    assert_stream(ZipLongest {
15        l: l.fuse(),
16        r: r.fuse(),
17        left_item: None,
18    })
19}
20
21pin_project! {
22    pub struct ZipLongest<L, R, Lt = <L as Stream>::Item> {
23        #[pin]
24        l: Fuse<L>,
25        #[pin]
26        r: Fuse<R>,
27        left_item: Option<Lt>,
28    }
29}
30
31impl<L: Stream<Item = Lt>, R: Stream<Item = Rt>, Lt, Rt> Stream for ZipLongest<L, R, Lt> {
32    type Item = EitherOrBoth<Lt, Rt>;
33
34    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
35        let this = self.project();
36        let l = match this.left_item.entry() {
37            Entry::Vacant(entry) => {
38                let Some(l) = ready!(this.l.poll_next(cx)) else {
39                    return Poll::Ready(ready!(this.r.poll_next(cx)).map(EitherOrBoth::Right));
40                };
41                entry.insert_entry(l)
42            }
43            Entry::Occupied(entry) => entry,
44        };
45        let Some(r) = ready!(this.r.poll_next(cx)) else {
46            return Poll::Ready(Some(EitherOrBoth::Left(l.remove())));
47        };
48        Poll::Ready(Some(EitherOrBoth::Both(l.remove(), r)))
49    }
50}
51
52#[cfg(test)]
53mod test {
54    use crate::{AsyncItertools, EitherOrBoth};
55
56    #[test]
57    fn left_longer() {
58        let l = futures_lite::stream::iter([1, 2, 3]);
59        let r = futures_lite::stream::iter([1, 2]);
60        assert!(futures_lite::stream::block_on(l.zip_longest(r)).eq([
61            EitherOrBoth::Both(1, 1),
62            EitherOrBoth::Both(2, 2),
63            EitherOrBoth::Left(3),
64        ]));
65    }
66
67    #[test]
68    fn equal_length() {
69        let l = futures_lite::stream::iter([1, 2, 3]);
70        let r = futures_lite::stream::iter([1, 2, 3]);
71        assert!(futures_lite::stream::block_on(l.zip_longest(r)).eq([
72            EitherOrBoth::Both(1, 1),
73            EitherOrBoth::Both(2, 2),
74            EitherOrBoth::Both(3, 3),
75        ]));
76    }
77
78    #[test]
79    fn right_longer() {
80        let l = futures_lite::stream::iter([1, 2]);
81        let r = futures_lite::stream::iter([1, 2, 3]);
82        assert!(futures_lite::stream::block_on(l.zip_longest(r)).eq([
83            EitherOrBoth::Both(1, 1),
84            EitherOrBoth::Both(2, 2),
85            EitherOrBoth::Right(3),
86        ]));
87    }
88}
async_itertools/features/zip_longest.rs

async_itertools/features/
zip_longest.rs
