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 if mem::replace(first, false) {
45 unsafe { context::raw::Live::new(fiber).delegate(context.as_mut()) }
48 } else {
49 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}