Skip to main content

dope_fiber/abi/
race.rs

1use core::task;
2
3use crate::{abi, context};
4
5#[pin_project::pin_project]
6#[must_use = "a fiber does nothing unless it is driven"]
7/// Resolves with the first completed child, preferring `left` when both run.
8/// A budget-skipped child resumes first on the next poll.
9/// Completion ends the owner's polling authority for this race.
10#[derive(Debug)]
11pub struct Race<L, R> {
12    #[pin]
13    left: L,
14    #[pin]
15    right: R,
16    resume: abi::Side,
17}
18
19impl<L, R> Race<L, R> {
20    pub const fn new(left: L, right: R) -> Self {
21        Self {
22            left,
23            right,
24            resume: abi::Side::Left,
25        }
26    }
27}
28
29impl<'d, L, R> abi::Fiber<'d> for Race<L, R>
30where
31    L: abi::Fiber<'d>,
32    R: abi::Fiber<'d>,
33{
34    type Output = Either<L::Output, R::Output>;
35
36    fn poll(call: context::PollCall<'_, '_, 'd, Self>) -> task::Poll<Self::Output> {
37        use core::task::Poll;
38
39        let (this, mut context) = call.into_parts();
40        let this = this.project();
41        match *this.resume {
42            abi::Side::Left => {
43                // SAFETY: this live race owns both unfinished children and
44                // transfers its existing admission to the first child.
45                if let Poll::Ready(output) =
46                    unsafe { context::raw::Live::new(this.left).delegate(context.as_mut()) }
47                {
48                    return Poll::Ready(Either::Left(output));
49                }
50                *this.resume = abi::Side::Right;
51                // SAFETY: this child is unfinished; Ready ends the race's poll.
52                match unsafe { context::raw::Live::new(this.right).try_poll(context) } {
53                    context::PollResult::Polled(Poll::Ready(output)) => {
54                        Poll::Ready(Either::Right(output))
55                    }
56                    context::PollResult::Polled(Poll::Pending) => {
57                        *this.resume = abi::Side::Left;
58                        Poll::Pending
59                    }
60                    context::PollResult::Exhausted => Poll::Pending,
61                }
62            }
63            abi::Side::Right => {
64                // SAFETY: the skipped child is unfinished and receives the
65                // existing admission; Ready ends the race's poll.
66                if let Poll::Ready(output) =
67                    unsafe { context::raw::Live::new(this.right).delegate(context.as_mut()) }
68                {
69                    return Poll::Ready(Either::Right(output));
70                }
71                *this.resume = abi::Side::Left;
72                // SAFETY: this child is unfinished; Ready ends the race's poll.
73                unsafe { context::raw::Live::new(this.left).poll(context) }.map(Either::Left)
74            }
75        }
76    }
77}
78
79const _: () = {
80    assert!(core::mem::size_of::<Race<(), ()>>() == core::mem::size_of::<abi::Side>());
81    assert!(core::mem::size_of::<Option<Race<(), ()>>>() == core::mem::size_of::<Race<(), ()>>());
82};
83
84#[derive(Debug)]
85pub enum Either<L, R> {
86    Left(L),
87    Right(R),
88}