heapless/binary_heap.rs
1//! A priority queue implemented with a binary heap.
2//!
3//! Insertion and popping the largest element have `O(log n)` time complexity. Checking the largest
4//! / smallest element is `O(1)`.
5
6// TODO not yet implemented
7// Converting a vector to a binary heap can be done in-place, and has `O(n)` complexity. A binary
8// heap can also be converted to a sorted vector in-place, allowing it to be used for an `O(n log
9// n)` in-place heapsort.
10
11use core::{
12 cmp::Ordering,
13 fmt,
14 marker::PhantomData,
15 mem::{self, ManuallyDrop},
16 ops::{Deref, DerefMut},
17 ptr, slice,
18};
19
20use crate::vec::Vec;
21
22/// Min-heap
23pub enum Min {}
24
25/// Max-heap
26pub enum Max {}
27
28/// The binary heap kind: min-heap or max-heap
29pub trait Kind: private::Sealed {
30 #[doc(hidden)]
31 fn ordering() -> Ordering;
32}
33
34impl Kind for Min {
35 fn ordering() -> Ordering {
36 Ordering::Less
37 }
38}
39
40impl Kind for Max {
41 fn ordering() -> Ordering {
42 Ordering::Greater
43 }
44}
45
46/// Sealed traits
47mod private {
48 pub trait Sealed {}
49}
50
51impl private::Sealed for Max {}
52impl private::Sealed for Min {}
53
54/// A priority queue implemented with a binary heap.
55///
56/// This can be either a min-heap or a max-heap.
57///
58/// It is a logic error for an item to be modified in such a way that the item's ordering relative
59/// to any other item, as determined by the `Ord` trait, changes while it is in the heap. This is
60/// normally only possible through `Cell`, `RefCell`, global state, I/O, or unsafe code.
61///
62/// ```
63/// use heapless::binary_heap::{BinaryHeap, Max};
64///
65/// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
66///
67/// // We can use peek to look at the next item in the heap. In this case,
68/// // there's no items in there yet so we get None.
69/// assert_eq!(heap.peek(), None);
70///
71/// // Let's add some scores...
72/// heap.push(1).unwrap();
73/// heap.push(5).unwrap();
74/// heap.push(2).unwrap();
75///
76/// // Now peek shows the most important item in the heap.
77/// assert_eq!(heap.peek(), Some(&5));
78///
79/// // We can check the length of a heap.
80/// assert_eq!(heap.len(), 3);
81///
82/// // We can iterate over the items in the heap, although they are returned in
83/// // a random order.
84/// for x in &heap {
85/// println!("{}", x);
86/// }
87///
88/// // If we instead pop these scores, they should come back in order.
89/// assert_eq!(heap.pop(), Some(5));
90/// assert_eq!(heap.pop(), Some(2));
91/// assert_eq!(heap.pop(), Some(1));
92/// assert_eq!(heap.pop(), None);
93///
94/// // We can clear the heap of any remaining items.
95/// heap.clear();
96///
97/// // The heap should now be empty.
98/// assert!(heap.is_empty())
99/// ```
100
101pub struct BinaryHeap<T, K, const N: usize> {
102 pub(crate) _kind: PhantomData<K>,
103 pub(crate) data: Vec<T, N>,
104}
105
106impl<T, K, const N: usize> BinaryHeap<T, K, N> {
107 /* Constructors */
108 /// Creates an empty BinaryHeap as a $K-heap.
109 ///
110 /// ```
111 /// use heapless::binary_heap::{BinaryHeap, Max};
112 ///
113 /// // allocate the binary heap on the stack
114 /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
115 /// heap.push(4).unwrap();
116 ///
117 /// // allocate the binary heap in a static variable
118 /// static mut HEAP: BinaryHeap<i32, Max, 8> = BinaryHeap::new();
119 /// ```
120 pub const fn new() -> Self {
121 Self {
122 _kind: PhantomData,
123 data: Vec::new(),
124 }
125 }
126}
127
128impl<T, K, const N: usize> BinaryHeap<T, K, N>
129where
130 T: Ord,
131 K: Kind,
132{
133 /* Public API */
134 /// Returns the capacity of the binary heap.
135 pub fn capacity(&self) -> usize {
136 self.data.capacity()
137 }
138
139 /// Drops all items from the binary heap.
140 ///
141 /// ```
142 /// use heapless::binary_heap::{BinaryHeap, Max};
143 ///
144 /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
145 /// heap.push(1).unwrap();
146 /// heap.push(3).unwrap();
147 ///
148 /// assert!(!heap.is_empty());
149 ///
150 /// heap.clear();
151 ///
152 /// assert!(heap.is_empty());
153 /// ```
154 pub fn clear(&mut self) {
155 self.data.clear()
156 }
157
158 /// Returns the length of the binary heap.
159 ///
160 /// ```
161 /// use heapless::binary_heap::{BinaryHeap, Max};
162 ///
163 /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
164 /// heap.push(1).unwrap();
165 /// heap.push(3).unwrap();
166 ///
167 /// assert_eq!(heap.len(), 2);
168 /// ```
169 pub fn len(&self) -> usize {
170 self.data.len()
171 }
172
173 /// Checks if the binary heap is empty.
174 ///
175 /// ```
176 /// use heapless::binary_heap::{BinaryHeap, Max};
177 ///
178 /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
179 ///
180 /// assert!(heap.is_empty());
181 ///
182 /// heap.push(3).unwrap();
183 /// heap.push(5).unwrap();
184 /// heap.push(1).unwrap();
185 ///
186 /// assert!(!heap.is_empty());
187 /// ```
188 pub fn is_empty(&self) -> bool {
189 self.len() == 0
190 }
191
192 /// Returns an iterator visiting all values in the underlying vector, in arbitrary order.
193 ///
194 /// ```
195 /// use heapless::binary_heap::{BinaryHeap, Max};
196 ///
197 /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
198 /// heap.push(1).unwrap();
199 /// heap.push(2).unwrap();
200 /// heap.push(3).unwrap();
201 /// heap.push(4).unwrap();
202 ///
203 /// // Print 1, 2, 3, 4 in arbitrary order
204 /// for x in heap.iter() {
205 /// println!("{}", x);
206 ///
207 /// }
208 /// ```
209 pub fn iter(&self) -> slice::Iter<'_, T> {
210 self.data.as_slice().iter()
211 }
212
213 /// Returns a mutable iterator visiting all values in the underlying vector, in arbitrary order.
214 ///
215 /// **WARNING** Mutating the items in the binary heap can leave the heap in an inconsistent
216 /// state.
217 pub fn iter_mut(&mut self) -> slice::IterMut<'_, T> {
218 self.data.as_mut_slice().iter_mut()
219 }
220
221 /// Returns the *top* (greatest if max-heap, smallest if min-heap) item in the binary heap, or
222 /// None if it is empty.
223 ///
224 /// ```
225 /// use heapless::binary_heap::{BinaryHeap, Max};
226 ///
227 /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
228 /// assert_eq!(heap.peek(), None);
229 ///
230 /// heap.push(1).unwrap();
231 /// heap.push(5).unwrap();
232 /// heap.push(2).unwrap();
233 /// assert_eq!(heap.peek(), Some(&5));
234 /// ```
235 pub fn peek(&self) -> Option<&T> {
236 self.data.as_slice().get(0)
237 }
238
239 /// Returns a mutable reference to the greatest item in the binary heap, or
240 /// `None` if it is empty.
241 ///
242 /// Note: If the `PeekMut` value is leaked, the heap may be in an
243 /// inconsistent state.
244 ///
245 /// # Examples
246 ///
247 /// Basic usage:
248 ///
249 /// ```
250 /// use heapless::binary_heap::{BinaryHeap, Max};
251 ///
252 /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
253 /// assert!(heap.peek_mut().is_none());
254 ///
255 /// heap.push(1);
256 /// heap.push(5);
257 /// heap.push(2);
258 /// {
259 /// let mut val = heap.peek_mut().unwrap();
260 /// *val = 0;
261 /// }
262 ///
263 /// assert_eq!(heap.peek(), Some(&2));
264 /// ```
265 pub fn peek_mut(&mut self) -> Option<PeekMut<'_, T, K, N>> {
266 if self.is_empty() {
267 None
268 } else {
269 Some(PeekMut {
270 heap: self,
271 sift: true,
272 })
273 }
274 }
275
276 /// Removes the *top* (greatest if max-heap, smallest if min-heap) item from the binary heap and
277 /// returns it, or None if it is empty.
278 ///
279 /// ```
280 /// use heapless::binary_heap::{BinaryHeap, Max};
281 ///
282 /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
283 /// heap.push(1).unwrap();
284 /// heap.push(3).unwrap();
285 ///
286 /// assert_eq!(heap.pop(), Some(3));
287 /// assert_eq!(heap.pop(), Some(1));
288 /// assert_eq!(heap.pop(), None);
289 /// ```
290 pub fn pop(&mut self) -> Option<T> {
291 if self.is_empty() {
292 None
293 } else {
294 Some(unsafe { self.pop_unchecked() })
295 }
296 }
297
298 /// Removes the *top* (greatest if max-heap, smallest if min-heap) item from the binary heap and
299 /// returns it, without checking if the binary heap is empty.
300 pub unsafe fn pop_unchecked(&mut self) -> T {
301 let mut item = self.data.pop_unchecked();
302
303 if !self.is_empty() {
304 mem::swap(&mut item, self.data.as_mut_slice().get_unchecked_mut(0));
305 self.sift_down_to_bottom(0);
306 }
307 item
308 }
309
310 /// Pushes an item onto the binary heap.
311 ///
312 /// ```
313 /// use heapless::binary_heap::{BinaryHeap, Max};
314 ///
315 /// let mut heap: BinaryHeap<_, Max, 8> = BinaryHeap::new();
316 /// heap.push(3).unwrap();
317 /// heap.push(5).unwrap();
318 /// heap.push(1).unwrap();
319 ///
320 /// assert_eq!(heap.len(), 3);
321 /// assert_eq!(heap.peek(), Some(&5));
322 /// ```
323 pub fn push(&mut self, item: T) -> Result<(), T> {
324 if self.data.is_full() {
325 return Err(item);
326 }
327
328 unsafe { self.push_unchecked(item) }
329 Ok(())
330 }
331
332 /// Pushes an item onto the binary heap without first checking if it's full.
333 pub unsafe fn push_unchecked(&mut self, item: T) {
334 let old_len = self.len();
335 self.data.push_unchecked(item);
336 self.sift_up(0, old_len);
337 }
338
339 /// Returns the underlying ```Vec<T,N>```. Order is arbitrary and time is O(1).
340 pub fn into_vec(self) -> Vec<T, N> {
341 self.data
342 }
343
344 /* Private API */
345 fn sift_down_to_bottom(&mut self, mut pos: usize) {
346 let end = self.len();
347 let start = pos;
348 unsafe {
349 let mut hole = Hole::new(self.data.as_mut_slice(), pos);
350 let mut child = 2 * pos + 1;
351 while child < end {
352 let right = child + 1;
353 // compare with the greater of the two children
354 if right < end && hole.get(child).cmp(hole.get(right)) != K::ordering() {
355 child = right;
356 }
357 hole.move_to(child);
358 child = 2 * hole.pos() + 1;
359 }
360 pos = hole.pos;
361 }
362 self.sift_up(start, pos);
363 }
364
365 fn sift_up(&mut self, start: usize, pos: usize) -> usize {
366 unsafe {
367 // Take out the value at `pos` and create a hole.
368 let mut hole = Hole::new(self.data.as_mut_slice(), pos);
369
370 while hole.pos() > start {
371 let parent = (hole.pos() - 1) / 2;
372 if hole.element().cmp(hole.get(parent)) != K::ordering() {
373 break;
374 }
375 hole.move_to(parent);
376 }
377 hole.pos()
378 }
379 }
380}
381
382/// Hole represents a hole in a slice i.e. an index without valid value
383/// (because it was moved from or duplicated).
384/// In drop, `Hole` will restore the slice by filling the hole
385/// position with the value that was originally removed.
386struct Hole<'a, T> {
387 data: &'a mut [T],
388 /// `elt` is always `Some` from new until drop.
389 elt: ManuallyDrop<T>,
390 pos: usize,
391}
392
393impl<'a, T> Hole<'a, T> {
394 /// Create a new Hole at index `pos`.
395 ///
396 /// Unsafe because pos must be within the data slice.
397 #[inline]
398 unsafe fn new(data: &'a mut [T], pos: usize) -> Self {
399 debug_assert!(pos < data.len());
400 let elt = ptr::read(data.get_unchecked(pos));
401 Hole {
402 data,
403 elt: ManuallyDrop::new(elt),
404 pos,
405 }
406 }
407
408 #[inline]
409 fn pos(&self) -> usize {
410 self.pos
411 }
412
413 /// Returns a reference to the element removed.
414 #[inline]
415 fn element(&self) -> &T {
416 &self.elt
417 }
418
419 /// Returns a reference to the element at `index`.
420 ///
421 /// Unsafe because index must be within the data slice and not equal to pos.
422 #[inline]
423 unsafe fn get(&self, index: usize) -> &T {
424 debug_assert!(index != self.pos);
425 debug_assert!(index < self.data.len());
426 self.data.get_unchecked(index)
427 }
428
429 /// Move hole to new location
430 ///
431 /// Unsafe because index must be within the data slice and not equal to pos.
432 #[inline]
433 unsafe fn move_to(&mut self, index: usize) {
434 debug_assert!(index != self.pos);
435 debug_assert!(index < self.data.len());
436 let ptr = self.data.as_mut_ptr();
437 let index_ptr: *const _ = ptr.add(index);
438 let hole_ptr = ptr.add(self.pos);
439 ptr::copy_nonoverlapping(index_ptr, hole_ptr, 1);
440 self.pos = index;
441 }
442}
443
444/// Structure wrapping a mutable reference to the greatest item on a
445/// `BinaryHeap`.
446///
447/// This `struct` is created by the [`peek_mut`] method on [`BinaryHeap`]. See
448/// its documentation for more.
449///
450/// [`peek_mut`]: struct.BinaryHeap.html#method.peek_mut
451/// [`BinaryHeap`]: struct.BinaryHeap.html
452pub struct PeekMut<'a, T, K, const N: usize>
453where
454 T: Ord,
455 K: Kind,
456{
457 heap: &'a mut BinaryHeap<T, K, N>,
458 sift: bool,
459}
460
461impl<T, K, const N: usize> Drop for PeekMut<'_, T, K, N>
462where
463 T: Ord,
464 K: Kind,
465{
466 fn drop(&mut self) {
467 if self.sift {
468 self.heap.sift_down_to_bottom(0);
469 }
470 }
471}
472
473impl<T, K, const N: usize> Deref for PeekMut<'_, T, K, N>
474where
475 T: Ord,
476 K: Kind,
477{
478 type Target = T;
479 fn deref(&self) -> &T {
480 debug_assert!(!self.heap.is_empty());
481 // SAFE: PeekMut is only instantiated for non-empty heaps
482 unsafe { self.heap.data.as_slice().get_unchecked(0) }
483 }
484}
485
486impl<T, K, const N: usize> DerefMut for PeekMut<'_, T, K, N>
487where
488 T: Ord,
489 K: Kind,
490{
491 fn deref_mut(&mut self) -> &mut T {
492 debug_assert!(!self.heap.is_empty());
493 // SAFE: PeekMut is only instantiated for non-empty heaps
494 unsafe { self.heap.data.as_mut_slice().get_unchecked_mut(0) }
495 }
496}
497
498impl<'a, T, K, const N: usize> PeekMut<'a, T, K, N>
499where
500 T: Ord,
501 K: Kind,
502{
503 /// Removes the peeked value from the heap and returns it.
504 pub fn pop(mut this: PeekMut<'a, T, K, N>) -> T {
505 let value = this.heap.pop().unwrap();
506 this.sift = false;
507 value
508 }
509}
510
511impl<'a, T> Drop for Hole<'a, T> {
512 #[inline]
513 fn drop(&mut self) {
514 // fill the hole again
515 unsafe {
516 let pos = self.pos;
517 ptr::write(self.data.get_unchecked_mut(pos), ptr::read(&*self.elt));
518 }
519 }
520}
521
522impl<T, K, const N: usize> Default for BinaryHeap<T, K, N>
523where
524 T: Ord,
525 K: Kind,
526{
527 fn default() -> Self {
528 Self::new()
529 }
530}
531
532impl<T, K, const N: usize> Clone for BinaryHeap<T, K, N>
533where
534 K: Kind,
535 T: Ord + Clone,
536{
537 fn clone(&self) -> Self {
538 Self {
539 _kind: self._kind,
540 data: self.data.clone(),
541 }
542 }
543}
544
545impl<T, K, const N: usize> fmt::Debug for BinaryHeap<T, K, N>
546where
547 K: Kind,
548 T: Ord + fmt::Debug,
549{
550 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
551 f.debug_list().entries(self.iter()).finish()
552 }
553}
554
555impl<'a, T, K, const N: usize> IntoIterator for &'a BinaryHeap<T, K, N>
556where
557 K: Kind,
558 T: Ord,
559{
560 type Item = &'a T;
561 type IntoIter = slice::Iter<'a, T>;
562
563 fn into_iter(self) -> Self::IntoIter {
564 self.iter()
565 }
566}
567
568#[cfg(test)]
569mod tests {
570 use std::vec::Vec;
571
572 use crate::binary_heap::{BinaryHeap, Max, Min};
573
574 #[test]
575 fn static_new() {
576 static mut _B: BinaryHeap<i32, Min, 16> = BinaryHeap::new();
577 }
578
579 #[test]
580 fn drop() {
581 droppable!();
582
583 {
584 let mut v: BinaryHeap<Droppable, Max, 2> = BinaryHeap::new();
585 v.push(Droppable::new()).ok().unwrap();
586 v.push(Droppable::new()).ok().unwrap();
587 v.pop().unwrap();
588 }
589
590 assert_eq!(Droppable::count(), 0);
591
592 {
593 let mut v: BinaryHeap<Droppable, Max, 2> = BinaryHeap::new();
594 v.push(Droppable::new()).ok().unwrap();
595 v.push(Droppable::new()).ok().unwrap();
596 }
597
598 assert_eq!(Droppable::count(), 0);
599
600 {
601 let mut v: BinaryHeap<Droppable, Min, 2> = BinaryHeap::new();
602 v.push(Droppable::new()).ok().unwrap();
603 v.push(Droppable::new()).ok().unwrap();
604 v.pop().unwrap();
605 }
606
607 assert_eq!(Droppable::count(), 0);
608
609 {
610 let mut v: BinaryHeap<Droppable, Min, 2> = BinaryHeap::new();
611 v.push(Droppable::new()).ok().unwrap();
612 v.push(Droppable::new()).ok().unwrap();
613 }
614
615 assert_eq!(Droppable::count(), 0);
616 }
617
618 #[test]
619 fn into_vec() {
620 droppable!();
621
622 let mut h: BinaryHeap<Droppable, Max, 2> = BinaryHeap::new();
623 h.push(Droppable::new()).ok().unwrap();
624 h.push(Droppable::new()).ok().unwrap();
625 h.pop().unwrap();
626
627 assert_eq!(Droppable::count(), 1);
628
629 let v = h.into_vec();
630
631 assert_eq!(Droppable::count(), 1);
632
633 core::mem::drop(v);
634
635 assert_eq!(Droppable::count(), 0);
636 }
637
638 #[test]
639 fn min() {
640 let mut heap = BinaryHeap::<_, Min, 16>::new();
641 heap.push(1).unwrap();
642 heap.push(2).unwrap();
643 heap.push(3).unwrap();
644 heap.push(17).unwrap();
645 heap.push(19).unwrap();
646 heap.push(36).unwrap();
647 heap.push(7).unwrap();
648 heap.push(25).unwrap();
649 heap.push(100).unwrap();
650
651 assert_eq!(
652 heap.iter().cloned().collect::<Vec<_>>(),
653 [1, 2, 3, 17, 19, 36, 7, 25, 100]
654 );
655
656 assert_eq!(heap.pop(), Some(1));
657
658 assert_eq!(
659 heap.iter().cloned().collect::<Vec<_>>(),
660 [2, 17, 3, 25, 19, 36, 7, 100]
661 );
662
663 assert_eq!(heap.pop(), Some(2));
664 assert_eq!(heap.pop(), Some(3));
665 assert_eq!(heap.pop(), Some(7));
666 assert_eq!(heap.pop(), Some(17));
667 assert_eq!(heap.pop(), Some(19));
668 assert_eq!(heap.pop(), Some(25));
669 assert_eq!(heap.pop(), Some(36));
670 assert_eq!(heap.pop(), Some(100));
671 assert_eq!(heap.pop(), None);
672
673 assert!(heap.peek_mut().is_none());
674
675 heap.push(1).unwrap();
676 heap.push(2).unwrap();
677 heap.push(10).unwrap();
678
679 {
680 let mut val = heap.peek_mut().unwrap();
681 *val = 7;
682 }
683
684 assert_eq!(heap.pop(), Some(2));
685 assert_eq!(heap.pop(), Some(7));
686 assert_eq!(heap.pop(), Some(10));
687 assert_eq!(heap.pop(), None);
688 }
689
690 #[test]
691 fn max() {
692 let mut heap = BinaryHeap::<_, Max, 16>::new();
693 heap.push(1).unwrap();
694 heap.push(2).unwrap();
695 heap.push(3).unwrap();
696 heap.push(17).unwrap();
697 heap.push(19).unwrap();
698 heap.push(36).unwrap();
699 heap.push(7).unwrap();
700 heap.push(25).unwrap();
701 heap.push(100).unwrap();
702
703 assert_eq!(
704 heap.iter().cloned().collect::<Vec<_>>(),
705 [100, 36, 19, 25, 3, 2, 7, 1, 17]
706 );
707
708 assert_eq!(heap.pop(), Some(100));
709
710 assert_eq!(
711 heap.iter().cloned().collect::<Vec<_>>(),
712 [36, 25, 19, 17, 3, 2, 7, 1]
713 );
714
715 assert_eq!(heap.pop(), Some(36));
716 assert_eq!(heap.pop(), Some(25));
717 assert_eq!(heap.pop(), Some(19));
718 assert_eq!(heap.pop(), Some(17));
719 assert_eq!(heap.pop(), Some(7));
720 assert_eq!(heap.pop(), Some(3));
721 assert_eq!(heap.pop(), Some(2));
722 assert_eq!(heap.pop(), Some(1));
723 assert_eq!(heap.pop(), None);
724
725 assert!(heap.peek_mut().is_none());
726
727 heap.push(1).unwrap();
728 heap.push(9).unwrap();
729 heap.push(10).unwrap();
730
731 {
732 let mut val = heap.peek_mut().unwrap();
733 *val = 7;
734 }
735
736 assert_eq!(heap.pop(), Some(9));
737 assert_eq!(heap.pop(), Some(7));
738 assert_eq!(heap.pop(), Some(1));
739 assert_eq!(heap.pop(), None);
740 }
741}