aionbot_core/
queue.rs

1use std::cmp::Reverse;
2use std::collections::{BinaryHeap, HashMap, VecDeque};
3use std::hash::Hash;
4
5#[derive(Clone, Eq, PartialEq)]
6struct EventEntry<T> {
7    priority: i8,
8    counter: usize,
9    item: T,
10}
11
12impl<T: Eq + PartialEq> EventEntry<T> {
13    fn new(priority: i8, counter: usize, item: T) -> Self {
14        Self {
15            priority,
16            counter,
17            item,
18        }
19    }
20}
21
22impl<T: Eq + PartialEq> Ord for EventEntry<T> {
23    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
24        // In order to make BinaryHeap a min-heap, we use `Reverse`.
25        Reverse(self.priority)
26            .cmp(&Reverse(other.priority))
27            .then(self.counter.cmp(&other.counter))
28    }
29}
30
31impl<T: Eq + PartialEq> PartialOrd for EventEntry<T> {
32    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
33        Some(self.cmp(other))
34    }
35}
36
37pub struct EventQueue<T: PartialEq + Hash + Clone> {
38    heap: BinaryHeap<EventEntry<T>>,
39    entry_finder: HashMap<T, Option<EventEntry<T>>>,
40    counter: usize,
41    order_queue: VecDeque<T>,
42}
43
44impl<T: Eq + Hash + Clone> EventQueue<T> {
45    pub fn new() -> Self {
46        Self {
47            heap: BinaryHeap::new(),
48            entry_finder: HashMap::new(),
49            counter: 0,
50            order_queue: VecDeque::new(),
51        }
52    }
53
54    pub fn push(&mut self, priority: i8, item: T) {
55        if self.entry_finder.contains_key(&item) {
56            self.remove(&item);
57        }
58        let entry = EventEntry::new(priority, self.counter, item.clone());
59        self.counter += 1;
60        self.entry_finder.insert(item.clone(), Some(entry.clone()));
61        self.heap.push(entry);
62        self.order_queue.push_front(item);
63    }
64
65    pub fn pop(&mut self) -> Option<T> {
66        while let Some(EventEntry { item, .. }) = self.heap.pop() {
67            if let Some(Some(_)) = self.entry_finder.remove(&item) {
68                self.order_queue.retain(|x| x != &item);
69                return Some(item);
70            }
71        }
72        None
73    }
74
75    pub fn remove(&mut self, item: &T) {
76        if let Some(Some(mut entry)) = self.entry_finder.remove(item) {
77            entry.item = item.clone(); // Invalidate the entry
78            self.order_queue.retain(|x| x != item);
79        }
80    }
81
82    pub fn is_empty(&self) -> bool {
83        self.heap.is_empty()
84    }
85
86    pub fn len(&self) -> usize {
87        self.heap.len()
88    }
89}
90
91impl<T: Eq + Hash + Clone> Default for EventQueue<T> {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_event_queue() {
103        let mut queue = EventQueue::new();
104        queue.push(1, "a");
105        queue.push(2, "b");
106        queue.push(3, "c");
107        queue.push(2, "d");
108        queue.push(1, "e");
109        assert_eq!(queue.pop(), Some("e"));
110        assert_eq!(queue.pop(), Some("a"));
111        assert_eq!(queue.pop(), Some("d"));
112        assert_eq!(queue.pop(), Some("b"));
113        assert_eq!(queue.pop(), Some("c"));
114        assert_eq!(queue.pop(), None);
115    }
116}
aionbot_core/queue.rs

aionbot_core/
queue.rs
