futures_concurrency/stream/chain/
tuple.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
use core::fmt;
use core::pin::Pin;
use core::task::{Context, Poll};

use futures_core::Stream;

use super::Chain;

macro_rules! impl_chain_for_tuple {
    ($mod_name: ident $StructName:ident $($F:ident)+) => {
        mod $mod_name {
            #[repr(usize)]
            enum Indexes {
                $($F,)+
            }

            $(
                pub(super) const $F: usize = Indexes::$F as usize;
            )+

            pub(super) const LEN: usize = [$(Indexes::$F,)+].len();
        }

        #[pin_project::pin_project]
        pub struct $StructName<$($F,)+> {
            index: usize,
            done: bool,
            $( #[pin] $F: $F,)+
        }

        impl<T, $($F,)+> Stream for $StructName<$($F,)+>
        where
            $($F: Stream<Item = T>,)+
        {
            type Item = T;

            fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
                let mut this = self.project();

                assert!(!*this.done, "Stream should not be polled after completion");

                loop {
                    if *this.index == $mod_name::LEN {
                        *this.done = true;
                        return Poll::Ready(None);
                    }

                    match *this.index {
                        $(
                            $mod_name::$F => {
                                let fut = unsafe { Pin::new_unchecked(&mut this.$F) };
                                match fut.poll_next(cx) {
                                    Poll::Ready(None) => {
                                        *this.index += 1;
                                        continue;
                                    }
                                    v @ (Poll::Pending | Poll::Ready(Some(_))) => return v,
                                }
                            },
                        )+
                        _  => unreachable!(),
                    }
                }
            }
        }

        impl<$($F,)+> fmt::Debug for $StructName<$($F,)+>
        where
            $($F: fmt::Debug,)+
        {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                f.debug_tuple("Chain")
                    $(.field(&self.$F))+
                    .finish()
            }
        }

        impl<T, $($F,)+> Chain for ($($F,)+)
        where
            $($F: Stream<Item = T>,)+
        {
            type Item = T;

            type Stream = $StructName<$($F,)+>;

            fn chain(self) -> Self::Stream {
                let ($($F,)*): ($($F,)*) = self;
                Self::Stream {
                    done: false,
                    index: 0,
                    $($F,)+
                }
            }
        }
    }
}

impl_chain_for_tuple! { chain_1 Chain1 A }
impl_chain_for_tuple! { chain_2 Chain2 A B }
impl_chain_for_tuple! { chain_3 Chain3 A B C }
impl_chain_for_tuple! { chain_4 Chain4 A B C D }
impl_chain_for_tuple! { chain_5 Chain5 A B C D E }
impl_chain_for_tuple! { chain_6 Chain6 A B C D E F }
impl_chain_for_tuple! { chain_7 Chain7 A B C D E F G }
impl_chain_for_tuple! { chain_8 Chain8 A B C D E F G H }
impl_chain_for_tuple! { chain_9 Chain9 A B C D E F G H I }
impl_chain_for_tuple! { chain_10 Chain10 A B C D E F G H I J }
impl_chain_for_tuple! { chain_11 Chain11 A B C D E F G H I J K }
impl_chain_for_tuple! { chain_12 Chain12 A B C D E F G H I J K L }

#[cfg(test)]
mod tests {
    use super::*;

    use futures_lite::future::block_on;
    use futures_lite::prelude::*;
    use futures_lite::stream;

    #[test]
    fn chain_3() {
        block_on(async {
            let a = stream::once(1);
            let b = stream::once(2);
            let c = stream::once(3);
            let mut s = (a, b, c).chain();

            assert_eq!(s.next().await, Some(1));
            assert_eq!(s.next().await, Some(2));
            assert_eq!(s.next().await, Some(3));
            assert_eq!(s.next().await, None);
        })
    }
}