use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use crate::time::SimDuration;
use crate::triggers::Timer;
pub enum Either<A, B> {
First(A),
Second(B),
}
pub struct First2<'a, A, B> {
a: Pin<Box<dyn Future<Output = A> + 'a>>,
b: Pin<Box<dyn Future<Output = B> + 'a>>,
}
pub fn first2<'a, FA, FB>(a: FA, b: FB) -> First2<'a, FA::Output, FB::Output>
where
FA: Future + 'a,
FB: Future + 'a,
{
First2 { a: Box::pin(a), b: Box::pin(b) }
}
impl<A, B> Future for First2<'_, A, B> {
type Output = Either<A, B>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if let Poll::Ready(v) = self.a.as_mut().poll(cx) {
return Poll::Ready(Either::First(v));
}
if let Poll::Ready(v) = self.b.as_mut().poll(cx) {
return Poll::Ready(Either::Second(v));
}
Poll::Pending
}
}
pub struct Join2<'a, A, B> {
a: Pin<Box<dyn Future<Output = A> + 'a>>,
b: Pin<Box<dyn Future<Output = B> + 'a>>,
ra: Option<A>,
rb: Option<B>,
}
pub fn join2<'a, FA, FB>(a: FA, b: FB) -> Join2<'a, FA::Output, FB::Output>
where
FA: Future + 'a,
FB: Future + 'a,
{
Join2 { a: Box::pin(a), b: Box::pin(b), ra: None, rb: None }
}
impl<A, B> Future for Join2<'_, A, B> {
type Output = (A, B);
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<(A, B)> {
let this = unsafe { self.get_unchecked_mut() };
if this.ra.is_none() {
if let Poll::Ready(v) = this.a.as_mut().poll(cx) {
this.ra = Some(v);
}
}
if this.rb.is_none() {
if let Poll::Ready(v) = this.b.as_mut().poll(cx) {
this.rb = Some(v);
}
}
if this.ra.is_some() && this.rb.is_some() {
Poll::Ready((this.ra.take().unwrap(), this.rb.take().unwrap()))
} else {
Poll::Pending
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TimeoutError;
impl fmt::Display for TimeoutError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "operation timed out")
}
}
impl std::error::Error for TimeoutError {}
pub async fn with_timeout<'a, F>(fut: F, d: SimDuration) -> Result<F::Output, TimeoutError>
where
F: Future + 'a,
{
match first2(fut, Timer::new(d)).await {
Either::First(v) => Ok(v),
Either::Second(()) => Err(TimeoutError),
}
}
pub struct JoinAll<'a, T> {
futs: Vec<Option<Pin<Box<dyn Future<Output = T> + 'a>>>>,
out: Vec<Option<T>>,
}
pub fn join_all<'a, T>(futs: Vec<Pin<Box<dyn Future<Output = T> + 'a>>>) -> JoinAll<'a, T> {
let n = futs.len();
JoinAll { futs: futs.into_iter().map(Some).collect(), out: (0..n).map(|_| None).collect() }
}
impl<T> Future for JoinAll<'_, T> {
type Output = Vec<T>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Vec<T>> {
let this = unsafe { self.get_unchecked_mut() };
let mut all_done = true;
for (i, slot) in this.futs.iter_mut().enumerate() {
if let Some(f) = slot {
match f.as_mut().poll(cx) {
Poll::Ready(v) => {
this.out[i] = Some(v);
*slot = None; }
Poll::Pending => all_done = false,
}
}
}
if all_done {
Poll::Ready(this.out.iter_mut().map(|o| o.take().unwrap()).collect())
} else {
Poll::Pending
}
}
}
#[macro_export]
macro_rules! first {
($a:expr, $b:expr $(,)?) => {
$crate::combinators::first2($a, $b)
};
($a:expr, $b:expr, $($rest:expr),+ $(,)?) => {
$crate::combinators::first2($a, $crate::first!($b, $($rest),+))
};
}
#[macro_export]
macro_rules! join {
($a:expr, $b:expr $(,)?) => {
$crate::combinators::join2($a, $b)
};
($a:expr, $b:expr, $($rest:expr),+ $(,)?) => {
$crate::combinators::join2($a, $crate::join!($b, $($rest),+))
};
}
#[cfg(test)]
mod tests {
use super::*;
use crate::testing::block_on;
use std::cell::RefCell;
use std::rc::Rc;
#[test]
fn join2_yields_both() {
block_on(async {
let (a, b) = join2(async { 1u8 }, async { "two" }).await;
assert_eq!(a, 1);
assert_eq!(b, "two");
});
}
#[test]
fn join_all_preserves_input_order() {
block_on(async {
let futs: Vec<Pin<Box<dyn Future<Output = u8>>>> =
vec![Box::pin(async { 1 }), Box::pin(async { 2 }), Box::pin(async { 3 })];
assert_eq!(join_all(futs).await, vec![1, 2, 3]);
});
}
#[test]
fn first2_returns_the_winner() {
block_on(async {
let ev = crate::sync::Event::new();
let waiter = ev.clone();
crate::executor::spawn(async move {
ev.set();
});
match first2(async { 7u8 }, async move { waiter.wait().await }).await {
Either::First(v) => assert_eq!(v, 7),
Either::Second(()) => panic!("the ready future should have won"),
}
});
}
#[test]
fn first2_drops_the_loser() {
struct Tattle(Rc<RefCell<bool>>);
impl Drop for Tattle {
fn drop(&mut self) {
*self.0.borrow_mut() = true;
}
}
let dropped = Rc::new(RefCell::new(false));
let flag = dropped.clone();
block_on(async move {
let never = crate::sync::Event::new();
let tattle = Tattle(flag);
let loser = async move {
let _t = tattle;
never.wait().await;
};
let _ = first2(async { 1u8 }, loser).await;
});
assert!(*dropped.borrow(), "the losing future was dropped, not left running");
}
#[test]
fn joined_futures_may_borrow() {
block_on(async {
let owned = vec![1u8, 2, 3];
let borrow_a = async { owned.len() };
let borrow_b = async { owned[0] as usize };
let (a, b) = join2(borrow_a, borrow_b).await;
assert_eq!((a, b), (3, 1));
});
}
}