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    /// Constructs the response only after admission; failure returns the input unchanged.
206    pub fn try_with<I, R>(
207        token: &mut region::Token<region::Identity<'region>>,
208        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
209        tag: K,
210        input: I,
211        make: impl FnOnce(I) -> (X::Context, R),
212    ) -> Result<(order::Depth, Self, R), (attach::Error, I)> {
213        let key = match lane.register(token, X::ACCEPTANCE, tag) {
214            Ok(key) => key,
215            Err(cause) => return Err((cause, input)),
216        };
217        let (extract, output) = make(input);
218        let response = Self {
219            handle: Handle {
220                reg: Some(Registration { lane, key }),
221            },
222            extract,
223        };
224        Ok((lane.depth(token), response, output))
225    }
226
227    /// Consumes typed group input only after every response, slot, and credit is admitted.
228    /// Any admission failure leaves every response and lane order unchanged.
229    pub fn try_group_by<I, Y, R, const N: usize>(
230        token: &mut region::Token<region::Identity<'region>>,
231        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
232        commands: &mut [Self; N],
233        terminal: &mut Response<'s, 'arena, 'credits, 'region, D, Y, E, K>,
234        input: I,
235        make: impl FnOnce(I) -> ([K; N], K, R),
236    ) -> Result<(order::Depth, R), (attach::Error, I)>
237    where
238        Y: Extract<'s, 'arena, 'credits, 'region, D, E>,
239    {
240        const { assert!(N > 0 && N < usize::MAX) };
241        for command in commands.iter() {
242            if let Err(cause) = command.handle.check(lane) {
243                return Err((cause, input));
244            }
245        }
246        if let Err(cause) = terminal.handle.check(lane) {
247            return Err((cause, input));
248        }
249        let (keys, key, output) =
250            lane.group_with(token, X::ACCEPTANCE, Y::ACCEPTANCE, input, make)?;
251        for (command, key) in commands.iter_mut().zip(keys) {
252            command.handle.reg = Some(Registration { lane, key });
253        }
254        terminal.handle.reg = Some(Registration { lane, key });
255        Ok((lane.depth(token), output))
256    }
257
258    pub const fn new() -> Self
259    where
260        X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
261    {
262        Self::with(())
263    }
264
265    pub const fn with(extract: X::Context) -> Self {
266        Self {
267            handle: Handle::new(),
268            extract,
269        }
270    }
271
272    /// Extracts an already available result once without installing a wake target.
273    /// A completed output detaches the same response ownership as `poll_reply`.
274    pub fn try_take(
275        &mut self,
276        token: &mut region::Token<region::Identity<'region>>,
277    ) -> Option<X::Output> {
278        let output = {
279            let (mut slot, _) = self.handle.slot(token)?;
280            X::extract(&self.extract, &mut slot)
281        };
282        if output.is_some() {
283            self.handle.detach();
284        }
285        output
286    }
287
288    /// Polls the attached response within the caller's admitted operation.
289    /// Extracting an output detaches the response before returning it; later
290    /// polls remain pending until an explicit attachment starts another reply.
291    pub fn poll_reply(
292        self: pin::Pin<&mut Self>,
293        context: pin::Pin<&mut context::Context<'_, 'region>>,
294    ) -> task::Poll<X::Output> {
295        let me = self.get_mut();
296        let output = me.handle.poll(context, |slot| {
297            slot.and_then(|slot| X::extract(&me.extract, slot))
298                .map_or(task::Poll::Pending, task::Poll::Ready)
299        });
300        if output.is_ready() {
301            me.handle.detach();
302        }
303        output
304    }
305
306    pub fn poll_next(
307        mut self: pin::Pin<&mut Self>,
308        context: pin::Pin<&mut context::Context<'_, 'region>>,
309    ) -> task::Poll<Option<X::Output>> {
310        let me = self.as_mut().get_mut();
311        let output = me.handle.poll(context, |slot| match slot {
312            Some(slot) => match X::extract(&me.extract, slot) {
313                Some(output) => task::Poll::Ready(Some(output)),
314                None if slot.state.is_finished() => task::Poll::Ready(None),
315                None => task::Poll::Pending,
316            },
317            None => task::Poll::Ready(None),
318        });
319        if matches!(output, task::Poll::Ready(None)) {
320            me.handle.detach();
321        }
322        output
323    }
324}
325
326impl<'s, 'arena, 'credits, 'region, D: Footprint, X, E, K> Default
327    for Response<'s, 'arena, 'credits, 'region, D, X, E, K>
328where
329    X: Extract<'s, 'arena, 'credits, 'region, D, E, Context = ()>,
330{
331    fn default() -> Self {
332        Self::new()
333    }
334}
335
336impl<
337    's,
338    'arena,
339    'credits,
340    'region,
341    D: Footprint,
342    X: Extract<'s, 'arena, 'credits, 'region, D, E>,
343    E,
344    K,
345> attach::Attach<'s, 'arena, 'credits, 'region, D, E, K>
346    for Response<'s, 'arena, 'credits, 'region, D, X, E, K>
347{
348    fn try_ordered(
349        &mut self,
350        token: &mut region::Token<region::Identity<'region>>,
351        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
352        tag: K,
353    ) -> Result<order::Depth, attach::Error> {
354        self.handle.try_ordered(token, lane, X::ACCEPTANCE, tag)
355    }
356
357    fn try_tagged(
358        &mut self,
359        token: &mut region::Token<region::Identity<'region>>,
360        lane: Lane<'s, 'arena, 'credits, 'region, D, E, K>,
361        tag: K,
362    ) -> Result<(), attach::Error> {
363        self.handle.try_attach(token, lane, X::ACCEPTANCE, tag)
364    }
365
366    fn detach(&mut self) {
367        self.handle.detach();
368    }
369
370    fn detach_order(&mut self, token: &region::Token<region::Identity<'region>>) -> order::State {
371        self.handle.detach_order(token)
372    }
373}
374
375impl<
376    's,
377    'arena,
378    'credits,
379    'region,
380    D: Footprint,
381    X: Extract<'s, 'arena, 'credits, 'region, D, E>,
382    E,
383    K,
384> abi::Fiber<'region> for Response<'s, 'arena, 'credits, 'region, D, X, E, K>
385{
386    type Output = X::Output;
387
388    fn poll(call: context::PollCall<'_, '_, 'region, Self>) -> task::Poll<X::Output> {
389        let (this, context) = call.into_parts();
390        this.poll_reply(context)
391    }
392}
393
394pub type Stream<'s, 'arena, 'credits, 'region, D, X, E = D, K = ()> =
395    Response<'s, 'arena, 'credits, 'region, D, X, E, K>;
dope_fiber/reply/recv/mod.rs

dope_fiber/reply/recv/
mod.rs
