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        extract: impl FnOnce(
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        let output = extract(Some(&mut slot));
127        if output.is_pending() {
128            signal.register(wake.completion());
129        }
130        output
131    }
132
133    fn slot<'a>(
134        &self,
135        token: &'a mut region::Token<region::Identity<'region>>,
136    ) -> Option<(
137        reply::Slot<'a, 's, 'arena, 'credits, 'region, D, E>,
138        &'a mut state::Signal<'region>,
139    )>
140    where
141        's: 'a,
142    {
143        let reg = self.reg.as_ref()?;
144        reg.lane.arena.slot(token, reg.key)
145    }
146
147    fn detach(&mut self) {
148        self.reg = None;
149    }
150}
151
152pub struct Response<
153    's,
154    'arena,
155    'credits,
156    'region,
157    D,
158    X: Extract<'s, 'arena, 'credits, 'region, D, E>,
159    E = D,
160    K = (),
161> {
162    handle: Handle<'s, 'arena, 'credits, 'region, D, E, K>,
163    extract: X::Context,
164}
165
166impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
167    Response<'s, 'arena, 'credits, 'region, D, X, E, K>
168{
169    /// Consumes typed group input only after every response, slot, and credit is admitted.
170    /// Any admission failure leaves every response and lane order unchanged.
171    pub fn try_group_by<I, Y, R, const N: usize>(
172        token: &mut region::Token<region::Identity<'region>>,
173        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
174        commands: &mut [Self; N],
175        terminal: &mut Response<'s, 'arena, 'credits, 'region, D, Y, E, K>,
176        input: I,
177        make: impl FnOnce(I) -> ([K; N], K, R),
178    ) -> Result<(order::Depth, R), (attach::Error, I)>
179    where
180        Y: Extract<'s, 'arena, 'credits, 'region, D, E>,
181    {
182        const { assert!(N > 0 && N < usize::MAX) };
183        for command in commands.iter() {
184            if let Err(cause) = command.handle.check(lane) {
185                return Err((cause, input));
186            }
187        }
188        if let Err(cause) = terminal.handle.check(lane) {
189            return Err((cause, input));
190        }
191        let (keys, key, output) =
192            lane.group_with(token, X::ACCEPTANCE, Y::ACCEPTANCE, input, make)?;
193        for (command, key) in commands.iter_mut().zip(keys) {
194            command.handle.reg = Some(Registration { lane, key });
195        }
196        terminal.handle.reg = Some(Registration { lane, key });
197        Ok((lane.depth(token), output))
198    }
199
200    pub const fn new() -> Self
201    where
202        X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
203    {
204        Self::with(())
205    }
206
207    pub const fn with(extract: X::Context) -> Self {
208        Self {
209            handle: Handle::new(),
210            extract,
211        }
212    }
213
214    /// Extracts an already available result once without installing a wake target.
215    /// A completed output detaches the same response ownership as `poll_reply`.
216    pub fn try_take(
217        &mut self,
218        token: &mut region::Token<region::Identity<'region>>,
219    ) -> Option<X::Output> {
220        let output = {
221            let (mut slot, _) = self.handle.slot(token)?;
222            X::extract(&self.extract, &mut slot)
223        };
224        if output.is_some() {
225            self.handle.detach();
226        }
227        output
228    }
229
230    /// Polls the attached response within the caller's admitted operation.
231    /// Extracting an output detaches the response before returning it; later
232    /// polls remain pending until an explicit attachment starts another reply.
233    pub fn poll_reply(
234        self: pin::Pin<&mut Self>,
235        context: pin::Pin<&mut context::Context<'_, 'region>>,
236    ) -> task::Poll<X::Output> {
237        let me = self.get_mut();
238        let output = me.handle.poll(context, |slot| {
239            slot.and_then(|slot| X::extract(&me.extract, slot))
240                .map_or(task::Poll::Pending, task::Poll::Ready)
241        });
242        if output.is_ready() {
243            me.handle.detach();
244        }
245        output
246    }
247}
248
249impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
250    Default for Response<'s, 'arena, 'credits, 'region, D, X, E, K>
251where
252    X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
253{
254    fn default() -> Self {
255        Self::new()
256    }
257}
258
259impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
260    attach::Attach<'s, 'arena, 'credits, 'region, D, E, K>
261    for Response<'s, 'arena, 'credits, 'region, D, X, E, K>
262{
263    fn try_ordered(
264        &mut self,
265        token: &mut region::Token<region::Identity<'region>>,
266        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
267        tag: K,
268    ) -> Result<order::Depth, attach::Error> {
269        self.handle.try_ordered(token, lane, X::ACCEPTANCE, tag)
270    }
271
272    fn try_tagged(
273        &mut self,
274        token: &mut region::Token<region::Identity<'region>>,
275        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
276        tag: K,
277    ) -> Result<(), attach::Error> {
278        self.handle.try_attach(token, lane, X::ACCEPTANCE, tag)
279    }
280
281    fn detach(&mut self) {
282        self.handle.detach();
283    }
284}
285
286impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
287    abi::Fiber<'region> for Response<'s, 'arena, 'credits, 'region, D, X, E, K>
288{
289    type Output = X::Output;
290
291    fn poll(call: context::PollCall<'_, '_, 'region, Self>) -> task::Poll<X::Output> {
292        let (this, context) = call.into_parts();
293        this.poll_reply(context)
294    }
295}
296
297pub struct Stream<
298    's,
299    'arena,
300    'credits,
301    'region,
302    D,
303    X: Extract<'s, 'arena, 'credits, 'region, D, E>,
304    E = D,
305    K = (),
306> {
307    handle: Handle<'s, 'arena, 'credits, 'region, D, E, K>,
308    extract: X::Context,
309}
310
311impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
312    Stream<'s, 'arena, 'credits, 'region, D, X, E, K>
313{
314    pub const fn new() -> Self
315    where
316        X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
317    {
318        Self::with(())
319    }
320
321    pub const fn with(extract: X::Context) -> Self {
322        Self {
323            handle: Handle::new(),
324            extract,
325        }
326    }
327
328    pub fn poll_next(
329        mut self: pin::Pin<&mut Self>,
330        context: pin::Pin<&mut context::Context<'_, 'region>>,
331    ) -> task::Poll<Option<X::Output>> {
332        let me = self.as_mut().get_mut();
333        let output = me.handle.poll(context, |slot| match slot {
334            Some(slot) => match X::extract(&me.extract, slot) {
335                Some(output) => task::Poll::Ready(Some(output)),
336                None if slot.state.is_finished() => task::Poll::Ready(None),
337                None => task::Poll::Pending,
338            },
339            None => task::Poll::Ready(None),
340        });
341        if matches!(output, task::Poll::Ready(None)) {
342            me.handle.detach();
343        }
344        output
345    }
346}
347
348impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
349    Default for Stream<'s, 'arena, 'credits, 'region, D, X, E, K>
350where
351    X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
352{
353    fn default() -> Self {
354        Self::new()
355    }
356}
357
358impl<'s, 'arena, 'credits, 'region, D, X: Extract<'s, 'arena, 'credits, 'region, D, E>, E, K>
359    attach::Attach<'s, 'arena, 'credits, 'region, D, E, K>
360    for Stream<'s, 'arena, 'credits, 'region, D, X, E, K>
361{
362    fn try_ordered(
363        &mut self,
364        token: &mut region::Token<region::Identity<'region>>,
365        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
366        tag: K,
367    ) -> Result<order::Depth, attach::Error> {
368        self.handle.try_ordered(token, lane, X::ACCEPTANCE, tag)
369    }
370
371    fn try_tagged(
372        &mut self,
373        token: &mut region::Token<region::Identity<'region>>,
374        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
375        tag: K,
376    ) -> Result<(), attach::Error> {
377        self.handle.try_attach(token, lane, X::ACCEPTANCE, tag)
378    }
379
380    fn detach(&mut self) {
381        self.handle.detach();
382    }
383}
dope_fiber/reply/recv/mod.rs

dope_fiber/reply/recv/
mod.rs
