Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Implement Chain for tuples up to 12 elements #124

Merged
merged 1 commit into from
Apr 5, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
132 changes: 132 additions & 0 deletions src/stream/chain/tuple.rs
Original file line number Diff line number Diff line change
@@ -1 +1,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);
})
}
}