use core::fmt;
use core::future::Future;
use core::iter::FromIterator;
use core::mem;
use core::pin::Pin;
use core::task::{Context, Poll};
use alloc::boxed::Box;
use alloc::vec::Vec;
use super::{TryFuture, TryMaybeDone};
fn iter_pin_mut<T>(slice: Pin<&mut [T]>) -> impl Iterator<Item = Pin<&mut T>> {
unsafe { slice.get_unchecked_mut() }
.iter_mut()
.map(|t| unsafe { Pin::new_unchecked(t) })
}
enum FinalState<E = ()> {
Pending,
AllDone,
Error(E)
}
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct TryJoinAll<F>
where
F: TryFuture,
{
elems: Pin<Box<[TryMaybeDone<F>]>>,
}
impl<F> fmt::Debug for TryJoinAll<F>
where
F: TryFuture + fmt::Debug,
F::Ok: fmt::Debug,
F::Error: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TryJoinAll")
.field("elems", &self.elems)
.finish()
}
}
pub fn try_join_all<I>(i: I) -> TryJoinAll<I::Item>
where
I: IntoIterator,
I::Item: TryFuture,
{
let elems: Box<[_]> = i.into_iter().map(TryMaybeDone::Future).collect();
TryJoinAll {
elems: elems.into(),
}
}
impl<F> Future for TryJoinAll<F>
where
F: TryFuture,
{
type Output = Result<Vec<F::Ok>, F::Error>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut state = FinalState::AllDone;
for elem in iter_pin_mut(self.elems.as_mut()) {
match elem.try_poll(cx) {
Poll::Pending => state = FinalState::Pending,
Poll::Ready(Ok(())) => {},
Poll::Ready(Err(e)) => {
state = FinalState::Error(e);
break;
}
}
}
match state {
FinalState::Pending => Poll::Pending,
FinalState::AllDone => {
let mut elems = mem::replace(&mut self.elems, Box::pin([]));
let results = iter_pin_mut(elems.as_mut())
.map(|e| e.take_output().unwrap())
.collect();
Poll::Ready(Ok(results))
},
FinalState::Error(e) => {
let _ = mem::replace(&mut self.elems, Box::pin([]));
Poll::Ready(Err(e))
},
}
}
}
impl<F: TryFuture> FromIterator<F> for TryJoinAll<F> {
fn from_iter<T: IntoIterator<Item = F>>(iter: T) -> Self {
try_join_all(iter)
}
}