| 1 | use super::assert_future; |
| 2 | use crate::future::TryFutureExt; |
| 3 | use alloc::vec::Vec; |
| 4 | use core::iter::FromIterator; |
| 5 | use core::mem; |
| 6 | use core::pin::Pin; |
| 7 | use futures_core::future::{Future, TryFuture}; |
| 8 | use futures_core::task::{Context, Poll}; |
| 9 | |
| 10 | /// Future for the [`select_ok`] function. |
| 11 | #[derive (Debug)] |
| 12 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
| 13 | pub struct SelectOk<Fut> { |
| 14 | inner: Vec<Fut>, |
| 15 | } |
| 16 | |
| 17 | impl<Fut: Unpin> Unpin for SelectOk<Fut> {} |
| 18 | |
| 19 | /// Creates a new future which will select the first successful future over a list of futures. |
| 20 | /// |
| 21 | /// The returned future will wait for any future within `iter` to be ready and Ok. Unlike |
| 22 | /// `select_all`, this will only return the first successful completion, or the last |
| 23 | /// failure. This is useful in contexts where any success is desired and failures |
| 24 | /// are ignored, unless all the futures fail. |
| 25 | /// |
| 26 | /// This function is only available when the `std` or `alloc` feature of this |
| 27 | /// library is activated, and it is activated by default. |
| 28 | /// |
| 29 | /// # Panics |
| 30 | /// |
| 31 | /// This function will panic if the iterator specified contains no items. |
| 32 | pub fn select_ok<I>(iter: I) -> SelectOk<I::Item> |
| 33 | where |
| 34 | I: IntoIterator, |
| 35 | I::Item: TryFuture + Unpin, |
| 36 | { |
| 37 | let ret: SelectOk = SelectOk { inner: iter.into_iter().collect() }; |
| 38 | assert!(!ret.inner.is_empty(), "iterator provided to select_ok was empty" ); |
| 39 | assert_future::< |
| 40 | Result<(<I::Item as TryFuture>::Ok, Vec<I::Item>), <I::Item as TryFuture>::Error>, |
| 41 | _, |
| 42 | >(ret) |
| 43 | } |
| 44 | |
| 45 | impl<Fut: TryFuture + Unpin> Future for SelectOk<Fut> { |
| 46 | type Output = Result<(Fut::Ok, Vec<Fut>), Fut::Error>; |
| 47 | |
| 48 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
| 49 | // loop until we've either exhausted all errors, a success was hit, or nothing is ready |
| 50 | loop { |
| 51 | let item = |
| 52 | self.inner.iter_mut().enumerate().find_map(|(i, f)| match f.try_poll_unpin(cx) { |
| 53 | Poll::Pending => None, |
| 54 | Poll::Ready(e) => Some((i, e)), |
| 55 | }); |
| 56 | match item { |
| 57 | Some((idx, res)) => { |
| 58 | // always remove Ok or Err, if it's not the last Err continue looping |
| 59 | drop(self.inner.remove(idx)); |
| 60 | match res { |
| 61 | Ok(e) => { |
| 62 | let rest = mem::take(&mut self.inner); |
| 63 | return Poll::Ready(Ok((e, rest))); |
| 64 | } |
| 65 | Err(e) => { |
| 66 | if self.inner.is_empty() { |
| 67 | return Poll::Ready(Err(e)); |
| 68 | } |
| 69 | } |
| 70 | } |
| 71 | } |
| 72 | None => { |
| 73 | // based on the filter above, nothing is ready, return |
| 74 | return Poll::Pending; |
| 75 | } |
| 76 | } |
| 77 | } |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | impl<Fut: TryFuture + Unpin> FromIterator<Fut> for SelectOk<Fut> { |
| 82 | fn from_iter<T: IntoIterator<Item = Fut>>(iter: T) -> Self { |
| 83 | select_ok(iter) |
| 84 | } |
| 85 | } |
| 86 | |