use super::assert_stream;
use crate::stream::{select_with_strategy, PollNext, SelectWithStrategy};
use core::pin::Pin;
use futures_core::stream::{FusedStream, Stream};
use futures_core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Select<St1, St2> {
#[pin]
inner: SelectWithStrategy<St1, St2, fn(&mut PollNext)-> PollNext, PollNext>,
}
}
pub fn select<St1, St2>(stream1: St1, stream2: St2) -> Select<St1, St2>
where
St1: Stream,
St2: Stream<Item = St1::Item>,
{
fn round_robin(last: &mut PollNext) -> PollNext {
last.toggle()
}
assert_stream::<St1::Item, _>(Select {
inner: select_with_strategy(stream1, stream2, round_robin),
})
}
impl<St1, St2> Select<St1, St2> {
pub fn get_ref(&self) -> (&St1, &St2) {
self.inner.get_ref()
}
pub fn get_mut(&mut self) -> (&mut St1, &mut St2) {
self.inner.get_mut()
}
pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut St1>, Pin<&mut St2>) {
let this = self.project();
this.inner.get_pin_mut()
}
pub fn into_inner(self) -> (St1, St2) {
self.inner.into_inner()
}
}
impl<St1, St2> FusedStream for Select<St1, St2>
where
St1: Stream,
St2: Stream<Item = St1::Item>,
{
fn is_terminated(&self) -> bool {
self.inner.is_terminated()
}
}
impl<St1, St2> Stream for Select<St1, St2>
where
St1: Stream,
St2: Stream<Item = St1::Item>,
{
type Item = St1::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St1::Item>> {
let this = self.project();
this.inner.poll_next(cx)
}
}