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