Skip to main content

dope_fiber/abi/
join.rs

1use core::{marker, mem, pin, task};
2
3use crate::{abi, context, task::storage::slots};
4
5#[pin_project::pin_project]
6#[must_use = "a fiber does nothing unless it is driven"]
7pub struct Join<'d, L, R>
8where
9    L: abi::Fiber<'d>,
10    R: abi::Fiber<'d>,
11{
12    #[pin]
13    left: slots::Task<L, L::Output>,
14    #[pin]
15    right: slots::Task<R, R::Output>,
16    next: abi::Side,
17    driver: marker::PhantomData<fn(&'d ()) -> &'d ()>,
18}
19
20impl<'d, L, R> Join<'d, L, R>
21where
22    L: abi::Fiber<'d>,
23    R: abi::Fiber<'d>,
24{
25    pub fn new(left: L, right: R) -> Self {
26        Self {
27            left: slots::Task::new(left),
28            right: slots::Task::new(right),
29            next: abi::Side::Left,
30            driver: marker::PhantomData,
31        }
32    }
33    fn poll_child<'child, F>(
34        child: pin::Pin<&'child mut slots::Task<F, F::Output>>,
35        mut context: pin::Pin<&mut context::Context<'_, 'd>>,
36        first: &mut bool,
37    ) -> Option<slots::raw::Done<'child, F, F::Output>>
38    where
39        F: abi::Fiber<'d>,
40    {
41        child
42            .poll(|fiber| {
43                // SAFETY: poll admits only the live slot and retains its completed
44                // output proof before another child can run.
45                if mem::replace(first, false) {
46                    // SAFETY: this is the join's first live child descent for its
47                    // admitted PollCall; the slot owns the terminal transition.
48                    unsafe { context::raw::Live::new(fiber).delegate(context.as_mut()) }
49                } else {
50                    // SAFETY: the live slot owns completion and pays a new credit.
51                    unsafe { context::raw::Live::new(fiber).poll(context.as_mut()) }
52                }
53            })
54            .map(|(done, _)| done)
55    }
56}
57
58impl<'d, L, R> abi::Fiber<'d> for Join<'d, L, R>
59where
60    L: abi::Fiber<'d>,
61    R: abi::Fiber<'d>,
62{
63    type Output = (L::Output, R::Output);
64
65    fn poll(call: context::PollCall<'_, '_, 'd, Self>) -> task::Poll<Self::Output> {
66        let (this, mut cx) = call.into_parts();
67        let this = this.project();
68        let side = *this.next;
69        *this.next = side.other();
70        let mut first = true;
71        let (left, right) = match side {
72            abi::Side::Left => (
73                Self::poll_child(this.left, cx.as_mut(), &mut first),
74                Self::poll_child(this.right, cx, &mut first),
75            ),
76            abi::Side::Right => {
77                let right = Self::poll_child(this.right, cx.as_mut(), &mut first);
78                (Self::poll_child(this.left, cx, &mut first), right)
79            }
80        };
81        let (Some(left), Some(right)) = (left, right) else {
82            return task::Poll::Pending;
83        };
84        task::Poll::Ready((left.take(), right.take()))
85    }
86}