Skip to content

Commit

Permalink
Use types to represent a select strategy instead of a boolean
Browse files Browse the repository at this point in the history
  • Loading branch information
Xaeroxe committed Aug 17, 2023
1 parent bc57a82 commit f73ee77
Show file tree
Hide file tree
Showing 5 changed files with 67 additions and 32 deletions.
3 changes: 3 additions & 0 deletions futures-util/src/future/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,9 @@ mod join_all;
#[cfg(feature = "alloc")]
pub use self::join_all::{join_all, JoinAll};

mod select_strategy;
pub use select_strategy::{Biased, Fair, IsBiased};

mod select;
pub use self::select::{select, select_biased, Select};

Expand Down
26 changes: 14 additions & 12 deletions futures-util/src/future/select.rs
Original file line number Diff line number Diff line change
@@ -1,18 +1,19 @@
use super::assert_future;
use super::{assert_future, Biased, Fair, IsBiased};
use crate::future::{Either, FutureExt};
use core::marker::PhantomData;
use core::pin::Pin;
use futures_core::future::{FusedFuture, Future};
use futures_core::task::{Context, Poll};

/// Future for the [`select()`] function.
#[must_use = "futures do nothing unless you `.await` or poll them"]
#[derive(Debug)]
pub struct Select<A, B> {
pub struct Select<A, B, BIASED = Fair> {
inner: Option<(A, B)>,
_biased: bool,
_phantom: PhantomData<BIASED>,
}

impl<A: Unpin, B: Unpin> Unpin for Select<A, B> {}
impl<A: Unpin, B: Unpin, BIASED> Unpin for Select<A, B, BIASED> {}

/// Waits for either one of two differently-typed futures to complete.
///
Expand Down Expand Up @@ -86,14 +87,14 @@ impl<A: Unpin, B: Unpin> Unpin for Select<A, B> {}
/// })
/// }
/// ```
pub fn select<A, B>(future1: A, future2: B) -> Select<A, B>
pub fn select<A, B>(future1: A, future2: B) -> Select<A, B, Fair>
where
A: Future + Unpin,
B: Future + Unpin,
{
assert_future::<Either<(A::Output, B), (B::Output, A)>, _>(Select {
inner: Some((future1, future2)),
_biased: false,
_phantom: PhantomData,
})
}

Expand Down Expand Up @@ -167,21 +168,22 @@ where
/// })
/// }
/// ```
pub fn select_biased<A, B>(future1: A, future2: B) -> Select<A, B>
pub fn select_biased<A, B>(future1: A, future2: B) -> Select<A, B, Biased>
where
A: Future + Unpin,
B: Future + Unpin,
{
assert_future::<Either<(A::Output, B), (B::Output, A)>, _>(Select {
inner: Some((future1, future2)),
_biased: true,
_phantom: PhantomData,
})
}

impl<A, B> Future for Select<A, B>
impl<A, B, BIASED> Future for Select<A, B, BIASED>
where
A: Future + Unpin,
B: Future + Unpin,
BIASED: IsBiased,
{
type Output = Either<(A::Output, B), (B::Output, A)>;

Expand All @@ -195,7 +197,6 @@ where
Some(value) => value,
}
}
let _biased = self._biased;

let (a, b) = self.inner.as_mut().expect("cannot poll Select twice");

Expand All @@ -208,7 +209,7 @@ where
}

#[cfg(feature = "std")]
if _biased || crate::gen_index(2) == 0 {
if BIASED::IS_BIASED || crate::gen_index(2) == 0 {
poll_wrap!(a, unwrap_option(self.inner.take()).1, Either::Left);
poll_wrap!(b, unwrap_option(self.inner.take()).0, Either::Right);
} else {
Expand All @@ -225,10 +226,11 @@ where
}
}

impl<A, B> FusedFuture for Select<A, B>
impl<A, B, BIASED> FusedFuture for Select<A, B, BIASED>
where
A: Future + Unpin,
B: Future + Unpin,
BIASED: IsBiased,
{
fn is_terminated(&self) -> bool {
self.inner.is_none()
Expand Down
22 changes: 12 additions & 10 deletions futures-util/src/future/select_ok.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
use super::assert_future;
use super::{Biased, Fair, IsBiased};
use crate::future::TryFutureExt;
use alloc::vec::Vec;
use core::iter::FromIterator;
use core::marker::PhantomData;
use core::mem;
use core::pin::Pin;
use futures_core::future::{Future, TryFuture};
Expand All @@ -10,12 +12,12 @@ use futures_core::task::{Context, Poll};
/// Future for the [`select_ok`] function.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct SelectOk<Fut> {
pub struct SelectOk<Fut, BIASED = Fair> {
inner: Vec<Fut>,
_biased: bool,
_phantom: PhantomData<BIASED>,
}

impl<Fut: Unpin> Unpin for SelectOk<Fut> {}
impl<Fut: Unpin, BIASED> Unpin for SelectOk<Fut, BIASED> {}

/// Creates a new future which will select the first successful future over a list of futures.
///
Expand All @@ -41,12 +43,12 @@ impl<Fut: Unpin> Unpin for SelectOk<Fut> {}
/// # Panics
///
/// This function will panic if the iterator specified contains no items.
pub fn select_ok<I>(iter: I) -> SelectOk<I::Item>
pub fn select_ok<I>(iter: I) -> SelectOk<I::Item, Fair>
where
I: IntoIterator,
I::Item: TryFuture + Unpin,
{
let ret = SelectOk { inner: iter.into_iter().collect(), _biased: false };
let ret = SelectOk { inner: iter.into_iter().collect(), _phantom: PhantomData };
assert!(!ret.inner.is_empty(), "iterator provided to select_ok was empty");
assert_future::<
Result<(<I::Item as TryFuture>::Ok, Vec<I::Item>), <I::Item as TryFuture>::Error>,
Expand All @@ -70,27 +72,27 @@ where
/// # Panics
///
/// This function will panic if the iterator specified contains no items.
pub fn select_ok_biased<I>(iter: I) -> SelectOk<I::Item>
pub fn select_ok_biased<I>(iter: I) -> SelectOk<I::Item, Biased>
where
I: IntoIterator,
I::Item: TryFuture + Unpin,
{
let ret = SelectOk { inner: iter.into_iter().collect(), _biased: true };
let ret = SelectOk { inner: iter.into_iter().collect(), _phantom: PhantomData };
assert!(!ret.inner.is_empty(), "iterator provided to select_ok was empty");
assert_future::<
Result<(<I::Item as TryFuture>::Ok, Vec<I::Item>), <I::Item as TryFuture>::Error>,
_,
>(ret)
}

impl<Fut: TryFuture + Unpin> Future for SelectOk<Fut> {
impl<Fut: TryFuture + Unpin, BIASED: IsBiased> Future for SelectOk<Fut, BIASED> {
type Output = Result<(Fut::Ok, Vec<Fut>), Fut::Error>;

fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Self { inner, _biased } = &mut *self;
let Self { inner, _phantom } = &mut *self;
#[cfg(feature = "std")]
{
if !*_biased {
if !BIASED::IS_BIASED {
crate::shuffle(inner);
}
}
Expand Down
25 changes: 25 additions & 0 deletions futures-util/src/future/select_strategy.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
/// When used with a select future, this will make the future biased.
/// When multiple futures are ready, the winner will be the first one
/// specified.
#[derive(Debug)]
pub struct Biased;

/// When used with a select future, this will make the future fair.
/// When multiple futures are ready, the winner will be pseudo-randomly
/// selected. This is the default behavior.
#[derive(Debug)]
pub struct Fair;

/// Reports whether the type is an instance of [`Biased`] or not.
pub trait IsBiased {
/// Contains the answer to our question: is this biased?
const IS_BIASED: bool;
}

impl IsBiased for Biased {
const IS_BIASED: bool = true;
}

impl IsBiased for Fair {
const IS_BIASED: bool = false;
}
23 changes: 13 additions & 10 deletions futures-util/src/future/try_select.rs
Original file line number Diff line number Diff line change
@@ -1,22 +1,24 @@
use super::{Biased, Fair, IsBiased};
use crate::future::{Either, TryFutureExt};
use core::marker::PhantomData;
use core::pin::Pin;
use futures_core::future::{Future, TryFuture};
use futures_core::task::{Context, Poll};

/// Future for the [`try_select()`] function.
#[must_use = "futures do nothing unless you `.await` or poll them"]
#[derive(Debug)]
pub struct TrySelect<A, B> {
pub struct TrySelect<A, B, BIASED = Fair> {
inner: Option<(A, B)>,
_biased: bool,
_phantom: PhantomData<BIASED>,
}

impl<A: Unpin, B: Unpin> Unpin for TrySelect<A, B> {}
impl<A: Unpin, B: Unpin, BIASED> Unpin for TrySelect<A, B, BIASED> {}

type EitherOk<A, B> = Either<(<A as TryFuture>::Ok, B), (<B as TryFuture>::Ok, A)>;
type EitherErr<A, B> = Either<(<A as TryFuture>::Error, B), (<B as TryFuture>::Error, A)>;

/// Waits for either one of two differently-typed futures to complete.
/// Waits for either onedsfasdf of two differently-typed futures to complete.
///
/// This function will return a new future which awaits for either one of both
/// futures to complete. The returned future will finish with both the value
Expand Down Expand Up @@ -55,14 +57,14 @@ type EitherErr<A, B> = Either<(<A as TryFuture>::Error, B), (<B as TryFuture>::E
/// })
/// }
/// ```
pub fn try_select<A, B>(future1: A, future2: B) -> TrySelect<A, B>
pub fn try_select<A, B>(future1: A, future2: B) -> TrySelect<A, B, Fair>
where
A: TryFuture + Unpin,
B: TryFuture + Unpin,
{
super::assert_future::<Result<EitherOk<A, B>, EitherErr<A, B>>, _>(TrySelect {
inner: Some((future1, future2)),
_biased: false,
_phantom: PhantomData,
})
}

Expand Down Expand Up @@ -103,21 +105,22 @@ where
/// })
/// }
/// ```
pub fn try_select_biased<A, B>(future1: A, future2: B) -> TrySelect<A, B>
pub fn try_select_biased<A, B>(future1: A, future2: B) -> TrySelect<A, B, Biased>
where
A: TryFuture + Unpin,
B: TryFuture + Unpin,
{
super::assert_future::<Result<EitherOk<A, B>, EitherErr<A, B>>, _>(TrySelect {
inner: Some((future1, future2)),
_biased: true,
_phantom: PhantomData,
})
}

impl<A: Unpin, B: Unpin> Future for TrySelect<A, B>
impl<A: Unpin, B: Unpin, BIASED> Future for TrySelect<A, B, BIASED>
where
A: TryFuture,
B: TryFuture,
BIASED: IsBiased,
{
type Output = Result<EitherOk<A, B>, EitherErr<A, B>>;

Expand All @@ -142,7 +145,7 @@ where

#[cfg(feature = "std")]
{
if self._biased || crate::gen_index(2) == 0 {
if BIASED::IS_BIASED || crate::gen_index(2) == 0 {
poll_wrap!(a, b, Either::Left, Either::Right)
} else {
poll_wrap!(b, a, Either::Right, Either::Left)
Expand Down

0 comments on commit f73ee77

Please sign in to comment.