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}