nautilus_persistence/backend/
binary_heap.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! A priority queue implemented with a binary heap.
17//!
18//! Vendored from the `binary-heap-plus` crate which depends on the unmaintained
19//! `compare` crate. Distributed here under MIT (see `licenses/` directory).
20//! Original source: <https://github.com/sekineh/binary-heap-plus-rs>
21
22#![deny(unsafe_op_in_unsafe_fn)]
23#![allow(
24    clippy::multiple_unsafe_ops_per_block,
25    reason = "vendored heap chains pointer ops; SAFETY comments justify each block"
26)]
27
28use std::{
29    fmt,
30    mem::{ManuallyDrop, swap},
31    ops::{Deref, DerefMut},
32    ptr,
33};
34
35use super::compare::Compare;
36
37/// A priority queue implemented with a binary heap.
38///
39/// This will be a max-heap by default, but the ordering is determined by the
40/// comparator `C`.
41pub struct BinaryHeap<T, C> {
42    data: Vec<T>,
43    cmp: C,
44}
45
46impl<T, C: Compare<T>> BinaryHeap<T, C> {
47    /// Creates a `BinaryHeap` from a `Vec` and comparator.
48    pub fn from_vec_cmp(vec: Vec<T>, cmp: C) -> Self {
49        let mut heap = Self { data: vec, cmp };
50        if !heap.data.is_empty() {
51            heap.rebuild();
52        }
53        heap
54    }
55
56    /// Returns a mutable reference to the greatest item in the binary heap, or
57    /// `None` if it is empty.
58    pub fn peek_mut(&mut self) -> Option<PeekMut<'_, T, C>> {
59        if self.is_empty() {
60            None
61        } else {
62            Some(PeekMut {
63                heap: self,
64                sift: false,
65            })
66        }
67    }
68
69    /// Removes the greatest item from the binary heap and returns it, or `None`
70    /// if it is empty.
71    pub fn pop(&mut self) -> Option<T> {
72        self.data.pop().map(|mut item| {
73            if !self.is_empty() {
74                swap(&mut item, &mut self.data[0]);
75                // SAFETY: !self.is_empty() means that self.len() > 0
76                unsafe { self.sift_down_to_bottom(0) };
77            }
78            item
79        })
80    }
81
82    /// Pushes an item onto the binary heap.
83    pub fn push(&mut self, item: T) {
84        let old_len = self.len();
85        self.data.push(item);
86        // SAFETY: Since we pushed a new item it means that
87        //  old_len = self.len() - 1 < self.len()
88        unsafe { self.sift_up(0, old_len) };
89    }
90
91    /// Sifts an element up towards the root until heap property is restored.
92    ///
93    /// # Safety
94    ///
95    /// The caller must guarantee that `pos < self.len()`.
96    unsafe fn sift_up(&mut self, start: usize, pos: usize) -> usize {
97        // SAFETY: The caller guarantees that pos < self.len()
98        let mut hole = unsafe { Hole::new(&mut self.data, pos) };
99
100        while hole.pos() > start {
101            let parent = (hole.pos() - 1) / 2;
102
103            // SAFETY: hole.pos() > start >= 0, which means hole.pos() > 0
104            //  and so hole.pos() - 1 can't underflow.
105            if self
106                .cmp
107                .compares_le(hole.element(), unsafe { hole.get(parent) })
108            {
109                break;
110            }
111
112            // SAFETY: Same as above
113            unsafe { hole.move_to(parent) };
114        }
115
116        hole.pos()
117    }
118
119    /// Sifts an element down within the range `[pos, end)`.
120    ///
121    /// # Safety
122    ///
123    /// The caller must guarantee that `pos < end <= self.len()`.
124    unsafe fn sift_down_range(&mut self, pos: usize, end: usize) {
125        // SAFETY: The caller guarantees that pos < end <= self.len().
126        let mut hole = unsafe { Hole::new(&mut self.data, pos) };
127        let mut child = 2 * hole.pos() + 1;
128
129        while child <= end.saturating_sub(2) {
130            // SAFETY: child < end - 1 < self.len() and
131            //  child + 1 < end <= self.len(), so they're valid indexes.
132            child +=
133                usize::from(unsafe { self.cmp.compares_le(hole.get(child), hole.get(child + 1)) });
134
135            // SAFETY: child is now either the old child or the old child+1
136            if self
137                .cmp
138                .compares_ge(hole.element(), unsafe { hole.get(child) })
139            {
140                return;
141            }
142
143            // SAFETY: same as above.
144            unsafe { hole.move_to(child) };
145            child = 2 * hole.pos() + 1;
146        }
147
148        // SAFETY: && short circuit
149        if child == end - 1
150            && self
151                .cmp
152                .compares_lt(hole.element(), unsafe { hole.get(child) })
153        {
154            unsafe { hole.move_to(child) };
155        }
156    }
157
158    /// Sifts an element down until heap property is restored.
159    ///
160    /// # Safety
161    ///
162    /// The caller must guarantee that `pos < self.len()`.
163    unsafe fn sift_down(&mut self, pos: usize) {
164        let len = self.len();
165        // SAFETY: pos < len is guaranteed by the caller
166        unsafe { self.sift_down_range(pos, len) };
167    }
168
169    /// Take an element at `pos` and move it all the way down the heap,
170    /// then sift it up to its position.
171    ///
172    /// # Safety
173    ///
174    /// The caller must guarantee that `pos < self.len()`.
175    unsafe fn sift_down_to_bottom(&mut self, mut pos: usize) {
176        let end = self.len();
177        let start = pos;
178
179        // SAFETY: The caller guarantees that pos < self.len().
180        let mut hole = unsafe { Hole::new(&mut self.data, pos) };
181        let mut child = 2 * hole.pos() + 1;
182
183        while child <= end.saturating_sub(2) {
184            // SAFETY: child < end - 1 < self.len() and
185            //  child + 1 < end <= self.len(), so they're valid indexes.
186            child +=
187                usize::from(unsafe { self.cmp.compares_le(hole.get(child), hole.get(child + 1)) });
188
189            // SAFETY: Same as above
190            unsafe { hole.move_to(child) };
191            child = 2 * hole.pos() + 1;
192        }
193
194        if child == end - 1 {
195            // SAFETY: child == end - 1 < self.len(), so it's a valid index
196            unsafe { hole.move_to(child) };
197        }
198        pos = hole.pos();
199        drop(hole);
200
201        // SAFETY: pos is the position in the hole
202        unsafe { self.sift_up(start, pos) };
203    }
204
205    /// Rebuilds the heap from an unordered vector.
206    fn rebuild(&mut self) {
207        let mut n = self.len() / 2;
208        while n > 0 {
209            n -= 1;
210            // SAFETY: n starts from self.len() / 2 and goes down to 0.
211            unsafe { self.sift_down(n) };
212        }
213    }
214}
215
216impl<T, C> BinaryHeap<T, C> {
217    /// Returns the length of the binary heap.
218    #[must_use]
219    pub fn len(&self) -> usize {
220        self.data.len()
221    }
222
223    /// Checks if the binary heap is empty.
224    #[must_use]
225    pub fn is_empty(&self) -> bool {
226        self.len() == 0
227    }
228
229    /// Drops all items from the binary heap.
230    pub fn clear(&mut self) {
231        self.data.clear();
232    }
233}
234
235impl<T: fmt::Debug, C> fmt::Debug for BinaryHeap<T, C> {
236    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
237        f.debug_list().entries(self.data.iter()).finish()
238    }
239}
240
241impl<T: Clone, C: Clone> Clone for BinaryHeap<T, C> {
242    fn clone(&self) -> Self {
243        Self {
244            data: self.data.clone(),
245            cmp: self.cmp.clone(),
246        }
247    }
248}
249
250/// Structure wrapping a mutable reference to the greatest item on a
251/// `BinaryHeap`.
252pub struct PeekMut<'a, T: 'a, C: 'a + Compare<T>> {
253    heap: &'a mut BinaryHeap<T, C>,
254    sift: bool,
255}
256
257impl<T: fmt::Debug, C: Compare<T>> fmt::Debug for PeekMut<'_, T, C> {
258    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
259        f.debug_tuple("PeekMut").field(&self.heap.data[0]).finish()
260    }
261}
262
263impl<T, C: Compare<T>> Drop for PeekMut<'_, T, C> {
264    fn drop(&mut self) {
265        if self.sift {
266            // SAFETY: PeekMut is only instantiated for non-empty heaps.
267            unsafe { self.heap.sift_down(0) };
268        }
269    }
270}
271
272impl<T, C: Compare<T>> Deref for PeekMut<'_, T, C> {
273    type Target = T;
274    fn deref(&self) -> &T {
275        debug_assert!(!self.heap.is_empty());
276        // SAFETY: PeekMut is only instantiated for non-empty heaps
277        unsafe { self.heap.data.get_unchecked(0) }
278    }
279}
280
281impl<T, C: Compare<T>> DerefMut for PeekMut<'_, T, C> {
282    fn deref_mut(&mut self) -> &mut T {
283        debug_assert!(!self.heap.is_empty());
284        self.sift = true;
285        // SAFETY: PeekMut is only instantiated for non-empty heaps
286        unsafe { self.heap.data.get_unchecked_mut(0) }
287    }
288}
289
290impl<T, C: Compare<T>> PeekMut<'_, T, C> {
291    /// Removes the peeked value from the heap and returns it.
292    ///
293    /// # Panics
294    ///
295    /// Panics if the wrapped heap is empty. Safe constructors only create [`PeekMut`]
296    /// for non-empty heaps.
297    #[must_use]
298    pub fn pop(mut this: Self) -> T {
299        let value = this.heap.pop().unwrap();
300        this.sift = false;
301        value
302    }
303}
304
305/// Hole represents a hole in a slice i.e., an index without valid value
306/// (because it was moved from or duplicated).
307struct Hole<'a, T: 'a> {
308    data: &'a mut [T],
309    elt: ManuallyDrop<T>,
310    pos: usize,
311}
312
313impl<'a, T> Hole<'a, T> {
314    /// Create a new `Hole` at index `pos`.
315    ///
316    /// # Safety
317    ///
318    /// `pos` must be within the data slice.
319    #[inline]
320    unsafe fn new(data: &'a mut [T], pos: usize) -> Self {
321        debug_assert!(pos < data.len());
322        // SAFETY: pos should be inside the slice
323        let elt = unsafe { ptr::read(data.get_unchecked(pos)) };
324        Hole {
325            data,
326            elt: ManuallyDrop::new(elt),
327            pos,
328        }
329    }
330
331    #[inline]
332    fn pos(&self) -> usize {
333        self.pos
334    }
335
336    /// Returns a reference to the element removed.
337    #[inline]
338    fn element(&self) -> &T {
339        &self.elt
340    }
341
342    /// Returns a reference to the element at `index`.
343    ///
344    /// # Safety
345    ///
346    /// `index` must be within the data slice and not equal to pos.
347    #[inline]
348    unsafe fn get(&self, index: usize) -> &T {
349        debug_assert_ne!(index, self.pos);
350        debug_assert!(index < self.data.len());
351        unsafe { self.data.get_unchecked(index) }
352    }
353
354    /// Move hole to new location.
355    ///
356    /// # Safety
357    ///
358    /// `index` must be within the data slice and not equal to pos.
359    #[inline]
360    unsafe fn move_to(&mut self, index: usize) {
361        debug_assert_ne!(index, self.pos);
362        debug_assert!(index < self.data.len());
363        // SAFETY: `index` and `pos` are bounds-checked by the debug assertions above.
364        unsafe {
365            let ptr = self.data.as_mut_ptr();
366            let index_ptr: *const _ = ptr.add(index);
367            let hole_ptr = ptr.add(self.pos);
368            ptr::copy_nonoverlapping(index_ptr, hole_ptr, 1);
369        }
370        self.pos = index;
371    }
372}
373
374impl<T> Drop for Hole<'_, T> {
375    #[inline]
376    fn drop(&mut self) {
377        // SAFETY: `pos` was valid when the hole was created and all moves
378        // maintain the invariant that `pos < data.len()`.
379        unsafe {
380            let pos = self.pos;
381            ptr::copy_nonoverlapping(
382                ptr::from_ref(&*self.elt),
383                self.data.get_unchecked_mut(pos),
384                1,
385            );
386        }
387    }
388}
389
390#[cfg(test)]
391mod tests {
392    use std::cmp::Ordering;
393
394    use rstest::rstest;
395
396    use super::*;
397
398    struct MaxComparator;
399
400    impl Compare<i32> for MaxComparator {
401        fn compare(&self, a: &i32, b: &i32) -> Ordering {
402            a.cmp(b)
403        }
404    }
405
406    struct MinComparator;
407
408    impl Compare<i32> for MinComparator {
409        fn compare(&self, a: &i32, b: &i32) -> Ordering {
410            b.cmp(a)
411        }
412    }
413
414    #[rstest]
415    fn test_max_heap() {
416        let mut heap = BinaryHeap::from_vec_cmp(vec![], MaxComparator);
417        heap.push(3);
418        heap.push(1);
419        heap.push(5);
420
421        assert_eq!(heap.pop(), Some(5));
422        assert_eq!(heap.pop(), Some(3));
423        assert_eq!(heap.pop(), Some(1));
424        assert_eq!(heap.pop(), None);
425    }
426
427    #[rstest]
428    fn test_min_heap() {
429        let mut heap = BinaryHeap::from_vec_cmp(vec![], MinComparator);
430        heap.push(3);
431        heap.push(1);
432        heap.push(5);
433
434        assert_eq!(heap.pop(), Some(1));
435        assert_eq!(heap.pop(), Some(3));
436        assert_eq!(heap.pop(), Some(5));
437        assert_eq!(heap.pop(), None);
438    }
439
440    #[rstest]
441    fn test_peek_mut() {
442        let mut heap = BinaryHeap::from_vec_cmp(vec![1, 5, 2], MaxComparator);
443
444        if let Some(mut val) = heap.peek_mut() {
445            *val = 0;
446        }
447
448        assert_eq!(heap.pop(), Some(2));
449    }
450
451    #[rstest]
452    fn test_peek_mut_pop() {
453        let mut heap = BinaryHeap::from_vec_cmp(vec![1, 5, 2], MaxComparator);
454
455        if let Some(val) = heap.peek_mut() {
456            let popped = PeekMut::pop(val);
457            assert_eq!(popped, 5);
458        }
459
460        assert_eq!(heap.pop(), Some(2));
461        assert_eq!(heap.pop(), Some(1));
462    }
463
464    #[rstest]
465    fn test_clear() {
466        let mut heap = BinaryHeap::from_vec_cmp(vec![1, 2, 3], MaxComparator);
467        assert!(!heap.is_empty());
468
469        heap.clear();
470        assert!(heap.is_empty());
471        assert_eq!(heap.len(), 0);
472    }
473
474    #[rstest]
475    fn test_from_vec() {
476        let heap = BinaryHeap::from_vec_cmp(vec![3, 1, 4, 1, 5, 9, 2, 6], MaxComparator);
477        let mut sorted = Vec::new();
478        let mut heap = heap;
479        while let Some(v) = heap.pop() {
480            sorted.push(v);
481        }
482        assert_eq!(sorted, vec![9, 6, 5, 4, 3, 2, 1, 1]);
483    }
484}
nautilus_persistence/backend/binary_heap.rs

nautilus_persistence/backend/
binary_heap.rs
