dfir_rs/util/
priority_stack.rs

Help
1//! A priority queue in which elements of the same priority are popped in a LIFO order.
2
3use smallvec::SmallVec;
4
5/// A priority stack in which elements of the same priority are popped in a LIFO order.
6// TODO(mingwei): Keep an upper bound on current priority to avoid scanning all stacks?
7#[derive(Debug, Clone)]
8pub struct PriorityStack<T> {
9    /// Note: inner stack `Vec`s may be empty.
10    stacks: Vec<SmallVec<[T; 1]>>,
11}
12
13impl<T> PriorityStack<T> {
14    /// Creates a new, empty `PriorityStack`.
15    pub fn new() -> Self {
16        Self {
17            stacks: Vec::default(),
18        }
19    }
20
21    /// Creates a new, empty `PriorityStack` with pre-allocated capacity up to the given priority.
22    pub fn with_priority_capacity(priority: usize) -> Self {
23        Self {
24            stacks: Vec::with_capacity(priority),
25        }
26    }
27
28    /// Pushes an element onto the stack with the given priority.
29    pub fn push(&mut self, priority: usize, item: T) {
30        if priority >= self.stacks.len() {
31            self.stacks.resize_with(priority + 1, Default::default);
32        }
33        self.stacks[priority].push(item);
34    }
35
36    /// Pops an element from the stack with the highest priority.
37    pub fn pop(&mut self) -> Option<T> {
38        self.stacks
39            .iter_mut()
40            .rev()
41            .filter_map(SmallVec::pop)
42            .next()
43    }
44
45    /// Pops an element from the stack and return `(priority, item)`.
46    pub fn pop_prio(&mut self) -> Option<(usize, T)> {
47        self.stacks
48            .iter_mut()
49            .enumerate()
50            .rev()
51            .filter_map(|(i, stack)| stack.pop().map(|x| (i, x)))
52            .next()
53    }
54
55    /// Returns the item with the highest priority without removing it.
56    pub fn peek(&self) -> Option<&T> {
57        self.stacks
58            .iter()
59            .rev()
60            .filter_map(|stack| stack.last())
61            .next()
62    }
63
64    /// Returns the item with the highest priority and its priority without removing it.
65    pub fn peek_prio(&self) -> Option<(usize, &T)> {
66        self.stacks
67            .iter()
68            .enumerate()
69            .rev()
70            .filter_map(|(i, stack)| stack.last().map(|x| (i, x)))
71            .next()
72    }
73
74    /// Returns the number of elements in the `PriorityStack`.
75    pub fn len(&self) -> usize {
76        self.stacks.iter().map(SmallVec::len).sum()
77    }
78
79    /// Returns true if the `PriorityStack` is empty.
80    pub fn is_empty(&self) -> bool {
81        self.stacks.is_empty()
82    }
83}
84
85impl<T> Default for PriorityStack<T> {
86    fn default() -> Self {
87        Self::new()
88    }
89}
90
91impl<T> Extend<(usize, T)> for PriorityStack<T> {
92    fn extend<I: IntoIterator<Item = (usize, T)>>(&mut self, iter: I) {
93        for (priority, item) in iter {
94            self.push(priority, item);
95        }
96    }
97}
dfir_rs/util/priority_stack.rs

dfir_rs/util/
priority_stack.rs
