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::raw::Entry<L, L::Output>,
14    #[pin]
15    right: slots::raw::Entry<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::raw::Entry::live(left),
28            right: slots::raw::Entry::live(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::raw::Entry<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        match child.poll(|fiber| {
42            // SAFETY: poll admits only the live slot and retains its completed
43            // output proof before another child can run.
44            if mem::replace(first, false) {
45                // SAFETY: this is the join's first live child descent for its
46                // admitted PollCall; the slot owns the terminal transition.
47                unsafe { context::raw::Live::new(fiber).delegate(context.as_mut()) }
48            } else {
49                // SAFETY: the live slot owns completion and pays a new credit.
50                unsafe { context::raw::Live::new(fiber).poll(context.as_mut()) }
51            }
52        }) {
53            Some(task::Poll::Ready((done, _))) => Some(done),
54            Some(task::Poll::Pending) | None => None,
55        }
56    }
57}
58
59impl<'d, L, R> abi::Fiber<'d> for Join<'d, L, R>
60where
61    L: abi::Fiber<'d>,
62    R: abi::Fiber<'d>,
63{
64    type Output = (L::Output, R::Output);
65
66    fn poll(call: context::PollCall<'_, '_, 'd, Self>) -> task::Poll<Self::Output> {
67        let (this, mut cx) = call.into_parts();
68        let this = this.project();
69        let side = *this.next;
70        *this.next = side.other();
71        let mut first = true;
72        let (left, right) = match side {
73            abi::Side::Left => (
74                Self::poll_child(this.left, cx.as_mut(), &mut first),
75                Self::poll_child(this.right, cx, &mut first),
76            ),
77            abi::Side::Right => {
78                let right = Self::poll_child(this.right, cx.as_mut(), &mut first);
79                (Self::poll_child(this.left, cx, &mut first), right)
80            }
81        };
82        let (Some(left), Some(right)) = (left, right) else {
83            return task::Poll::Pending;
84        };
85        task::Poll::Ready((left.take(), right.take()))
86    }
87}