Skip to main content

dope_core/driver/schedule/ready/wake/
queue.rs

1use std::{mem, num};
2
3use o3::depot::queue::resident::{
4    self, fifo,
5    fifo::{Index, Scope},
6};
7
8use crate::{
9    driver,
10    driver::schedule::{
11        app,
12        ready::{
13            self, wake,
14            wake::{
15                Drain, Once, Waker,
16                fanout::storage::Step,
17                resident::{drain, one},
18            },
19        },
20    },
21};
22
23/// Bounded FIFO wakes with a constructor-proved lane extent.
24pub struct Queue<'d, L = num::NonZero<usize>> {
25    driver: driver::Reference<'d>,
26    inner: fifo::Fifo<ready::Retained<'d>, L>,
27}
28
29/// Linear ownership of one resident queue slot.
30#[repr(transparent)]
31pub struct Registration<'target, 'd: 'target, L = num::NonZero<usize>> {
32    pub(super) inner: fifo::Registration<'target, ready::Retained<'d>, L>,
33}
34
35#[derive(Clone, Copy, Debug, Eq, PartialEq)]
36#[must_use = "a rejected wake registration must be handled"]
37pub enum Error {
38    Lane(resident::Bounds<num::NonZero<u32>>),
39    Active { position: wake::Position },
40    Draining { position: wake::Position },
41    Full { position: wake::Position },
42}
43
44impl From<fifo::Reject> for Error {
45    fn from(error: fifo::Reject) -> Self {
46        match error {
47            fifo::Reject::Active { index } => Self::Active {
48                position: wake::Position::new(index),
49            },
50            fifo::Reject::Draining { index } => Self::Draining {
51                position: wake::Position::new(index),
52            },
53            fifo::Reject::Full { index } => Self::Full {
54                position: wake::Position::new(index),
55            },
56        }
57    }
58}
59
60impl<'d, L> Queue<'d, L> {
61    pub fn try_new(
62        driver: driver::Reference<'d>,
63        capacity: usize,
64        lanes: L,
65    ) -> Result<Self, resident::Error>
66    where
67        L: resident::Extent,
68    {
69        Ok(Self {
70            driver,
71            inner: fifo::Fifo::try_new(capacity, lanes)?,
72        })
73    }
74}
75
76impl<'d> Queue<'d> {
77    /// Creates one queue and borrows its source under a fresh invariant scope.
78    pub fn try_with<S, R>(
79        driver: driver::Reference<'d>,
80        capacity: usize,
81        lanes: num::NonZero<usize>,
82        source: S,
83        run: impl for<'scope> FnOnce(&'scope S, Queue<'d, Scope<'scope>>) -> R,
84    ) -> Result<R, resident::Error> {
85        fifo::Fifo::try_with(capacity, lanes, source, |source, inner| {
86            run(source, Queue { driver, inner })
87        })
88    }
89
90    pub fn try_register<'target>(
91        &'target self,
92        position: usize,
93        slot: &mut Registration<'target, 'd>,
94        wake: &Waker<'d>,
95    ) -> Result<(), Error> {
96        let inner = self.inner.try_lane(position).map_err(Error::Lane)?;
97        inner
98            .try_keep(&mut slot.inner, *wake.key())
99            .map_err(Error::from)
100    }
101
102    pub fn try_wake_one(
103        &self,
104        position: usize,
105    ) -> Result<Once, resident::Bounds<num::NonZero<u32>>> {
106        let lane = self.inner.try_lane(position)?;
107        Ok(one(self.driver, lane.take_one()))
108    }
109}
110
111impl<'scope, 'd> Queue<'d, Scope<'scope>> {
112    pub fn try_index(
113        &self,
114        position: usize,
115    ) -> Result<Index<'scope>, resident::Bounds<num::NonZero<u32>>> {
116        self.inner.try_index(position)
117    }
118
119    pub fn try_register<'target>(
120        &'target self,
121        index: Index<'scope>,
122        slot: &mut Registration<'target, 'd, Scope<'scope>>,
123        wake: &Waker<'d>,
124    ) -> Result<(), fifo::Reject> {
125        self.inner
126            .lane(index)
127            .try_keep(&mut slot.inner, *wake.key())
128    }
129
130    pub fn wake_one(&self, index: Index<'scope>) -> Once {
131        one(self.driver, self.inner.lane(index).take_one())
132    }
133}
134
135impl<'d, const N: usize> Queue<'d, resident::Fixed<N>> {
136    pub fn try_register<'target, const I: usize>(
137        &'target self,
138        slot: &mut Registration<'target, 'd, resident::Fixed<N>>,
139        wake: &Waker<'d>,
140    ) -> Result<(), fifo::Reject> {
141        self.inner
142            .lane::<I>()
143            .try_keep(&mut slot.inner, *wake.key())
144    }
145
146    pub fn wake_one<const I: usize>(&self) -> Once {
147        one(self.driver, self.inner.lane::<I>().take_one())
148    }
149}
150
151impl<'d, L> Queue<'d, L> {
152    /// Registers one wake in the constructor-proved first lane.
153    pub fn try_first<'target>(
154        &'target self,
155        slot: &mut Registration<'target, 'd, L>,
156        wake: &Waker<'d>,
157    ) -> Result<(), fifo::Reject> {
158        self.inner.first().try_keep(&mut slot.inner, *wake.key())
159    }
160
161    /// Wakes one registration from lane zero.
162    pub fn wake_first(&self) -> Once {
163        one(self.driver, self.inner.first().take_one())
164    }
165
166    pub fn wake_next(&self) -> Once {
167        let Some(target) = self.inner.take_next() else {
168            return Once::Empty;
169        };
170        target.wake(self.driver);
171        Once::Woke
172    }
173
174    pub fn has_waiters(&self) -> bool {
175        self.inner.has_waiters()
176    }
177
178    pub fn begin_drain(&self) {
179        self.inner.begin_drain();
180    }
181
182    pub fn is_draining(&self) -> bool {
183        self.inner.is_draining()
184    }
185
186    pub fn drain_one(&self, work: app::Work<'_, 'd>) -> Drain {
187        drain::one(
188            self.driver,
189            work,
190            || self.inner.is_draining() && self.inner.has_waiters(),
191            || match self.inner.drain_next() {
192                resident::Drain::Done => None,
193                resident::Drain::Item(target) => Some(Step::Item(target)),
194            },
195        )
196    }
197}
198
199impl<L> Default for Registration<'_, '_, L> {
200    fn default() -> Self {
201        Self::new()
202    }
203}
204
205impl<L> Registration<'_, '_, L> {
206    pub const fn new() -> Self {
207        Self {
208            inner: fifo::Registration::new(),
209        }
210    }
211
212    pub fn is_active(&self) -> bool {
213        self.inner.is_active()
214    }
215
216    pub fn clear(&mut self) {
217        *self = Self::new();
218    }
219}
220
221const _: () = {
222    assert!(mem::size_of::<Registration<'static, 'static>>() == 2 * mem::size_of::<usize>());
223    assert!(mem::align_of::<Registration<'static, 'static>>() == mem::align_of::<usize>());
224    assert!(
225        size_of::<Registration<'static, 'static, resident::Fixed<1>>>() == 2 * size_of::<usize>()
226    );
227    assert!(size_of::<Queue<'static, resident::Fixed<1>>>() == size_of::<Queue<'static>>());
228    assert!(size_of::<Queue<'static, Scope<'static>>>() == size_of::<Queue<'static>>());
229    assert!(size_of::<Index<'static>>() == size_of::<u32>());
230    assert!(size_of::<Registration<'static, 'static, Scope<'static>>>() == 2 * size_of::<usize>());
231};