priority_set/
lib.rs

1//! A fixed size priority set suitable for no_std use.
2//!
3//! Example:
4//!
5//! ```rust
6//! # use priority_set::*;
7//!
8//! #[derive(PartialEq, Debug)]
9//! enum Command {
10//!     QueryServerA,
11//!     QueryServerB,
12//! }
13//!
14//! fn main() {
15//!     // Create a priority set with 10 slots
16//!     let mut p: PrioritySet<Command, 10> = PrioritySet::new();
17//!
18//!     // Insert two items
19//!     p.insert(Priority(10), Command::QueryServerA);
20//!     p.insert(Priority(20), Command::QueryServerB);
21//!     p.insert(Priority(30), Command::QueryServerA);
22//!
23//!     // We inserted a duplicate command, so its priority was updated, but no new item was added
24//!     assert_eq!(p.len(), 2);
25//!
26//!     // Pops the highest priority item, which is QueryServerA with Priority(30)
27//!     assert_eq!(p.pop(), Some(Command::QueryServerA));
28//!     assert_eq!(p.pop(), Some(Command::QueryServerB));
29//!     assert_eq!(p.pop(), None);
30//! }
31
32#![cfg_attr(not(test), no_std)]
33
34#[cfg(test)]
35#[macro_use(quickcheck)]
36extern crate quickcheck_macros;
37
38use core::mem;
39use core::cmp;
40
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
43pub struct Priority(pub usize);
44
45impl Priority {
46    pub const MIN: Priority = Priority(usize::MIN);
47    pub const MAX: Priority = Priority(usize::MAX);
48}
49
50impl From<usize> for Priority {
51    fn from(i: usize) -> Self {
52        Self(i)
53    }
54}
55
56impl From<Priority> for usize {
57    fn from(p: Priority) -> Self {
58        p.0
59    }
60}
61
62#[derive(Debug, PartialEq)]
63pub struct PriorityEntry<I: PartialEq> {
64    pub item: I,
65    pub priority: Priority,
66}
67
68impl<I: PartialEq + Clone> Clone for PriorityEntry<I> {
69    fn clone(&self) -> Self {
70        Self {
71            item: self.item.clone(),
72            priority: self.priority.clone(),
73        }
74    }
75}
76
77/// The result of a [PrioritySet::insert] operation.
78#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
79pub enum InsertResult {
80    /// The item was inserted into an empty slot.
81    Inserted,
82    /// The item already exists in the set, and the priority was updated.
83    Updated,
84    /// We replaced an existing entry with a lower priority.
85    Replaced,
86    /// There was no space for this item, and all other items had higher priority, so we dropped it.
87    Dropped,
88}
89
90
91/// PrioritySet is a fixed size Priority Set.
92///
93/// It never allocates, instead dropping the lowest priority item when
94/// inserting a new one.
95#[derive(Debug)]
96pub struct PrioritySet<I: PartialEq, const N: usize> {
97    items: [Option<PriorityEntry<I>>; N],
98}
99
100
101impl<I: PartialEq, const N: usize> PrioritySet<I, N> {
102    /// Initializes a new empty `PrioritySet` with `N` free slots.
103    pub fn new() -> Self {
104        Self {
105            items: array_init_none(),
106        }
107    }
108
109    /// Counts the number of occupied entries in the set
110    pub fn len(&self) -> usize {
111        self.items
112            .iter()
113            .filter(|x| x.is_some())
114            .count()
115    }
116
117    /// Returns `true` if all slots in the priority set are occupied
118    pub fn is_full(&self) -> bool {
119        self.len() == N
120    }
121
122    /// Inserts an item with priority.
123    ///
124    /// If the item already exists in the set, the priority is updated. The highest priority is chosen
125    /// between the existing priority and the new priority.
126    ///
127    /// If the set is full the least prioritary item is dropped. If all items are of higher priority
128    /// than the item being inserted, no change occurs.
129    ///
130    /// Returns `true` if the item was inserted, `false` if it was dropped.
131    pub fn insert(&mut self, priority: Priority, item: I) -> InsertResult {
132        let new_entry = PriorityEntry {
133            priority,
134            item,
135        };
136
137        // Check if item exists, and update its priority
138        if let Some(entry) = self.entry_mut(&new_entry.item) {
139            // The inserted priority was lower than the existing one, so we drop this insert
140            if entry.priority > priority {
141                return InsertResult::Dropped;
142            }
143
144            entry.priority = cmp::max(entry.priority, priority);
145            return InsertResult::Updated;
146        }
147
148        // Try to find an open slot.
149        let empty_slot = self.items
150            .iter_mut()
151            .find(|s| s.is_none());
152
153        if let Some(slot) = empty_slot {
154            let _ = mem::replace(slot, Some(new_entry));
155            return InsertResult::Inserted;
156        }
157
158        // If we can't find a open slot, lets find the one that has the lowest priority
159        // and that has a lower priority than the item being inserted.
160        let replacement_slot = self.items
161            .iter_mut()
162            .min_by_key(|slot| slot.as_ref().unwrap().priority)
163            .and_then(|slot| if slot.as_ref().unwrap().priority > priority {
164                None
165            } else {
166                Some(slot)
167            });
168
169        if let Some(slot) = replacement_slot {
170            let _ = mem::replace(slot, Some(new_entry));
171            return InsertResult::Replaced;
172        }
173
174        return InsertResult::Dropped;
175    }
176
177    /// Gets the priority of a item, if it exists
178    pub fn priority(&self, item: &I) -> Option<Priority> {
179        self.entry(item)
180            .map(|entry| entry.priority)
181
182    }
183
184    /// Returns an iterator over the entries
185    ///
186    /// Iteration order is not guaranteed.
187    pub fn iter(&self) -> impl Iterator<Item = &PriorityEntry<I>> {
188        self.items
189            .iter()
190            .flatten()
191    }
192
193    /// Returns a mutable iterator over the entries
194    ///
195    /// Iteration order is not guaranteed.
196    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut PriorityEntry<I>> {
197        self.items
198            .iter_mut()
199            .flatten()
200    }
201
202    /// Finds an item entry
203    pub fn entry(&self, item: &I) -> Option<&PriorityEntry<I>> {
204        self.iter().find(|e| e.item == *item)
205    }
206
207    /// Returns a mutable entry to an item
208    pub fn entry_mut(&mut self, item: &I) -> Option<&mut PriorityEntry<I>> {
209        self.iter_mut().find(|e| e.item == *item)
210    }
211
212    /// Pops the highest priority item from the set
213    pub fn pop(&mut self) -> Option<I> {
214        self.pop_entry()
215            .map(|entry| entry.item)
216    }
217
218    /// Pops the highest priority item entry from the set
219    pub fn pop_entry(&mut self) -> Option<PriorityEntry<I>> {
220        let slot = self.items
221            .iter_mut()
222            .filter(|entry| entry.is_some())
223            .max_by_key(|slot| slot.as_ref().unwrap().priority);
224
225        if let Some(entry) = slot {
226            mem::replace(entry, None)
227        } else {
228            None
229        }
230    }
231}
232
233impl<I: PartialEq + Clone, const N: usize> Clone for PrioritySet<I, N> {
234    fn clone(&self) -> Self {
235        PrioritySet {
236            items: self.items.clone()
237        }
238    }
239}
240
241/// Initializes an array of fixed size `N` with `None`.
242fn array_init_none<T, const N: usize>() -> [Option<T>; N] {
243    let mut items: [Option<T>; N] = unsafe { core::mem::zeroed() };
244    for slot in items.iter_mut() {
245        *slot = None;
246    }
247    items
248}
249
250
251#[cfg(test)]
252mod test {
253    use super::*;
254
255    #[test]
256    fn new() {
257        let p: PrioritySet<i32, 10> = PrioritySet::new();
258
259        assert_eq!(p.len(), 0);
260        assert!(!p.is_full());
261    }
262
263    #[test]
264    fn insert_increases_len() {
265        let mut p: PrioritySet<i32, 10> = PrioritySet::new();
266
267        assert_eq!(p.len(), 0);
268        p.insert(Priority(10), 10);
269        assert_eq!(p.len(), 1);
270    }
271
272    #[test]
273    fn insert_updates_on_duplicate_items() {
274        let mut p: PrioritySet<i32, 10> = PrioritySet::new();
275
276        assert_eq!(p.insert(Priority(10), 10), InsertResult::Inserted);
277        assert_eq!(p.insert(Priority(20), 10), InsertResult::Updated);
278        assert_eq!(p.len(), 1);
279    }
280
281    #[test]
282    fn insert_drops_when_full_with_higher_priority_items() {
283        let mut p: PrioritySet<i32, 2> = PrioritySet::new();
284
285        assert_eq!(p.insert(Priority(10), 10), InsertResult::Inserted);
286        assert_eq!(p.insert(Priority(20), 11), InsertResult::Inserted);
287        assert_eq!(p.insert(Priority(5), 12), InsertResult::Dropped);
288
289        assert_eq!(p.entry(&12), None);
290    }
291
292    #[test]
293    fn insert_replaces_items_with_lower_priority_when_full() {
294        let mut p: PrioritySet<i32, 2> = PrioritySet::new();
295
296        assert_eq!(p.len(), 0);
297        assert_eq!(p.insert(Priority(10), 10), InsertResult::Inserted);
298        assert_eq!(p.insert(Priority(20), 11), InsertResult::Inserted);
299
300        // Replaces item 10, which only has 10 priority
301        assert_eq!(p.insert(Priority(15), 12), InsertResult::Replaced);
302
303        assert!(p.entry(&11).is_some());
304        assert!(p.entry(&12).is_some());
305    }
306
307    #[test]
308    fn insert_replaces_the_lowest_priority_item() {
309        let mut p: PrioritySet<i32, 3> = PrioritySet::new();
310
311        assert_eq!(p.insert(Priority(20), 10), InsertResult::Inserted);
312        assert_eq!(p.insert(Priority(10), 11), InsertResult::Inserted);
313        assert_eq!(p.insert(Priority(15), 12), InsertResult::Inserted);
314        assert_eq!(p.insert(Priority(30), 13), InsertResult::Replaced);
315
316        assert!(p.entry(&10).is_some());
317        assert!(p.entry(&11).is_none());
318        assert!(p.entry(&12).is_some());
319        assert!(p.entry(&13).is_some());
320    }
321
322    #[test]
323    fn insert_updates_the_priority_of_an_item_if_it_already_exists() {
324        let mut p: PrioritySet<i32, 3> = PrioritySet::new();
325
326        assert_eq!(p.insert(Priority(10), 10), InsertResult::Inserted);
327        assert_eq!(p.insert(Priority(20), 10), InsertResult::Updated);
328        assert_eq!(p.insert(Priority(5), 10), InsertResult::Dropped);
329
330        assert_eq!(p.priority(&10), Some(Priority(20)));
331    }
332
333    #[test]
334    fn pop_gets_the_highest_priority_item() {
335        let mut p: PrioritySet<i32, 3> = PrioritySet::new();
336
337        p.insert(Priority(10), 10);
338        p.insert(Priority(20), 20);
339        p.insert(Priority(15), 30);
340
341        assert_eq!(p.pop(), Some(20));
342        assert_eq!(p.pop(), Some(30));
343        assert_eq!(p.pop(), Some(10));
344        assert_eq!(p.len(), 0);
345    }
346
347    #[test]
348    fn iter_iterates_over_all_entries() {
349        let mut p: PrioritySet<i32, 10> = PrioritySet::new();
350
351        assert_eq!(p.insert(Priority(10), 10), InsertResult::Inserted);
352        assert_eq!(p.insert(Priority(20), 11), InsertResult::Inserted);
353
354        let mut values: Vec<_> = p.iter().collect();
355        values.sort_by_key(|i| i.priority);
356
357        assert_eq!(values, [
358            &PriorityEntry {
359                priority: Priority(10),
360                item: 10
361            },
362            &PriorityEntry {
363                priority: Priority(20),
364                item: 11
365            }
366        ]);
367    }
368}
369
370#[cfg(test)]
371mod quickcheck_tests {
372    use super::*;
373    use std::collections::HashSet;
374
375    #[quickcheck]
376    fn insert_never_crashes(input: Vec<(usize, i32)>) -> bool {
377        let mut p: PrioritySet<i32, 16> = PrioritySet::new();
378        for (priority, item) in input {
379            p.insert(Priority(priority), item);
380        }
381        true
382    }
383
384    #[quickcheck]
385    fn is_a_set(input: Vec<(usize, i32)>) -> bool {
386        let input: Vec<_> = input.into_iter().take(32).collect();
387
388        let mut p: PrioritySet<i32, 32> = PrioritySet::new();
389        for (priority, item) in input.iter() {
390            p.insert(Priority(*priority), *item);
391        }
392
393        let mut set = HashSet::new();
394        for (_, item) in input.iter() {
395            set.insert(*item);
396        }
397
398        let mut p_items: Vec<_> = p.iter().map(|p| p.item).collect();
399        p_items.sort();
400
401        let mut h_items: Vec<_> = set.iter().map(|i| *i).collect();
402        h_items.sort();
403
404        p_items == h_items
405    }
406
407
408    #[quickcheck]
409    fn pop_sorts_by_priority(input: Vec<(usize, i32)>) -> bool {
410        let input: Vec<_> = input.into_iter().take(32).collect();
411
412        let mut p: PrioritySet<i32, 32> = PrioritySet::new();
413        for (priority, item) in input.iter() {
414            p.insert(Priority(*priority), *item);
415        }
416
417        let mut prev_priority = Priority::MAX;
418        loop {
419            match p.pop_entry() {
420                Some(popped) => {
421                    if popped.priority > prev_priority {
422                        return false;
423                    }
424
425                    prev_priority = popped.priority;
426                },
427                None => {
428                    return true;
429                }
430            };
431        }
432    }
433}