dope_fiber/reply/recv/
mod.rs

1//! Fiber adapters for reply attachment and extraction.
2
3pub mod attach;
4pub mod outcome;
5
6use std::{mem, pin, task};
7
8use o3::base::cell::region;
9
10use crate::{
11    abi, context,
12    reply::{
13        self, Lane,
14        arena::{retire::Retired, state},
15        order,
16    },
17};
18
19/// Declares how an extractor consumes data items from its reply slot.
20#[derive(Clone, Copy, Eq, PartialEq)]
21#[repr(u8)]
22pub enum Acceptance {
23    /// Items can be extracted while the reply is still arriving.
24    Payload,
25    /// Data items are not consumed.
26    Unit,
27    /// The complete reply must fit without waiting for shared credits.
28    Collected,
29}
30
31const _: () = assert!(mem::size_of::<Acceptance>() == mem::size_of::<bool>());
32
33/// Extracts one logical value from a reply slot.
34pub trait Extract<'s, 'arena, 'credits, 'region, D, E = D> {
35    type Context: Unpin;
36    type Output;
37    const ACCEPTANCE: Acceptance = Acceptance::Payload;
38
39    fn extract(
40        context: &Self::Context,
41        slot: &mut reply::Slot<'_, 's, 'arena, 'credits, 'region, D, E>,
42    ) -> Option<Self::Output>;
43}
44
45struct Registration<'s, 'arena, 'credits, 'region, D, E, K> {
46    lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
47    key: state::EntryKey,
48}
49
50impl<D, E, K> Drop for Registration<'_, '_, '_, '_, D, E, K> {
51    fn drop(&mut self) {
52        self.lane.defer_retire(Retired::new(self.key));
53    }
54}
55
56struct Handle<'s, 'arena, 'credits, 'region, D, E, K> {
57    reg: Option<Registration<'s, 'arena, 'credits, 'region, D, E, K>>,
58}
59
60impl<'s, 'arena, 'credits, 'region, D, E, K> Handle<'s, 'arena, 'credits, 'region, D, E, K> {
61    const fn new() -> Self {
62        Self { reg: None }
63    }
64
65    fn register(
66        &self,
67        token: &mut region::Token<region::Identity<'region>>,
68        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
69        accept: Acceptance,
70        tag: K,
71    ) -> Result<Registration<'s, 'arena, 'credits, 'region, D, E, K>, attach::Error> {
72        self.check(lane)?;
73        let key = lane.register(token, accept, tag)?;
74        Ok(Registration { lane, key })
75    }
76
77    fn check(
78        &self,
79        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
80    ) -> Result<(), attach::Error> {
81        if let Some(current) = self.reg.as_ref() {
82            return Err(attach::Error::Attached {
83                current: current.lane.index,
84                requested: lane.index,
85            });
86        }
87        Ok(())
88    }
89
90    fn try_ordered(
91        &mut self,
92        token: &mut region::Token<region::Identity<'region>>,
93        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
94        accept: Acceptance,
95        tag: K,
96    ) -> Result<order::Depth, attach::Error> {
97        let reg = self.register(token, lane, accept, tag)?;
98        let depth = reg.lane.depth(token);
99        self.reg = Some(reg);
100        Ok(depth)
101    }
102
103    fn try_attach(
104        &mut self,
105        token: &mut region::Token<region::Identity<'region>>,
106        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
107        accept: Acceptance,
108        tag: K,
109    ) -> Result<(), attach::Error> {
110        let reg = self.register(token, lane, accept, tag)?;
111        self.reg = Some(reg);
112        Ok(())
113    }
114
115    fn poll<O>(
116        &self,
117        mut context: pin::Pin<&mut context::Context<'_, 'region>>,
118        mut extract: impl FnMut(
119            Option<&mut reply::Slot<'_, 's, 'arena, 'credits, 'region, D, E>>,
120        ) -> task::Poll<O>,
121    ) -> task::Poll<O> {
122        let wake = context.as_ref().waker();
123        let Some((mut slot, signal)) = self.slot(context.as_mut().region_token()) else {
124            return extract(None);
125        };
126        loop {
127            if let task::Poll::Ready(output) = extract(Some(&mut slot)) {
128                return task::Poll::Ready(output);
129            }
130            if signal.poll(wake.completion()).is_pending() {
131                return task::Poll::Pending;
132            }
133        }
134    }
135
136    fn slot<'a>(
137        &self,
138        token: &'a mut region::Token<region::Identity<'region>>,
139    ) -> Option<(
140        reply::Slot<'a, 's, 'arena, 'credits, 'region, D, E>,
141        &'a mut state::Signal<'region>,
142    )>
143    where
144        's: 'a,
145    {
146        let reg = self.reg.as_ref()?;
147        reg.lane.arena.slot(token, reg.key)
148    }
149
150    fn detach(&mut self) {
151        self.reg = None;
152    }
153}
154
155pub struct Response<
156    's,
157    'arena,
158    'credits,
159    'region,
160    D,
161    X: Extract<'s, 'arena, 'credits, 'region, D, E>,
162    E = D,
163    K = (),
164> {
165    handle: Handle<'s, 'arena, 'credits, 'region, D, E, K>,
166    extract: X::Context,
167}
168
169impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
170    Response<'s, 'arena, 'credits, 'region, D, X, E, K>
171{
172    /// Consumes typed group input only after every response, slot, and credit is admitted.
173    /// Any admission failure leaves every response and lane order unchanged.
174    pub fn try_group_by<I, Y, R, const N: usize>(
175        token: &mut region::Token<region::Identity<'region>>,
176        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
177        commands: &mut [Self; N],
178        terminal: &mut Response<'s, 'arena, 'credits, 'region, D, Y, E, K>,
179        input: I,
180        make: impl FnOnce(I) -> ([K; N], K, R),
181    ) -> Result<(order::Depth, R), (attach::Error, I)>
182    where
183        Y: Extract<'s, 'arena, 'credits, 'region, D, E>,
184    {
185        const { assert!(N > 0 && N < usize::MAX) };
186        for command in commands.iter() {
187            if let Err(cause) = command.handle.check(lane) {
188                return Err((cause, input));
189            }
190        }
191        if let Err(cause) = terminal.handle.check(lane) {
192            return Err((cause, input));
193        }
194        let (keys, key, output) =
195            lane.group_with(token, X::ACCEPTANCE, Y::ACCEPTANCE, input, make)?;
196        for (command, key) in commands.iter_mut().zip(keys) {
197            command.handle.reg = Some(Registration { lane, key });
198        }
199        terminal.handle.reg = Some(Registration { lane, key });
200        Ok((lane.depth(token), output))
201    }
202
203    pub const fn new() -> Self
204    where
205        X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
206    {
207        Self::with(())
208    }
209
210    pub const fn with(extract: X::Context) -> Self {
211        Self {
212            handle: Handle::new(),
213            extract,
214        }
215    }
216
217    /// Extracts an already available result once without installing a wake target.
218    /// A completed output detaches the same response ownership as `poll_reply`.
219    pub fn try_take(
220        &mut self,
221        token: &mut region::Token<region::Identity<'region>>,
222    ) -> Option<X::Output> {
223        let output = {
224            let (mut slot, _) = self.handle.slot(token)?;
225            X::extract(&self.extract, &mut slot)
226        };
227        if output.is_some() {
228            self.handle.detach();
229        }
230        output
231    }
232
233    /// Polls the attached response within the caller's admitted operation.
234    /// Extracting an output detaches the response before returning it; later
235    /// polls remain pending until an explicit attachment starts another reply.
236    pub fn poll_reply(
237        self: pin::Pin<&mut Self>,
238        context: pin::Pin<&mut context::Context<'_, 'region>>,
239    ) -> task::Poll<X::Output> {
240        let me = self.get_mut();
241        let output = me.handle.poll(context, |slot| {
242            slot.and_then(|slot| X::extract(&me.extract, slot))
243                .map_or(task::Poll::Pending, task::Poll::Ready)
244        });
245        if output.is_ready() {
246            me.handle.detach();
247        }
248        output
249    }
250}
251
252impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
253    Default for Response<'s, 'arena, 'credits, 'region, D, X, E, K>
254where
255    X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
256{
257    fn default() -> Self {
258        Self::new()
259    }
260}
261
262impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
263    attach::Attach<'s, 'arena, 'credits, 'region, D, E, K>
264    for Response<'s, 'arena, 'credits, 'region, D, X, E, K>
265{
266    fn try_ordered(
267        &mut self,
268        token: &mut region::Token<region::Identity<'region>>,
269        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
270        tag: K,
271    ) -> Result<order::Depth, attach::Error> {
272        self.handle.try_ordered(token, lane, X::ACCEPTANCE, tag)
273    }
274
275    fn try_tagged(
276        &mut self,
277        token: &mut region::Token<region::Identity<'region>>,
278        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
279        tag: K,
280    ) -> Result<(), attach::Error> {
281        self.handle.try_attach(token, lane, X::ACCEPTANCE, tag)
282    }
283
284    fn detach(&mut self) {
285        self.handle.detach();
286    }
287}
288
289impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
290    abi::Fiber<'region> for Response<'s, 'arena, 'credits, 'region, D, X, E, K>
291{
292    type Output = X::Output;
293
294    fn poll(call: context::PollCall<'_, '_, 'region, Self>) -> task::Poll<X::Output> {
295        let (this, context) = call.into_parts();
296        this.poll_reply(context)
297    }
298}
299
300pub struct Stream<
301    's,
302    'arena,
303    'credits,
304    'region,
305    D,
306    X: Extract<'s, 'arena, 'credits, 'region, D, E>,
307    E = D,
308    K = (),
309> {
310    handle: Handle<'s, 'arena, 'credits, 'region, D, E, K>,
311    extract: X::Context,
312}
313
314impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
315    Stream<'s, 'arena, 'credits, 'region, D, X, E, K>
316{
317    pub const fn new() -> Self
318    where
319        X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
320    {
321        Self::with(())
322    }
323
324    pub const fn with(extract: X::Context) -> Self {
325        Self {
326            handle: Handle::new(),
327            extract,
328        }
329    }
330
331    pub fn poll_next(
332        mut self: pin::Pin<&mut Self>,
333        context: pin::Pin<&mut context::Context<'_, 'region>>,
334    ) -> task::Poll<Option<X::Output>> {
335        let me = self.as_mut().get_mut();
336        let output = me.handle.poll(context, |slot| match slot {
337            Some(slot) => match X::extract(&me.extract, slot) {
338                Some(output) => task::Poll::Ready(Some(output)),
339                None if slot.state.is_finished() => task::Poll::Ready(None),
340                None => task::Poll::Pending,
341            },
342            None => task::Poll::Ready(None),
343        });
344        if matches!(output, task::Poll::Ready(None)) {
345            me.handle.detach();
346        }
347        output
348    }
349}
350
351impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
352    Default for Stream<'s, 'arena, 'credits, 'region, D, X, E, K>
353where
354    X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
355{
356    fn default() -> Self {
357        Self::new()
358    }
359}
360
361impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
362    attach::Attach<'s, 'arena, 'credits, 'region, D, E, K>
363    for Stream<'s, 'arena, 'credits, 'region, D, X, E, K>
364{
365    fn try_ordered(
366        &mut self,
367        token: &mut region::Token<region::Identity<'region>>,
368        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
369        tag: K,
370    ) -> Result<order::Depth, attach::Error> {
371        self.handle.try_ordered(token, lane, X::ACCEPTANCE, tag)
372    }
373
374    fn try_tagged(
375        &mut self,
376        token: &mut region::Token<region::Identity<'region>>,
377        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
378        tag: K,
379    ) -> Result<(), attach::Error> {
380        self.handle.try_attach(token, lane, X::ACCEPTANCE, tag)
381    }
382
383    fn detach(&mut self) {
384        self.handle.detach();
385    }
386}
dope_fiber/reply/recv/mod.rs

dope_fiber/reply/recv/
mod.rs
