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