orengine_utils/array_queue.rs
1//! This module contains the [`ArrayQueue`].
2use crate::hints::{assert_hint, likely, unlikely};
3use std::mem::MaybeUninit;
4use std::ops::{Deref, DerefMut};
5use std::{mem, ptr};
6use std::ptr::{slice_from_raw_parts, slice_from_raw_parts_mut};
7
8/// `ArrayQueue` is a queue, but it uses an array on a stack and can't be resized.
9///
10/// # Example
11///
12/// ```rust
13/// use orengine_utils::ArrayQueue;
14///
15/// let mut queue = ArrayQueue::<u32, 2>::new();
16///
17/// queue.push(1).unwrap();
18/// queue.push(2).unwrap();
19///
20/// assert_eq!(queue.pop(), Some(1));
21/// assert_eq!(queue.pop(), Some(2));
22/// assert_eq!(queue.pop(), None);
23/// ```
24pub struct ArrayQueue<T, const N: usize> {
25 array: [MaybeUninit<T>; N],
26 len: usize,
27 head: usize,
28}
29
30impl<T, const N: usize> ArrayQueue<T, N> {
31 /// Creates a new `ArrayQueue`.
32 pub fn new() -> Self {
33 #[allow(
34 clippy::uninit_assumed_init,
35 reason = "We guarantee that the array is initialized, when reading from it"
36 )]
37 {
38 Self {
39 array: [const { MaybeUninit::uninit() }; N],
40 len: 0,
41 head: 0,
42 }
43 }
44 }
45
46 /// Returns the capacity of the queue.
47 pub const fn capacity(&self) -> usize {
48 N
49 }
50
51 /// Returns the number of elements in the queue.
52 pub fn len(&self) -> usize {
53 self.len
54 }
55
56 /// Returns `true` if the queue is empty.
57 pub fn is_empty(&self) -> bool {
58 self.len == 0
59 }
60
61 /// Returns an index of the underlying array for the provided index.
62 #[inline]
63 fn to_physical_idx_from_head(&self, idx: usize) -> usize {
64 let logical_index = self.head + idx;
65
66 debug_assert!(logical_index < N || (logical_index - N) < N);
67
68 if unlikely(logical_index >= N) {
69 logical_index - N
70 } else {
71 logical_index
72 }
73 }
74
75 /// Returns a pair of slices that represent the occupied region of the queue.
76 ///
77 /// # Example
78 ///
79 /// ```rust
80 /// use orengine_utils::ArrayQueue;
81 ///
82 /// let mut array_queue = ArrayQueue::<u32, 4>::new();
83 ///
84 /// array_queue.push(1).unwrap();
85 /// array_queue.push(2).unwrap();
86 ///
87 /// let should_be: (&[u32], &[u32]) = (&[1, 2], &[]);
88 ///
89 /// assert_eq!(array_queue.as_slices(), should_be);
90 ///
91 /// array_queue.push(3).unwrap();
92 /// array_queue.push(4).unwrap();
93 ///
94 /// assert_eq!(array_queue.pop(), Some(1));
95 /// assert_eq!(array_queue.pop(), Some(2));
96 ///
97 /// array_queue.push(5).unwrap();
98 ///
99 /// let should_be: (&[u32], &[u32]) = (&[3, 4], &[5]);
100 ///
101 /// assert_eq!(array_queue.as_slices(), should_be);
102 /// ```
103 pub fn as_slices(&self) -> (&[T], &[T]) {
104 let phys_head = self.to_physical_idx_from_head(0);
105 let phys_tail = self.to_physical_idx_from_head(self.len());
106
107 if phys_tail > phys_head {
108 (
109 unsafe { &*slice_from_raw_parts(self.array.as_ptr().add(phys_head).cast(), self.len) },
110 &[],
111 )
112 } else {
113 (
114 unsafe { &*slice_from_raw_parts(self.array.as_ptr().add(phys_head).cast(), N - phys_head) },
115 unsafe { &*slice_from_raw_parts(self.array.as_ptr().cast(), phys_tail) },
116 )
117 }
118 }
119
120 /// Returns a pair of mutable slices that represent the occupied region of the queue.
121 ///
122 /// # Example
123 ///
124 /// ```rust
125 /// use orengine_utils::ArrayQueue;
126 ///
127 /// let mut array_queue = ArrayQueue::<u32, 4>::new();
128 ///
129 /// array_queue.push(1).unwrap();
130 /// array_queue.push(2).unwrap();
131 ///
132 /// let should_be: (&mut [u32], &mut [u32]) = (&mut [1, 2], &mut []);
133 ///
134 /// assert_eq!(array_queue.as_mut_slices(), should_be);
135 ///
136 /// array_queue.push(3).unwrap();
137 /// array_queue.push(4).unwrap();
138 ///
139 /// assert_eq!(array_queue.pop(), Some(1));
140 /// assert_eq!(array_queue.pop(), Some(2));
141 ///
142 /// array_queue.push(5).unwrap();
143 ///
144 /// let should_be: (&mut [u32], &mut [u32]) = (&mut [3, 4], &mut [5]);
145 ///
146 /// assert_eq!(array_queue.as_mut_slices(), should_be);
147 /// ```
148 pub fn as_mut_slices(&mut self) -> (&mut [T], &mut [T]) {
149 let phys_head = self.to_physical_idx_from_head(0);
150 let phys_tail = self.to_physical_idx_from_head(self.len());
151
152 if phys_tail > phys_head {
153 (
154 unsafe { &mut *slice_from_raw_parts_mut(self.array.as_mut_ptr().add(phys_head).cast(), self.len) },
155 &mut [],
156 )
157 } else {
158 (
159 unsafe { &mut *slice_from_raw_parts_mut(self.array.as_mut_ptr().add(phys_head).cast(), N - phys_head) },
160 unsafe { &mut *slice_from_raw_parts_mut(self.array.as_mut_ptr().cast(), phys_tail) },
161 )
162 }
163 }
164
165 /// Increases the head index by the specified number and decreases the length by the same number.
166 ///
167 /// # Safety
168 ///
169 /// The caller must ensure usage of items that become available after this function.
170 ///
171 /// # Example
172 ///
173 /// ```rust
174 /// use orengine_utils::ArrayQueue;
175 ///
176 /// let mut queue = ArrayQueue::from([1, 2, 3, 4]);
177 ///
178 /// queue.pop().unwrap();
179 /// queue.push(5).unwrap();
180 ///
181 /// let slices = queue.as_mut_slices();
182 /// let should_be: (&mut [u32], &mut [u32]) = (&mut [2, 3, 4], &mut [5]);
183 /// assert_eq!(slices, should_be);
184 ///
185 /// for item in slices.0.iter_mut() {
186 /// // Do something with items
187 /// unsafe { std::ptr::drop_in_place(item); } // Ensure the items are dropped
188 /// }
189 ///
190 /// // Here head is 1 and len is 4
191 ///
192 /// let slices = queue.as_slices();
193 /// let as_previous: (&[u32], &[u32]) = (&[2, 3, 4], &[5]); // But the queue is still the same, while 3 elements were read
194 /// assert_eq!(slices, as_previous);
195 ///
196 /// unsafe { queue.inc_head_by(3); }
197 ///
198 /// // Here head is 0 (4 is wrapped around), and len is 1
199 ///
200 /// let slices = queue.as_slices();
201 /// let should_be: (&[u32], &[u32]) = (&[5], &[]);
202 /// assert_eq!(slices, should_be); // Now it is valid
203 /// ```
204 pub unsafe fn inc_head_by(&mut self, number: usize) {
205 self.head = self.to_physical_idx_from_head(number);
206 self.len -= number;
207 }
208
209 /// Decreases the length by the specified number.
210 ///
211 /// # Safety
212 ///
213 /// The caller must ensure that the length is not less than the specified number.
214 /// And the caller must ensure usage of items that become available after this function.
215 ///
216 /// # Example
217 ///
218 /// ```rust
219 /// use orengine_utils::ArrayQueue;
220 ///
221 /// let mut queue = ArrayQueue::from([1, 2, 3, 4]);
222 ///
223 /// queue.pop().unwrap();
224 /// queue.pop().unwrap();
225 /// queue.push(5).unwrap();
226 /// queue.push(6).unwrap();
227 ///
228 /// let slices = queue.as_mut_slices();
229 /// let should_be: (&mut [u32], &mut [u32]) = (&mut [3, 4], &mut [5, 6]);
230 /// assert_eq!(slices, should_be);
231 ///
232 /// for item in slices.1.iter_mut() {
233 /// // Do something with items
234 /// unsafe { std::ptr::drop_in_place(item); } // Ensure the items are dropped
235 /// }
236 ///
237 /// // Here head is 2 and len is 4
238 ///
239 /// let slices = queue.as_slices();
240 /// let as_previous: (&[u32], &[u32]) = (&[3, 4], &[5, 6]); // But the queue is still the same, while 2 elements were read
241 /// assert_eq!(slices, as_previous);
242 ///
243 /// unsafe { queue.dec_len_by(2); }
244 ///
245 /// // Here head is 2 and len is 2
246 ///
247 /// let slices = queue.as_slices();
248 /// let should_be: (&[u32], &[u32]) = (&[3, 4], &[]);
249 /// assert_eq!(slices, should_be); // Now it is valid
250 /// ```
251 pub unsafe fn dec_len_by(&mut self, number: usize) {
252 debug_assert!(self.len >= number);
253
254 self.len -= number;
255 }
256
257 /// Appends an element to the back of the queue.
258 ///
259 /// # Safety
260 ///
261 /// The caller must ensure that the queue is not full.
262 pub unsafe fn push_unchecked(&mut self, value: T) {
263 assert_hint(self.len() < N, "Tried to push to a full array stack");
264
265 let idx = self.to_physical_idx_from_head(self.len());
266
267 unsafe { ptr::write(self.array.get_unchecked_mut(idx), MaybeUninit::new(value)) };
268
269 self.len += 1;
270 }
271
272 /// Appends an element to the back of the queue or returns `Err(value)` if the queue is full.
273 pub fn push(&mut self, value: T) -> Result<(), T> {
274 if likely(self.len() < N) {
275 unsafe { self.push_unchecked(value) };
276
277 Ok(())
278 } else {
279 Err(value)
280 }
281 }
282
283 /// Pushes the provided value to the front of the queue.
284 ///
285 /// # Safety
286 ///
287 /// The caller must ensure that the queue is not full.
288 pub unsafe fn push_priority_value_unchecked(&mut self, value: T) {
289 assert_hint(self.len() < N, "Tried to push to a full array stack");
290
291 let phys_head = self.to_physical_idx_from_head(0);
292 let idx = if likely(phys_head > 0) {
293 phys_head - 1
294 } else {
295 N - 1
296 };
297
298 unsafe { ptr::write(self.array.get_unchecked_mut(idx), MaybeUninit::new(value)) };
299
300 self.head = idx;
301 self.len += 1;
302 }
303
304 /// Pushes the provided value to the front of the queue
305 /// or returns `Err(value)` if the queue is full.
306 ///
307 /// # Example
308 ///
309 /// ```rust
310 /// use orengine_utils::ArrayQueue;
311 ///
312 /// let mut queue = ArrayQueue::<u32, 3>::new();
313 ///
314 /// queue.push_priority_value(1).unwrap(); // [1, _, _]
315 /// queue.push(2).unwrap(); // [1, 2, _]
316 /// queue.push_priority_value(3).unwrap(); // [3, 1, 2]
317 ///
318 /// assert_eq!(queue.pop(), Some(3));
319 /// assert_eq!(queue.pop(), Some(1));
320 /// assert_eq!(queue.pop(), Some(2));
321 /// ```
322 pub fn push_priority_value(&mut self, value: T) -> Result<(), T> {
323 if likely(self.len() < N) {
324 unsafe { self.push_priority_value_unchecked(value) };
325
326 Ok(())
327 } else {
328 Err(value)
329 }
330 }
331
332 /// Removes the first element and returns it, or `None` if the queue is empty.
333 pub fn pop(&mut self) -> Option<T> {
334 if !self.is_empty() {
335 self.len -= 1;
336
337 let idx = self.head;
338 self.head = self.to_physical_idx_from_head(1);
339
340 assert_hint(
341 self.array.len() > idx,
342 &format!("idx: {}, len: {}", idx, self.array.len()),
343 );
344
345 Some(unsafe { self.array.get_unchecked_mut(idx).assume_init_read() })
346 } else {
347 None
348 }
349 }
350
351 /// Removes the last element and returns it, or `None` if the queue is empty.
352 ///
353 /// # Example
354 ///
355 /// ```rust
356 /// use orengine_utils::ArrayQueue;
357 ///
358 /// let mut queue = ArrayQueue::<u32, 3>::new();
359 ///
360 /// queue.push(1).unwrap(); // [1, _, _]
361 /// queue.push(2).unwrap(); // [1, 2, _]
362 /// queue.push(3).unwrap(); // [1, 2, 3]
363 ///
364 /// assert_eq!(queue.pop_less_priority_value(), Some(3));
365 /// assert_eq!(queue.pop(), Some(1));
366 /// assert_eq!(queue.pop(), Some(2));
367 /// ```
368 pub fn pop_less_priority_value(&mut self) -> Option<T> {
369 if !self.is_empty() {
370 self.len -= 1;
371
372 let idx = self.to_physical_idx_from_head(self.len());
373
374 Some(unsafe { self.array.get_unchecked_mut(idx).assume_init_read() })
375 } else {
376 None
377 }
378 }
379
380 /// Clears with calling the provided function on each element.
381 pub fn clear_with<F>(&mut self, mut f: F)
382 where
383 F: FnMut(T),
384 {
385 for i in 0..self.len {
386 let idx = self.to_physical_idx_from_head(i);
387
388 let value = unsafe { self.array.get_unchecked_mut(idx).assume_init_read() };
389
390 f(value);
391 }
392
393 self.len = 0;
394 }
395
396 /// Drops all elements in the queue and set the length to 0.
397 pub fn clear(&mut self) {
398 if mem::needs_drop::<T>() {
399 for i in 0..self.len {
400 let idx = self.to_physical_idx_from_head(i);
401
402 unsafe { ptr::drop_in_place(self.array.get_unchecked_mut(idx)) };
403 }
404 }
405
406 self.len = 0;
407 }
408
409 /// Returns a reference iterator over the queue.
410 pub fn iter(&self) -> impl ExactSizeIterator<Item = &T> {
411 struct Iter<'array_queue, T, const N: usize> {
412 queue: &'array_queue ArrayQueue<T, N>,
413 iterated: usize,
414 }
415
416 impl<'array_queue, T, const N: usize> Iterator for Iter<'array_queue, T, N> {
417 type Item = &'array_queue T;
418
419 fn next(&mut self) -> Option<Self::Item> {
420 if self.iterated < self.queue.len {
421 let idx = self.queue.to_physical_idx_from_head(self.iterated);
422
423 self.iterated += 1;
424
425 Some(unsafe { self.queue.array.get_unchecked(idx).assume_init_ref() })
426 } else {
427 None
428 }
429 }
430
431 fn size_hint(&self) -> (usize, Option<usize>) {
432 (self.queue.len - self.iterated, Some(self.queue.len - self.iterated))
433 }
434 }
435
436 impl<T, const N: usize> ExactSizeIterator for Iter<'_, T, N> {
437 fn len(&self) -> usize {
438 self.queue.len
439 }
440 }
441
442 Iter {
443 queue: self,
444 iterated: 0,
445 }
446 }
447
448 /// Returns a mutable reference iterator over the queue.
449 pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
450 struct IterMut<'array_queue, T, const N: usize> {
451 queue: &'array_queue mut ArrayQueue<T, N>,
452 iterated: usize,
453 }
454
455 impl<'array_queue, T, const N: usize> Iterator for IterMut<'array_queue, T, N> {
456 type Item = &'array_queue mut T;
457
458 fn next(&mut self) -> Option<Self::Item> {
459 if self.iterated < self.queue.len {
460 let idx = self.queue.to_physical_idx_from_head(self.iterated);
461
462 self.iterated += 1;
463
464 Some(unsafe { &mut *(self.queue.array.get_unchecked_mut(idx) as *mut _ as *mut T) })
465 } else {
466 None
467 }
468 }
469
470 fn size_hint(&self) -> (usize, Option<usize>) {
471 (self.queue.len - self.iterated, Some(self.queue.len - self.iterated))
472 }
473 }
474
475 impl<T, const N: usize> ExactSizeIterator for IterMut<'_, T, N> {
476 fn len(&self) -> usize {
477 self.queue.len
478 }
479 }
480
481 IterMut {
482 queue: self,
483 iterated: 0,
484 }
485 }
486
487 /// Refills the queue with elements provided by the function.
488 ///
489 /// # Safety
490 ///
491 /// The caller must ensure that the queue is empty before refilling.
492 pub unsafe fn refill_with(&mut self, f: impl FnOnce(&mut [MaybeUninit<T>; N]) -> usize) {
493 debug_assert!(self.is_empty(), "ArrayQueue should be empty before refilling");
494
495 let filled = f(&mut self.array);
496
497 debug_assert!(filled <= N, "Filled more than the capacity");
498
499 self.len = filled;
500 self.head = 0;
501 }
502
503 /// Returns a pointer to the underlying array.
504 fn as_slice_ptr(&self) -> *const [T; N] {
505 (&raw const self.array).cast()
506 }
507
508 /// Returns a mutable pointer to the underlying array.
509 fn as_mut_slice_ptr(&mut self) -> *mut [T; N] {
510 (&raw mut self.array).cast()
511 }
512}
513
514impl<T, const N: usize> Deref for ArrayQueue<T, N> {
515 type Target = [T];
516
517 fn deref(&self) -> &Self::Target {
518 unsafe { &*self.as_slice_ptr() }
519 }
520}
521
522impl<T, const N: usize> AsRef<[T]> for ArrayQueue<T, N> {
523 fn as_ref(&self) -> &[T] {
524 unsafe { &*self.as_slice_ptr() }
525 }
526}
527
528impl<T, const N: usize> DerefMut for ArrayQueue<T, N> {
529 fn deref_mut(&mut self) -> &mut Self::Target {
530 unsafe { &mut *self.as_mut_slice_ptr() }
531 }
532}
533
534impl<T, const N: usize> AsMut<[T]> for ArrayQueue<T, N> {
535 fn as_mut(&mut self) -> &mut [T] {
536 unsafe { &mut *self.as_mut_slice_ptr() }
537 }
538}
539
540impl<T, const N: usize> Default for ArrayQueue<T, N> {
541 fn default() -> Self {
542 Self::new()
543 }
544}
545
546impl<T, const N: usize> From<[T; N]> for ArrayQueue<T, N> {
547 fn from(array: [T; N]) -> Self {
548 Self {
549 array: unsafe { (&raw const array).cast::<[MaybeUninit<T>; N]>().read() },
550 len: N,
551 head: 0,
552 }
553 }
554}
555
556impl<T, const N: usize> Drop for ArrayQueue<T, N> {
557 fn drop(&mut self) {
558 self.clear();
559 }
560}
561
562#[cfg(test)]
563mod tests {
564 use super::*;
565
566 #[test]
567 fn test_array_queue() {
568 let mut queue = ArrayQueue::<u32, 4>::new();
569
570 unsafe {
571 queue.push_unchecked(1);
572 assert_eq!(queue.len(), 1);
573
574 queue.push_unchecked(2);
575 assert_eq!(queue.len(), 2);
576
577 queue.push(3).unwrap();
578 assert_eq!(queue.len(), 3);
579
580 assert_eq!(queue.pop(), Some(1));
581 assert_eq!(queue.len(), 2);
582
583 queue.push_unchecked(4);
584 assert_eq!(queue.len(), 3);
585
586 queue.push_unchecked(5);
587 assert_eq!(queue.len(), 4);
588
589 assert_eq!(queue.push(6), Err(6));
590
591 assert_eq!(queue.pop(), Some(2));
592 assert_eq!(queue.pop(), Some(3));
593 assert_eq!(queue.pop(), Some(4));
594 assert_eq!(queue.pop(), Some(5));
595 assert_eq!(queue.pop(), None);
596 }
597 }
598
599 #[test]
600 fn test_array_queue_iterators() {
601 let mut queue = ArrayQueue::<u32, 4>::new();
602
603 unsafe {
604 queue.push_unchecked(1);
605 queue.push_unchecked(2);
606 queue.push_unchecked(3);
607 queue.push_unchecked(4);
608 }
609
610 assert_eq!(queue.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
611 assert_eq!(queue.iter_mut().collect::<Vec<_>>(), vec![&mut 1, &mut 2, &mut 3, &mut 4]);
612 }
613
614 #[test]
615 fn test_array_queue_refill_with() {
616 let mut queue = ArrayQueue::<u32, 4>::new();
617
618 unsafe {
619 queue.refill_with(|array| {
620 array.copy_from_slice(&[
621 MaybeUninit::new(1),
622 MaybeUninit::new(2),
623 MaybeUninit::new(3),
624 MaybeUninit::new(4)
625 ]);
626
627 4
628 });
629 };
630
631 assert_eq!(queue.len(), 4);
632 assert_eq!(queue.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
633 }
634}