logo
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
//! Binary heap with a contextful comparator and position tracking
//!
//! The implementation is mostly based on the Rust standard library's
//! `BinaryHeap`.
use core::marker::Destruct;

mod helpers;
#[cfg(test)]
mod tests;
mod veclike;
pub use self::veclike::*;

/// Context type for [`BinaryHeap`]'s operations.
pub trait BinaryHeapCtx<Element> {
    /// Return `true` iff `x < y`.
    fn lt(&mut self, x: &Element, y: &Element) -> bool;

    /// Called when the element `e` is moved to the new position `new_index`.
    #[default_method_body_is_const]
    fn on_move(&mut self, e: &mut Element, new_index: usize) {
        let _ = (e, new_index);
    }
}

impl<T: ~const Ord + ~const PartialOrd> const BinaryHeapCtx<T> for () {
    fn lt(&mut self, x: &T, y: &T) -> bool {
        *x < *y
    }
}

/// Min-heap.
pub trait BinaryHeap: VecLike {
    /// Remove the least item from the heap and return it.
    fn heap_pop<Ctx>(&mut self, ctx: Ctx) -> Option<Self::Element>
    where
        Ctx: ~const BinaryHeapCtx<Self::Element> + ~const Destruct;

    /// Remove the item at the specified position and return it.
    fn heap_remove<Ctx>(&mut self, i: usize, ctx: Ctx) -> Option<Self::Element>
    where
        Ctx: ~const BinaryHeapCtx<Self::Element> + ~const Destruct;

    /// Push an item onto the heap and return its position.
    fn heap_push<Ctx>(&mut self, item: Self::Element, ctx: Ctx) -> usize
    where
        Ctx: ~const BinaryHeapCtx<Self::Element> + ~const Destruct;
}

impl<T> const BinaryHeap for T
where
    // `~const Deref` isn't implied because of
    // [ref:veclike_const_supertrait]
    T: ~const VecLike + ~const core::ops::Deref + ~const core::ops::DerefMut,
    T::Element: ~const Destruct,
{
    fn heap_pop<Ctx>(&mut self, ctx: Ctx) -> Option<Self::Element>
    where
        Ctx: ~const BinaryHeapCtx<Self::Element> + ~const Destruct,
    {
        self.heap_remove(0, ctx)
    }

    fn heap_remove<Ctx>(&mut self, i: usize, mut ctx: Ctx) -> Option<Self::Element>
    where
        Ctx: ~const BinaryHeapCtx<Self::Element> + ~const Destruct,
    {
        if i >= self.len() {
            return None;
        }

        if let Some(mut item) = self.pop() {
            let slice = &mut **self;
            if i < slice.len() {
                // Swap the last item with the item at `i`
                core::mem::swap(&mut slice[i], &mut item);
                ctx.on_move(&mut slice[i], i);

                let should_sift_up = i > 0 && ctx.lt(&slice[i], &slice[(i - 1) / 2]);

                // Sift down or up the item at `i`, restoring the invariant
                // Safety: `i` points to an element within `slice`.
                unsafe {
                    if should_sift_up {
                        sift_up(slice, 0, i, ctx);
                    } else {
                        sift_down(slice, i, ctx);
                    }
                }
            }
            Some(item)
        } else {
            debug_assert!(false);
            None
        }
    }

    fn heap_push<Ctx>(&mut self, item: Self::Element, ctx: Ctx) -> usize
    where
        Ctx: ~const BinaryHeapCtx<Self::Element> + ~const Destruct,
    {
        let i = self.len();
        self.push(item);

        let slice = &mut **self;
        assert!(i < slice.len());

        // Safety: `i` points to an element within `slice`.
        unsafe { sift_up(slice, 0, i, ctx) }
    }
}

// The implementations of sift_up and sift_down use unsafe blocks in
// order to move an element out of the vector (leaving behind a
// hole), shift along the others and move the removed element back into the
// vector at the final location of the hole.
// The `Hole` type is used to represent this, and make sure
// the hole is filled back at the end of its scope, even on panic.
// Using a hole reduces the constant factor compared to using swaps,
// which involves twice as many moves.
/// Sift-up operation
///
/// # Safety
///
/// `pos` must point to an element within `this`.
const unsafe fn sift_up<Element, Ctx>(
    this: &mut [Element],
    start: usize,
    pos: usize,
    mut ctx: Ctx,
) -> usize
where
    Ctx: ~const BinaryHeapCtx<Element> + ~const Destruct,
    Element: ~const Destruct,
{
    unsafe {
        // Take out the value at `pos` and create a hole.
        let mut hole = helpers::Hole::new(this, pos);

        while hole.pos() > start {
            let parent = (hole.pos() - 1) / 2;
            if !ctx.lt(hole.element(), hole.get(parent)) {
                break;
            }

            let prev_pos = hole.pos();
            hole.move_to(parent);

            // `[prev_pos]` is now filled with the element moved from `[parent]`
            ctx.on_move(hole.get_mut(prev_pos), prev_pos);
        }

        // Report the final position of the newly-inserted element
        let pos = hole.pos();
        ctx.on_move(hole.element_mut(), pos);

        pos
    }
}

/// Take an element at `pos` and move it down the heap,
/// while its children are larger.
///
/// # Safety
///
/// `pos` must point to an element within `this`.
const unsafe fn sift_down<Element, Ctx>(this: &mut [Element], pos: usize, mut ctx: Ctx)
where
    Ctx: ~const BinaryHeapCtx<Element> + ~const Destruct,
    Element: ~const Destruct,
{
    let end = this.len();
    unsafe {
        let mut hole = helpers::Hole::new(this, pos);
        let mut child = 2 * pos + 1;
        while child < end {
            let right = child + 1;
            // compare with the lesser of the two children
            if right < end && !ctx.lt(hole.get(child), hole.get(right)) {
                child = right;
            }

            // if we are already in order, stop.
            if !ctx.lt(hole.get(child), hole.element()) {
                break;
            }

            let prev_pos = hole.pos();
            hole.move_to(child);

            // `[prev_pos]` is now filled with the element moved from `[child]`
            ctx.on_move(hole.get_mut(prev_pos), prev_pos);

            child = 2 * hole.pos() + 1;
        }

        // Report the final position of `hole.element_mut()`
        let pos = hole.pos();
        ctx.on_move(hole.element_mut(), pos);
    }
}