Skip to main content

qsm_core/
priority_queue.rs

1/// Bucket-based priority queue for O(1) push/pop operations
2/// Used for region growing phase unwrapping where priorities are discrete (0-255)
3///
4/// Higher priority = better quality = processed first (starts from highest bin)
5
6pub struct BucketQueue<T> {
7    bins: Vec<Vec<T>>,
8    current_priority: isize,  // Can go negative when exhausted
9    count: usize,
10    n_bins: usize,
11}
12
13impl<T> BucketQueue<T> {
14    pub fn new(n_bins: usize) -> Self {
15        BucketQueue {
16            bins: (0..n_bins).map(|_| Vec::new()).collect(),
17            current_priority: (n_bins - 1) as isize,  // Start at highest priority
18            count: 0,
19            n_bins,
20        }
21    }
22
23    #[inline]
24    #[allow(dead_code)]
25    pub fn is_empty(&self) -> bool {
26        self.count == 0
27    }
28
29    #[inline]
30    pub fn push(&mut self, priority: usize, item: T) {
31        let priority = priority.min(self.n_bins - 1);
32        self.bins[priority].push(item);
33        self.count += 1;
34        // Update current_priority if this is higher (better quality)
35        if (priority as isize) > self.current_priority {
36            self.current_priority = priority as isize;
37        }
38    }
39
40    #[inline]
41    pub fn pop(&mut self) -> Option<T> {
42        if self.count == 0 {
43            return None;
44        }
45
46        // Find next non-empty bin starting from current priority (going DOWN)
47        while self.current_priority >= 0 && self.bins[self.current_priority as usize].is_empty() {
48            self.current_priority -= 1;
49        }
50
51        if self.current_priority < 0 {
52            return None;
53        }
54
55        self.count -= 1;
56        self.bins[self.current_priority as usize].pop()
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn test_bucket_queue() {
66        let mut queue: BucketQueue<(usize, usize, usize)> = BucketQueue::new(256);
67        assert!(queue.is_empty());
68
69        queue.push(100, (1, 2, 3));
70        queue.push(50, (4, 5, 6));
71        queue.push(200, (7, 8, 9));
72
73        assert!(!queue.is_empty());
74
75        // Should pop HIGHEST priority first (better quality edges first)
76        assert_eq!(queue.pop(), Some((7, 8, 9)));  // priority 200
77        assert_eq!(queue.pop(), Some((1, 2, 3)));  // priority 100
78        assert_eq!(queue.pop(), Some((4, 5, 6)));  // priority 50
79        assert!(queue.is_empty());
80        assert_eq!(queue.pop(), None);
81    }
82}