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, slice_from_raw_parts_mut};
8use core::{fmt, mem, ptr};
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 (
182 unsafe {
183 &mut *slice_from_raw_parts_mut(
184 self.array.as_mut_ptr().add(phys_head).cast(),
185 N - phys_head,
186 )
187 },
188 unsafe {
189 &mut *slice_from_raw_parts_mut(self.array.as_mut_ptr().cast(), phys_tail)
190 },
191 )
192 }
193 }
194
195 /// Increases the head index by the specified number and decreases the length by the same number.
196 ///
197 /// # Safety
198 ///
199 /// The caller must ensure usage of items that become available after this function.
200 ///
201 /// # Example
202 ///
203 /// ```rust
204 /// use orengine_utils::ArrayQueue;
205 ///
206 /// let mut queue = ArrayQueue::from([1, 2, 3, 4]);
207 ///
208 /// queue.pop().unwrap();
209 /// queue.push(5).unwrap();
210 ///
211 /// let slices = queue.as_mut_slices();
212 /// let should_be: (&mut [u32], &mut [u32]) = (&mut [2, 3, 4], &mut [5]);
213 /// assert_eq!(slices, should_be);
214 ///
215 /// for item in slices.0.iter_mut() {
216 /// // Do something with items
217 /// unsafe { std::ptr::drop_in_place(item); } // Ensure the items are dropped
218 /// }
219 ///
220 /// // Here head is 1 and len is 4
221 ///
222 /// let slices = queue.as_slices();
223 /// let as_previous: (&[u32], &[u32]) = (&[2, 3, 4], &[5]); // But the queue is still the same, while 3 elements were read
224 /// assert_eq!(slices, as_previous);
225 ///
226 /// unsafe { queue.inc_head_by(3); }
227 ///
228 /// // Here head is 0 (4 is wrapped around), and len is 1
229 ///
230 /// let slices = queue.as_slices();
231 /// let should_be: (&[u32], &[u32]) = (&[5], &[]);
232 /// assert_eq!(slices, should_be); // Now it is valid
233 /// ```
234 pub unsafe fn inc_head_by(&mut self, number: usize) {
235 self.head = self.to_physical_idx_from_head(number);
236 self.len -= number;
237 }
238
239 /// Decreases the length by the specified number.
240 ///
241 /// # Safety
242 ///
243 /// The caller must ensure that the length is not less than the specified number.
244 /// And the caller must ensure usage of items that become available after this function.
245 ///
246 /// # Example
247 ///
248 /// ```rust
249 /// use orengine_utils::ArrayQueue;
250 ///
251 /// let mut queue = ArrayQueue::from([1, 2, 3, 4]);
252 ///
253 /// queue.pop().unwrap();
254 /// queue.pop().unwrap();
255 /// queue.push(5).unwrap();
256 /// queue.push(6).unwrap();
257 ///
258 /// let slices = queue.as_mut_slices();
259 /// let should_be: (&mut [u32], &mut [u32]) = (&mut [3, 4], &mut [5, 6]);
260 /// assert_eq!(slices, should_be);
261 ///
262 /// for item in slices.1.iter_mut() {
263 /// // Do something with items
264 /// unsafe { std::ptr::drop_in_place(item); } // Ensure the items are dropped
265 /// }
266 ///
267 /// // Here head is 2 and len is 4
268 ///
269 /// let slices = queue.as_slices();
270 /// let as_previous: (&[u32], &[u32]) = (&[3, 4], &[5, 6]); // But the queue is still the same, while 2 elements were read
271 /// assert_eq!(slices, as_previous);
272 ///
273 /// unsafe { queue.dec_len_by(2); }
274 ///
275 /// // Here head is 2 and len is 2
276 ///
277 /// let slices = queue.as_slices();
278 /// let should_be: (&[u32], &[u32]) = (&[3, 4], &[]);
279 /// assert_eq!(slices, should_be); // Now it is valid
280 /// ```
281 pub unsafe fn dec_len_by(&mut self, number: usize) {
282 debug_assert!(self.len >= number);
283
284 self.len -= number;
285 }
286
287 /// Appends an element to the back of the queue.
288 ///
289 /// # Safety
290 ///
291 /// The caller must ensure that the queue is not full.
292 pub unsafe fn push_unchecked(&mut self, value: T) {
293 assert_hint(self.len() < N, "Tried to push to a full array stack");
294
295 let idx = self.to_physical_idx_from_head(self.len());
296
297 unsafe { ptr::write(self.array.get_unchecked_mut(idx), MaybeUninit::new(value)) };
298
299 self.len += 1;
300 }
301
302 /// Appends an element to the back of the queue or returns `Err(value)` if the queue is full.
303 pub fn push(&mut self, value: T) -> Result<(), T> {
304 if likely(self.len() < N) {
305 unsafe { self.push_unchecked(value) };
306
307 Ok(())
308 } else {
309 Err(value)
310 }
311 }
312
313 /// Pushes the provided value to the front of the queue.
314 ///
315 /// # Safety
316 ///
317 /// The caller must ensure that the queue is not full.
318 pub unsafe fn push_priority_value_unchecked(&mut self, value: T) {
319 assert_hint(self.len() < N, "Tried to push to a full array stack");
320
321 let phys_head = self.to_physical_idx_from_head(0);
322 let idx = if likely(phys_head > 0) {
323 phys_head - 1
324 } else {
325 N - 1
326 };
327
328 unsafe { ptr::write(self.array.get_unchecked_mut(idx), MaybeUninit::new(value)) };
329
330 self.head = idx;
331 self.len += 1;
332 }
333
334 /// Pushes the provided value to the front of the queue
335 /// or returns `Err(value)` if the queue is full.
336 ///
337 /// # Example
338 ///
339 /// ```rust
340 /// use orengine_utils::ArrayQueue;
341 ///
342 /// let mut queue = ArrayQueue::<u32, 3>::new();
343 ///
344 /// queue.push_priority_value(1).unwrap(); // [1, _, _]
345 /// queue.push(2).unwrap(); // [1, 2, _]
346 /// queue.push_priority_value(3).unwrap(); // [3, 1, 2]
347 ///
348 /// assert_eq!(queue.pop(), Some(3));
349 /// assert_eq!(queue.pop(), Some(1));
350 /// assert_eq!(queue.pop(), Some(2));
351 /// ```
352 pub fn push_priority_value(&mut self, value: T) -> Result<(), T> {
353 if likely(self.len() < N) {
354 unsafe { self.push_priority_value_unchecked(value) };
355
356 Ok(())
357 } else {
358 Err(value)
359 }
360 }
361
362 /// Removes the first element and returns it, or `None` if the queue is empty.
363 pub fn pop(&mut self) -> Option<T> {
364 if !self.is_empty() {
365 self.len -= 1;
366
367 let idx = self.head;
368 self.head = self.to_physical_idx_from_head(1);
369
370 assert_hint(
371 self.array.len() > idx,
372 &format!("idx: {}, len: {}", idx, self.array.len()),
373 );
374
375 Some(unsafe { self.array.get_unchecked_mut(idx).assume_init_read() })
376 } else {
377 None
378 }
379 }
380
381 /// Removes the last element and returns it, or `None` if the queue is empty.
382 ///
383 /// # Example
384 ///
385 /// ```rust
386 /// use orengine_utils::ArrayQueue;
387 ///
388 /// let mut queue = ArrayQueue::<u32, 3>::new();
389 ///
390 /// queue.push(1).unwrap(); // [1, _, _]
391 /// queue.push(2).unwrap(); // [1, 2, _]
392 /// queue.push(3).unwrap(); // [1, 2, 3]
393 ///
394 /// assert_eq!(queue.pop_less_priority_value(), Some(3));
395 /// assert_eq!(queue.pop(), Some(1));
396 /// assert_eq!(queue.pop(), Some(2));
397 /// ```
398 pub fn pop_less_priority_value(&mut self) -> Option<T> {
399 if !self.is_empty() {
400 self.len -= 1;
401
402 let idx = self.to_physical_idx_from_head(self.len());
403
404 Some(unsafe { self.array.get_unchecked_mut(idx).assume_init_read() })
405 } else {
406 None
407 }
408 }
409
410 /// Pushes a slice to the queue.
411 ///
412 /// It returns an error if the queue does not have enough space.
413 ///
414 /// # Safety
415 ///
416 /// It `T` is not `Copy`, the caller should [`forget`](mem::forget) the values.
417 #[inline]
418 pub unsafe fn extend_from_slice(&mut self, slice: &[T]) -> Result<(), NotEnoughSpace> {
419 if unlikely(self.len() + slice.len() > self.capacity()) {
420 return Err(NotEnoughSpace);
421 }
422
423 let phys_tail = self.to_physical_idx_from_head(self.len());
424 let right_space = self.capacity() - phys_tail;
425 let ptr = (&raw mut self.array).cast::<T>();
426
427 unsafe {
428 if slice.len() <= right_space {
429 // fits in one memcpy
430 ptr::copy_nonoverlapping(slice.as_ptr(), ptr.add(phys_tail), slice.len());
431 } else {
432 // wraparound case
433 ptr::copy_nonoverlapping(slice.as_ptr(), ptr.add(phys_tail), right_space);
434
435 ptr::copy_nonoverlapping(
436 slice.as_ptr().add(right_space),
437 ptr,
438 slice.len() - right_space,
439 );
440 }
441 }
442
443 self.len += slice.len();
444
445 Ok(())
446 }
447
448 /// Clears with calling the provided function on each element.
449 pub fn clear_with<F>(&mut self, mut f: F)
450 where
451 F: FnMut(T),
452 {
453 for i in 0..self.len {
454 let idx = self.to_physical_idx_from_head(i);
455
456 let value = unsafe { self.array.get_unchecked_mut(idx).assume_init_read() };
457
458 f(value);
459 }
460
461 self.len = 0;
462 }
463
464 /// Drops all elements in the queue and set the length to 0.
465 pub fn clear(&mut self) {
466 if mem::needs_drop::<T>() {
467 for i in 0..self.len {
468 let idx = self.to_physical_idx_from_head(i);
469
470 unsafe { ptr::drop_in_place(self.array.get_unchecked_mut(idx)) };
471 }
472 }
473
474 self.len = 0;
475 }
476
477 /// Returns a reference iterator over the queue.
478 pub fn iter(&self) -> impl ExactSizeIterator<Item = &T> {
479 struct Iter<'array_queue, T, const N: usize> {
480 queue: &'array_queue ArrayQueue<T, N>,
481 iterated: usize,
482 }
483
484 impl<'array_queue, T, const N: usize> Iterator for Iter<'array_queue, T, N> {
485 type Item = &'array_queue T;
486
487 fn next(&mut self) -> Option<Self::Item> {
488 if self.iterated < self.queue.len {
489 let idx = self.queue.to_physical_idx_from_head(self.iterated);
490
491 self.iterated += 1;
492
493 Some(unsafe { self.queue.array.get_unchecked(idx).assume_init_ref() })
494 } else {
495 None
496 }
497 }
498
499 fn size_hint(&self) -> (usize, Option<usize>) {
500 (
501 self.queue.len - self.iterated,
502 Some(self.queue.len - self.iterated),
503 )
504 }
505 }
506
507 impl<T, const N: usize> ExactSizeIterator for Iter<'_, T, N> {
508 fn len(&self) -> usize {
509 self.queue.len
510 }
511 }
512
513 Iter {
514 queue: self,
515 iterated: 0,
516 }
517 }
518
519 /// Returns a mutable reference iterator over the queue.
520 pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
521 struct IterMut<'array_queue, T, const N: usize> {
522 queue: &'array_queue mut ArrayQueue<T, N>,
523 iterated: usize,
524 }
525
526 impl<'array_queue, T, const N: usize> Iterator for IterMut<'array_queue, T, N> {
527 type Item = &'array_queue mut T;
528
529 fn next(&mut self) -> Option<Self::Item> {
530 if self.iterated < self.queue.len {
531 let idx = self.queue.to_physical_idx_from_head(self.iterated);
532
533 self.iterated += 1;
534
535 Some(unsafe {
536 &mut *ptr::from_mut(self.queue.array.get_unchecked_mut(idx)).cast::<T>()
537 })
538 } else {
539 None
540 }
541 }
542
543 fn size_hint(&self) -> (usize, Option<usize>) {
544 (
545 self.queue.len - self.iterated,
546 Some(self.queue.len - self.iterated),
547 )
548 }
549 }
550
551 impl<T, const N: usize> ExactSizeIterator for IterMut<'_, T, N> {
552 fn len(&self) -> usize {
553 self.queue.len
554 }
555 }
556
557 IterMut {
558 queue: self,
559 iterated: 0,
560 }
561 }
562
563 /// Refills the queue with elements provided by the function.
564 ///
565 /// # Safety
566 ///
567 /// The caller must ensure that the queue is empty before refilling.
568 pub unsafe fn refill_with(&mut self, f: impl FnOnce(&mut [MaybeUninit<T>; N]) -> usize) {
569 debug_assert!(
570 self.is_empty(),
571 "ArrayQueue should be empty before refilling"
572 );
573
574 let filled = f(&mut self.array);
575
576 debug_assert!(filled <= N, "Filled more than the capacity");
577
578 self.len = filled;
579 self.head = 0;
580 }
581}
582
583impl<T, const N: usize> Default for ArrayQueue<T, N> {
584 fn default() -> Self {
585 Self::new()
586 }
587}
588
589impl<T, const N: usize> From<[T; N]> for ArrayQueue<T, N> {
590 fn from(array: [T; N]) -> Self {
591 Self {
592 array: unsafe { (&raw const array).cast::<[MaybeUninit<T>; N]>().read() },
593 len: N,
594 head: 0,
595 }
596 }
597}
598
599impl<T, const N: usize> Drop for ArrayQueue<T, N> {
600 fn drop(&mut self) {
601 self.clear();
602 }
603}
604
605#[cfg(test)]
606mod tests {
607 use super::*;
608 use alloc::vec;
609 use alloc::vec::Vec;
610
611 #[test]
612 fn test_array_queue() {
613 let mut queue = ArrayQueue::<u32, 4>::new();
614
615 unsafe {
616 queue.push_unchecked(1);
617 assert_eq!(queue.len(), 1);
618
619 queue.push_unchecked(2);
620 assert_eq!(queue.len(), 2);
621
622 queue.push(3).unwrap();
623 assert_eq!(queue.len(), 3);
624
625 assert_eq!(queue.pop(), Some(1));
626 assert_eq!(queue.len(), 2);
627
628 queue.push_unchecked(4);
629 assert_eq!(queue.len(), 3);
630
631 queue.push_unchecked(5);
632 assert_eq!(queue.len(), 4);
633
634 assert_eq!(queue.push(6), Err(6));
635
636 assert_eq!(queue.pop(), Some(2));
637 assert_eq!(queue.pop(), Some(3));
638 assert_eq!(queue.pop(), Some(4));
639 assert_eq!(queue.pop(), Some(5));
640 assert_eq!(queue.pop(), None);
641 }
642 }
643
644 #[test]
645 fn test_array_queue_iterators() {
646 let mut queue = ArrayQueue::<u32, 4>::new();
647
648 unsafe {
649 queue.push_unchecked(1);
650 queue.push_unchecked(2);
651 queue.push_unchecked(3);
652 queue.push_unchecked(4);
653 }
654
655 assert_eq!(queue.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
656 assert_eq!(
657 queue.iter_mut().collect::<Vec<_>>(),
658 vec![&mut 1, &mut 2, &mut 3, &mut 4]
659 );
660 }
661
662 #[test]
663 fn test_array_queue_refill_with() {
664 let mut queue = ArrayQueue::<u32, 4>::new();
665
666 unsafe {
667 queue.refill_with(|array| {
668 array.copy_from_slice(&[
669 MaybeUninit::new(1),
670 MaybeUninit::new(2),
671 MaybeUninit::new(3),
672 MaybeUninit::new(4),
673 ]);
674
675 4
676 });
677 };
678
679 assert_eq!(queue.len(), 4);
680 assert_eq!(queue.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
681 }
682
683 #[test]
684 fn test_array_queue_extend_from_slice() {
685 // --- CASE 1: without wraparound ---
686 {
687 let mut q = ArrayQueue::<usize, 8>::new();
688
689 unsafe {
690 q.extend_from_slice(&[1, 2, 3]).unwrap();
691 }
692
693 assert_eq!(q.len(), 3);
694
695 assert_eq!(q.pop().unwrap(), 1);
696 assert_eq!(q.pop().unwrap(), 2);
697 assert_eq!(q.pop().unwrap(), 3);
698
699 // With the updated head index
700
701 unsafe {
702 q.extend_from_slice(&[1, 2, 3]).unwrap();
703 }
704
705 assert_eq!(q.len(), 3);
706
707 assert_eq!(q.pop().unwrap(), 1);
708 assert_eq!(q.pop().unwrap(), 2);
709 assert_eq!(q.pop().unwrap(), 3);
710 }
711
712 // --- CASE 2: wraparound ---
713 {
714 let mut q = ArrayQueue::<usize, 8>::new();
715
716 unsafe {
717 q.extend_from_slice(&[1, 2, 3, 4, 5, 6, 7]).unwrap();
718 };
719
720 q.pop().unwrap();
721 q.pop().unwrap();
722 q.pop().unwrap();
723 q.pop().unwrap();
724
725 assert_eq!(q.len(), 3);
726
727 unsafe {
728 q.extend_from_slice(&[50, 51, 52, 53, 54]).unwrap();
729 }
730
731 assert_eq!(q.len(), 8);
732
733 let mut actual = Vec::new();
734 while let Some(value) = q.pop() {
735 actual.push(value);
736 }
737
738 assert_eq!(&actual, &[5, 6, 7, 50, 51, 52, 53, 54]);
739 }
740
741 // --- CASE 4: overflow → Err ---
742 {
743 let mut q = ArrayQueue::<usize, 4>::new();
744
745 unsafe {
746 q.extend_from_slice(&[1, 2, 3]).unwrap();
747 };
748
749 assert!(unsafe { q.extend_from_slice(&[9, 9]) }.is_err());
750 assert_eq!(q.len(), 3, "len must remain unchanged");
751 }
752 }
753}