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 [`BinaryHeap::peek_mut`].
448/// See its documentation for more.
449pub struct PeekMut<'a, T, K, const N: usize>
450where
451 T: Ord,
452 K: Kind,
453{
454 heap: &'a mut BinaryHeap<T, K, N>,
455 sift: bool,
456}
457
458impl<T, K, const N: usize> Drop for PeekMut<'_, T, K, N>
459where
460 T: Ord,
461 K: Kind,
462{
463 fn drop(&mut self) {
464 if self.sift {
465 self.heap.sift_down_to_bottom(pos:0);
466 }
467 }
468}
469
470impl<T, K, const N: usize> Deref for PeekMut<'_, T, K, N>
471where
472 T: Ord,
473 K: Kind,
474{
475 type Target = T;
476 fn deref(&self) -> &T {
477 debug_assert!(!self.heap.is_empty());
478 // SAFE: PeekMut is only instantiated for non-empty heaps
479 unsafe { self.heap.data.as_slice().get_unchecked(index:0) }
480 }
481}
482
483impl<T, K, const N: usize> DerefMut for PeekMut<'_, T, K, N>
484where
485 T: Ord,
486 K: Kind,
487{
488 fn deref_mut(&mut self) -> &mut T {
489 debug_assert!(!self.heap.is_empty());
490 // SAFE: PeekMut is only instantiated for non-empty heaps
491 unsafe { self.heap.data.as_mut_slice().get_unchecked_mut(index:0) }
492 }
493}
494
495impl<'a, T, K, const N: usize> PeekMut<'a, T, K, N>
496where
497 T: Ord,
498 K: Kind,
499{
500 /// Removes the peeked value from the heap and returns it.
501 pub fn pop(mut this: PeekMut<'a, T, K, N>) -> T {
502 let value: T = this.heap.pop().unwrap();
503 this.sift = false;
504 value
505 }
506}
507
508impl<'a, T> Drop for Hole<'a, T> {
509 #[inline]
510 fn drop(&mut self) {
511 // fill the hole again
512 unsafe {
513 let pos: usize = self.pos;
514 ptr::write(self.data.get_unchecked_mut(pos), src:ptr::read(&*self.elt));
515 }
516 }
517}
518
519impl<T, K, const N: usize> Default for BinaryHeap<T, K, N>
520where
521 T: Ord,
522 K: Kind,
523{
524 fn default() -> Self {
525 Self::new()
526 }
527}
528
529impl<T, K, const N: usize> Clone for BinaryHeap<T, K, N>
530where
531 K: Kind,
532 T: Ord + Clone,
533{
534 fn clone(&self) -> Self {
535 Self {
536 _kind: self._kind,
537 data: self.data.clone(),
538 }
539 }
540}
541
542impl<T, K, const N: usize> fmt::Debug for BinaryHeap<T, K, N>
543where
544 K: Kind,
545 T: Ord + fmt::Debug,
546{
547 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
548 f.debug_list().entries(self.iter()).finish()
549 }
550}
551
552impl<'a, T, K, const N: usize> IntoIterator for &'a BinaryHeap<T, K, N>
553where
554 K: Kind,
555 T: Ord,
556{
557 type Item = &'a T;
558 type IntoIter = slice::Iter<'a, T>;
559
560 fn into_iter(self) -> Self::IntoIter {
561 self.iter()
562 }
563}
564
565#[cfg(test)]
566mod tests {
567 use std::vec::Vec;
568
569 use crate::binary_heap::{BinaryHeap, Max, Min};
570
571 #[test]
572 fn static_new() {
573 static mut _B: BinaryHeap<i32, Min, 16> = BinaryHeap::new();
574 }
575
576 #[test]
577 fn drop() {
578 droppable!();
579
580 {
581 let mut v: BinaryHeap<Droppable, Max, 2> = BinaryHeap::new();
582 v.push(Droppable::new()).ok().unwrap();
583 v.push(Droppable::new()).ok().unwrap();
584 v.pop().unwrap();
585 }
586
587 assert_eq!(Droppable::count(), 0);
588
589 {
590 let mut v: BinaryHeap<Droppable, Max, 2> = BinaryHeap::new();
591 v.push(Droppable::new()).ok().unwrap();
592 v.push(Droppable::new()).ok().unwrap();
593 }
594
595 assert_eq!(Droppable::count(), 0);
596
597 {
598 let mut v: BinaryHeap<Droppable, Min, 2> = BinaryHeap::new();
599 v.push(Droppable::new()).ok().unwrap();
600 v.push(Droppable::new()).ok().unwrap();
601 v.pop().unwrap();
602 }
603
604 assert_eq!(Droppable::count(), 0);
605
606 {
607 let mut v: BinaryHeap<Droppable, Min, 2> = BinaryHeap::new();
608 v.push(Droppable::new()).ok().unwrap();
609 v.push(Droppable::new()).ok().unwrap();
610 }
611
612 assert_eq!(Droppable::count(), 0);
613 }
614
615 #[test]
616 fn into_vec() {
617 droppable!();
618
619 let mut h: BinaryHeap<Droppable, Max, 2> = BinaryHeap::new();
620 h.push(Droppable::new()).ok().unwrap();
621 h.push(Droppable::new()).ok().unwrap();
622 h.pop().unwrap();
623
624 assert_eq!(Droppable::count(), 1);
625
626 let v = h.into_vec();
627
628 assert_eq!(Droppable::count(), 1);
629
630 core::mem::drop(v);
631
632 assert_eq!(Droppable::count(), 0);
633 }
634
635 #[test]
636 fn min() {
637 let mut heap = BinaryHeap::<_, Min, 16>::new();
638 heap.push(1).unwrap();
639 heap.push(2).unwrap();
640 heap.push(3).unwrap();
641 heap.push(17).unwrap();
642 heap.push(19).unwrap();
643 heap.push(36).unwrap();
644 heap.push(7).unwrap();
645 heap.push(25).unwrap();
646 heap.push(100).unwrap();
647
648 assert_eq!(
649 heap.iter().cloned().collect::<Vec<_>>(),
650 [1, 2, 3, 17, 19, 36, 7, 25, 100]
651 );
652
653 assert_eq!(heap.pop(), Some(1));
654
655 assert_eq!(
656 heap.iter().cloned().collect::<Vec<_>>(),
657 [2, 17, 3, 25, 19, 36, 7, 100]
658 );
659
660 assert_eq!(heap.pop(), Some(2));
661 assert_eq!(heap.pop(), Some(3));
662 assert_eq!(heap.pop(), Some(7));
663 assert_eq!(heap.pop(), Some(17));
664 assert_eq!(heap.pop(), Some(19));
665 assert_eq!(heap.pop(), Some(25));
666 assert_eq!(heap.pop(), Some(36));
667 assert_eq!(heap.pop(), Some(100));
668 assert_eq!(heap.pop(), None);
669
670 assert!(heap.peek_mut().is_none());
671
672 heap.push(1).unwrap();
673 heap.push(2).unwrap();
674 heap.push(10).unwrap();
675
676 {
677 let mut val = heap.peek_mut().unwrap();
678 *val = 7;
679 }
680
681 assert_eq!(heap.pop(), Some(2));
682 assert_eq!(heap.pop(), Some(7));
683 assert_eq!(heap.pop(), Some(10));
684 assert_eq!(heap.pop(), None);
685 }
686
687 #[test]
688 fn max() {
689 let mut heap = BinaryHeap::<_, Max, 16>::new();
690 heap.push(1).unwrap();
691 heap.push(2).unwrap();
692 heap.push(3).unwrap();
693 heap.push(17).unwrap();
694 heap.push(19).unwrap();
695 heap.push(36).unwrap();
696 heap.push(7).unwrap();
697 heap.push(25).unwrap();
698 heap.push(100).unwrap();
699
700 assert_eq!(
701 heap.iter().cloned().collect::<Vec<_>>(),
702 [100, 36, 19, 25, 3, 2, 7, 1, 17]
703 );
704
705 assert_eq!(heap.pop(), Some(100));
706
707 assert_eq!(
708 heap.iter().cloned().collect::<Vec<_>>(),
709 [36, 25, 19, 17, 3, 2, 7, 1]
710 );
711
712 assert_eq!(heap.pop(), Some(36));
713 assert_eq!(heap.pop(), Some(25));
714 assert_eq!(heap.pop(), Some(19));
715 assert_eq!(heap.pop(), Some(17));
716 assert_eq!(heap.pop(), Some(7));
717 assert_eq!(heap.pop(), Some(3));
718 assert_eq!(heap.pop(), Some(2));
719 assert_eq!(heap.pop(), Some(1));
720 assert_eq!(heap.pop(), None);
721
722 assert!(heap.peek_mut().is_none());
723
724 heap.push(1).unwrap();
725 heap.push(9).unwrap();
726 heap.push(10).unwrap();
727
728 {
729 let mut val = heap.peek_mut().unwrap();
730 *val = 7;
731 }
732
733 assert_eq!(heap.pop(), Some(9));
734 assert_eq!(heap.pop(), Some(7));
735 assert_eq!(heap.pop(), Some(1));
736 assert_eq!(heap.pop(), None);
737 }
738}
739