Skip to main content

conc_util/future/
select.rs

1use core::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use crate::Either;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub struct Select<A: Future, B: Future>(A, B);
10
11impl<A: Future, B: Future> Select<A, B> {
12    pub fn new(a: A, b: B) -> Self {
13        Self(a, b)
14    }
15}
16
17impl<A: Future, B: Future> Future for Select<A, B> {
18    type Output = Either<A::Output, B::Output>;
19
20    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
21        // SAFETY: doing projections, fields of pinned are pinned
22        let (a, b) = unsafe {
23            let ptr = self.get_unchecked_mut() as *mut Self;
24            (
25                Pin::new_unchecked(&mut (*ptr).0),
26                Pin::new_unchecked(&mut (*ptr).1),
27            )
28        };
29
30        if let Poll::Ready(a) = a.poll(cx) {
31            Poll::Ready(Either::A(a))
32        } else if let Poll::Ready(b) = b.poll(cx) {
33            Poll::Ready(Either::B(b))
34        } else {
35            Poll::Pending
36        }
37    }
38}