utils/
queue.rs

1//! Queue data structures.
2
3use std::fmt::{Debug, Formatter};
4
5/// A priority queue using a circular array of buckets for priorities within a sliding window.
6///
7/// New items must have a priority within `current_priority..current_priority + N`.
8/// `current_priority` is only updated on pop and when the first item is pushed.
9///
10/// Items are popped in LIFO order within each priority bucket as [`Vec`] is used internally to
11/// avoid the extra overhead of [`VecDeque`](std::collections::VecDeque).
12///
13/// # Example
14/// ```
15/// # use utils::queue::BucketQueue;
16/// let mut queue = BucketQueue::<i32, 8>::new();
17/// queue.push(0, 100);
18/// queue.push(2, 200);
19/// queue.push(7, 400);
20/// queue.push(2, 300);
21/// assert_eq!(queue.pop(), Some(100));
22/// assert_eq!(queue.pop(), Some(300));
23/// assert_eq!(queue.pop(), Some(200));
24/// assert_eq!(queue.peek(), Some(&400));
25/// assert_eq!(queue.peek_entry(), Some((7, &400)));
26/// assert_eq!(queue.pop_entry(), Some((7, 400)));
27/// assert_eq!(queue.pop(), None);
28/// assert_eq!(queue.peek(), None);
29/// ```
30#[must_use]
31#[derive(Clone)]
32pub struct BucketQueue<T, const N: usize> {
33    buckets: [Vec<T>; N],
34    current_priority: usize,
35    size: usize,
36}
37
38impl<T, const N: usize> BucketQueue<T, N> {
39    #[inline]
40    pub fn new() -> Self {
41        Self::default()
42    }
43
44    #[inline]
45    pub fn with_capacity(per_bucket_capacity: usize) -> Self {
46        BucketQueue {
47            buckets: std::array::from_fn(|_| Vec::with_capacity(per_bucket_capacity)),
48            current_priority: 0,
49            size: 0,
50        }
51    }
52
53    #[inline]
54    pub fn push(&mut self, priority: usize, value: T) {
55        if self.size == 0 {
56            self.current_priority = priority;
57        } else {
58            assert!(
59                priority >= self.current_priority && priority < self.current_priority + N,
60                "priority {priority} out of range {}..{}",
61                self.current_priority,
62                self.current_priority + N,
63            );
64        }
65
66        self.buckets[priority % N].push(value);
67        self.size += 1;
68    }
69
70    #[inline]
71    #[must_use]
72    pub fn pop(&mut self) -> Option<T> {
73        self.pop_entry().map(|(_, v)| v)
74    }
75
76    #[inline]
77    #[must_use]
78    pub fn pop_entry(&mut self) -> Option<(usize, T)> {
79        if self.size == 0 {
80            return None;
81        }
82
83        loop {
84            let idx = self.current_priority % N;
85            if let Some(value) = self.buckets[idx].pop() {
86                self.size -= 1;
87                return Some((self.current_priority, value));
88            }
89            self.current_priority += 1;
90        }
91    }
92
93    #[inline]
94    #[must_use]
95    pub fn peek(&self) -> Option<&T> {
96        self.peek_entry().map(|(_, v)| v)
97    }
98
99    #[inline]
100    #[must_use]
101    pub fn peek_entry(&self) -> Option<(usize, &T)> {
102        if self.size == 0 {
103            return None;
104        }
105
106        // peek takes &self so can't advance current_priority
107        for priority in self.current_priority..self.current_priority + N {
108            if let Some(value) = self.buckets[priority % N].last() {
109                return Some((priority, value));
110            }
111        }
112        unreachable!()
113    }
114
115    #[inline]
116    pub fn clear(&mut self) {
117        self.buckets.iter_mut().for_each(Vec::clear);
118        self.size = 0;
119    }
120
121    #[inline]
122    #[must_use]
123    pub fn len(&self) -> usize {
124        self.size
125    }
126
127    #[inline]
128    #[must_use]
129    pub fn is_empty(&self) -> bool {
130        self.size == 0
131    }
132}
133
134impl<T, const N: usize> Default for BucketQueue<T, N> {
135    #[inline]
136    fn default() -> Self {
137        BucketQueue {
138            buckets: std::array::from_fn(|_| Vec::new()),
139            current_priority: 0,
140            size: 0,
141        }
142    }
143}
144
145impl<T: Debug, const N: usize> Debug for BucketQueue<T, N> {
146    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
147        f.debug_list()
148            .entries(
149                (self.current_priority..self.current_priority + N).flat_map(|priority| {
150                    self.buckets[priority % N]
151                        .iter()
152                        .rev() // Match LIFO order within buckets
153                        .map(move |value| (priority, value))
154                }),
155            )
156            .finish()
157    }
158}