kernel/mem/heap.rs
1//! This module provides a binary heap implementation.
2
3use super::array::Vec;
4use crate::utils::KernelError;
5
6/// An array-based binary heap, with N elements stored inline.
7pub struct BinaryHeap<T, const N: usize> {
8 vec: Vec<T, N>,
9}
10
11impl<T: Clone + Copy + Ord, const N: usize> BinaryHeap<T, N> {
12 /// Create a new empty binary heap.
13 pub const fn new() -> Self {
14 Self { vec: Vec::new() }
15 }
16
17 /// Push a value onto the binary heap.
18 ///
19 /// `value` - The value to push onto the binary heap.
20 ///
21 /// Returns `Ok(())` if the value was pushed onto the binary heap, or an error if the heap cannot be extended (e.g. OOM).
22 pub fn push(&mut self, value: T) -> Result<(), KernelError> {
23 self.vec.push(value)?;
24 self.sift_up(self.len() - 1);
25 Ok(())
26 }
27
28 /// Pop the smallest value from the binary heap.
29 ///
30 /// Returns the smallest value in the binary heap, or `None` if the heap is empty.
31 pub fn pop(&mut self) -> Option<T> {
32 if self.is_empty() {
33 return None;
34 }
35
36 let value = self.peek().cloned();
37 self.vec.swap(0, self.len() - 1);
38 self.vec.pop();
39 self.sift_down(0);
40 value
41 }
42
43 /// Sift the value at the given index up the binary heap.
44 ///
45 /// `index` - The index of the value to sift up.
46 fn sift_up(&mut self, mut index: usize) {
47 // We move up the heap until we reach the root or the parent is smaller than the current value.
48 while index > 0 {
49 let parent = (index - 1) / 2;
50 if self.vec.at(parent) <= self.vec.at(index) {
51 break;
52 }
53 self.vec.swap(parent, index);
54 index = parent;
55 }
56 }
57
58 /// Sift the value at the given index down the binary heap.
59 ///
60 /// `index` - The index of the value to sift down.
61 fn sift_down(&mut self, mut index: usize) {
62 // We move down the heap until we reach a leaf or the value is smaller than both children.
63 while index < self.len() {
64 let left = 2 * index + 1;
65 let right = 2 * index + 2;
66 let mut smallest = index;
67
68 if left < self.len() && self.vec.at(left) < self.vec.at(smallest) {
69 smallest = left;
70 }
71
72 if right < self.len() && self.vec.at(right) < self.vec.at(smallest) {
73 smallest = right;
74 }
75
76 if smallest == index {
77 break;
78 }
79
80 self.vec.swap(smallest, index);
81 index = smallest;
82 }
83 }
84
85 /// Check if the binary heap is empty.
86 ///
87 /// Returns `true` if the binary heap is empty, `false` otherwise.
88 pub fn is_empty(&self) -> bool {
89 self.len() == 0
90 }
91
92 /// Peek at the smallest value in the binary heap.
93 ///
94 /// Returns the smallest value in the binary heap, or `None` if the heap is empty.
95 pub fn peek(&self) -> Option<&T> {
96 if self.is_empty() {
97 return None;
98 }
99 self.vec.at(0)
100 }
101
102 /// Get the number of elements in the binary heap.
103 pub fn len(&self) -> usize {
104 self.vec.len()
105 }
106}