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

dope_fiber/reply/recv/
mod.rs
