90 lines
1.8 KiB
Rust
90 lines
1.8 KiB
Rust
use std::{
|
|
future::Future,
|
|
pin::Pin,
|
|
sync::{
|
|
atomic::{AtomicBool, Ordering},
|
|
Arc,
|
|
},
|
|
task::{Context, Poll, Wake, Waker},
|
|
};
|
|
|
|
pub enum Either<L, R> {
|
|
Left(L),
|
|
Right(R),
|
|
}
|
|
|
|
pub fn join<L, R>(left: L, right: R) -> impl Future<Output = Either<L::Output, R::Output>> + Unpin
|
|
where
|
|
L: Future + Unpin,
|
|
R: Future + Unpin,
|
|
{
|
|
Select {
|
|
left: SelectState {
|
|
woken: Arc::new(AtomicBool::new(true)),
|
|
future: left,
|
|
},
|
|
right: SelectState {
|
|
woken: Arc::new(AtomicBool::new(true)),
|
|
future: right,
|
|
},
|
|
}
|
|
}
|
|
|
|
struct SelectWaker {
|
|
woken: Arc<AtomicBool>,
|
|
parent: Waker,
|
|
}
|
|
|
|
struct SelectState<T> {
|
|
woken: Arc<AtomicBool>,
|
|
future: T,
|
|
}
|
|
|
|
struct Select<L, R> {
|
|
left: SelectState<L>,
|
|
right: SelectState<R>,
|
|
}
|
|
|
|
impl<T: Future> SelectState<T> {
|
|
fn check_for_poll(&self) -> bool {
|
|
self.woken.swap(false, Ordering::AcqRel)
|
|
}
|
|
}
|
|
|
|
impl Wake for SelectWaker {
|
|
fn wake(self: Arc<Self>) {
|
|
self.wake_by_ref();
|
|
}
|
|
|
|
fn wake_by_ref(self: &Arc<Self>) {
|
|
self.woken.store(true, Ordering::Release);
|
|
self.parent.wake_by_ref();
|
|
}
|
|
}
|
|
|
|
impl<L, R> Future for Select<L, R>
|
|
where
|
|
L: Future + Unpin,
|
|
R: Future + Unpin,
|
|
{
|
|
type Output = Either<L::Output, R::Output>;
|
|
|
|
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
let this = self.get_mut();
|
|
|
|
if this.left.check_for_poll() {
|
|
if let Poll::Ready(out) = Pin::new(&mut this.left.future).poll(cx) {
|
|
return Poll::Ready(Either::Left(out));
|
|
}
|
|
}
|
|
|
|
if this.right.check_for_poll() {
|
|
if let Poll::Ready(out) = Pin::new(&mut this.right.future).poll(cx) {
|
|
return Poll::Ready(Either::Right(out));
|
|
}
|
|
}
|
|
|
|
Poll::Pending
|
|
}
|
|
}
|