use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
pub fn select_two<F1, F2, T1, T2>(future1: F1, future2: F2) -> SelectTwo<F1, F2>
where
F1: Future<Output = T1> + Unpin,
F2: Future<Output = T2> + Unpin,
{
SelectTwo {
future1: Some(future1),
future2: Some(future2),
}
}
pub struct SelectTwo<F1, F2> {
future1: Option<F1>,
future2: Option<F2>,
}
impl<F1, F2, T1, T2> Future for SelectTwo<F1, F2>
where
F1: Future<Output = T1> + Unpin,
F2: Future<Output = T2> + Unpin,
{
type Output = SelectTwoOutput<T1, T2>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if let Some(future1) = &mut self.future1 {
match Pin::new(future1).poll(cx) {
Poll::Ready(value) => {
self.future1 = None;
return Poll::Ready(SelectTwoOutput::First(value));
},
Poll::Pending => {},
}
}
if let Some(future2) = &mut self.future2 {
match Pin::new(future2).poll(cx) {
Poll::Ready(value) => {
self.future2 = None;
return Poll::Ready(SelectTwoOutput::Second(value));
},
Poll::Pending => {},
}
}
Poll::Pending
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SelectTwoOutput<T1, T2> {
First(T1),
Second(T2),
}
pub fn select_multiple<F>(futures: Vec<F>) -> SelectMultiple<F>
where
F: Future + Unpin,
{
SelectMultiple { futures }
}
pub struct SelectMultiple<F> {
futures: Vec<F>,
}
impl<F> Future for SelectMultiple<F>
where
F: Future + Unpin,
{
type Output = SelectMultipleOutput<F::Output>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
for (index, future) in self.futures.iter_mut().enumerate() {
match Pin::new(future).poll(cx) {
Poll::Ready(value) => {
return Poll::Ready(SelectMultipleOutput::Completed(index, value));
},
Poll::Pending => {},
}
}
Poll::Pending
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SelectMultipleOutput<T> {
Completed(usize, T),
}
pub struct SelectMultipleBinding<T> {
pub index: usize,
pub output: Option<T>,
}
#[cfg(test)]
mod tests {
use super::*;
use std::task::{Context, Waker};
#[test]
fn test_select_two_output_first() {
let output: SelectTwoOutput<i32, i32> = SelectTwoOutput::First(42);
assert!(matches!(output, SelectTwoOutput::First(42)));
}
#[test]
fn test_select_two_output_second() {
let output: SelectTwoOutput<i32, i32> = SelectTwoOutput::Second(100);
assert!(matches!(output, SelectTwoOutput::Second(100)));
}
#[test]
fn test_select_multiple_output() {
let output = SelectMultipleOutput::Completed(0, 42);
assert!(matches!(output, SelectMultipleOutput::Completed(0, 42)));
}
#[test]
fn test_ready_future() {
struct Ready;
impl Future for Ready {
type Output = i32;
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
Poll::Ready(42)
}
}
let noop_waker = Waker::noop();
let mut context = Context::from_waker(&noop_waker);
let mut future = Box::pin(Ready);
assert!(matches!(Pin::new(&mut future).poll(&mut context), Poll::Ready(42)));
}
}