osiris/types/
heap.rs

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