futures_timer/native/
heap.rs

Help
1//! A simple binary heap with support for removal of arbitrary elements
2//!
3//! This heap is used to manage timer state in the event loop. All timeouts go
4//! into this heap and we also cancel timeouts from this heap. The crucial
5//! feature of this heap over the standard library's `BinaryHeap` is the ability
6//! to remove arbitrary elements. (e.g. when a timer is canceled)
7//!
8//! Note that this heap is not at all optimized right now, it should hopefully
9//! just work.
10
11use std::mem;
12
13pub struct Heap<T> {
14    // Binary heap of items, plus the slab index indicating what position in the
15    // list they're in.
16    items: Vec<(T, usize)>,
17
18    // A map from a slab index (assigned to an item above) to the actual index
19    // in the array the item appears at.
20    index: Vec<SlabSlot<usize>>,
21    next_index: usize,
22}
23
24enum SlabSlot<T> {
25    Empty { next: usize },
26    Full { value: T },
27}
28
29pub struct Slot {
30    idx: usize,
31}
32
33impl<T: Ord> Heap<T> {
34    pub fn new() -> Heap<T> {
35        Heap {
36            items: Vec::new(),
37            index: Vec::new(),
38            next_index: 0,
39        }
40    }
41
42    /// Pushes an element onto this heap, returning a slot token indicating
43    /// where it was pushed on to.
44    ///
45    /// The slot can later get passed to `remove` to remove the element from the
46    /// heap, but only if the element was previously not removed from the heap.
47    pub fn push(&mut self, t: T) -> Slot {
48        self.assert_consistent();
49        let len = self.items.len();
50        let slot = SlabSlot::Full { value: len };
51        let slot_idx = if self.next_index == self.index.len() {
52            self.next_index += 1;
53            self.index.push(slot);
54            self.index.len() - 1
55        } else {
56            match mem::replace(&mut self.index[self.next_index], slot) {
57                SlabSlot::Empty { next } => mem::replace(&mut self.next_index, next),
58                SlabSlot::Full { .. } => panic!(),
59            }
60        };
61        self.items.push((t, slot_idx));
62        self.percolate_up(len);
63        self.assert_consistent();
64        Slot { idx: slot_idx }
65    }
66
67    pub fn peek(&self) -> Option<&T> {
68        self.assert_consistent();
69        self.items.get(0).map(|i| &i.0)
70    }
71
72    pub fn pop(&mut self) -> Option<T> {
73        self.assert_consistent();
74        if self.items.is_empty() {
75            return None;
76        }
77        let slot = Slot {
78            idx: self.items[0].1,
79        };
80        Some(self.remove(slot))
81    }
82
83    pub fn remove(&mut self, slot: Slot) -> T {
84        self.assert_consistent();
85        let empty = SlabSlot::Empty {
86            next: self.next_index,
87        };
88        let idx = match mem::replace(&mut self.index[slot.idx], empty) {
89            SlabSlot::Full { value } => value,
90            SlabSlot::Empty { .. } => panic!(),
91        };
92        self.next_index = slot.idx;
93        let (item, slot_idx) = self.items.swap_remove(idx);
94        debug_assert_eq!(slot.idx, slot_idx);
95        if idx < self.items.len() {
96            set_index(&mut self.index, self.items[idx].1, idx);
97            if self.items[idx].0 < item {
98                self.percolate_up(idx);
99            } else {
100                self.percolate_down(idx);
101            }
102        }
103        self.assert_consistent();
104        item
105    }
106
107    fn percolate_up(&mut self, mut idx: usize) -> usize {
108        while idx > 0 {
109            let parent = (idx - 1) / 2;
110            if self.items[idx].0 >= self.items[parent].0 {
111                break;
112            }
113            let (a, b) = self.items.split_at_mut(idx);
114            mem::swap(&mut a[parent], &mut b[0]);
115            set_index(&mut self.index, a[parent].1, parent);
116            set_index(&mut self.index, b[0].1, idx);
117            idx = parent;
118        }
119        idx
120    }
121
122    fn percolate_down(&mut self, mut idx: usize) -> usize {
123        loop {
124            let left = 2 * idx + 1;
125            let right = 2 * idx + 2;
126
127            let mut swap_left = true;
128            match (self.items.get(left), self.items.get(right)) {
129                (Some(left), None) => {
130                    if left.0 >= self.items[idx].0 {
131                        break;
132                    }
133                }
134                (Some(left), Some(right)) => {
135                    if left.0 < self.items[idx].0 {
136                        if right.0 < left.0 {
137                            swap_left = false;
138                        }
139                    } else if right.0 < self.items[idx].0 {
140                        swap_left = false;
141                    } else {
142                        break;
143                    }
144                }
145
146                (None, None) => break,
147                (None, Some(_right)) => panic!("not possible"),
148            }
149
150            let (a, b) = if swap_left {
151                self.items.split_at_mut(left)
152            } else {
153                self.items.split_at_mut(right)
154            };
155            mem::swap(&mut a[idx], &mut b[0]);
156            set_index(&mut self.index, a[idx].1, idx);
157            set_index(&mut self.index, b[0].1, a.len());
158            idx = a.len();
159        }
160        idx
161    }
162
163    fn assert_consistent(&self) {
164        if !cfg!(assert_timer_heap_consistent) {
165            return;
166        }
167
168        assert_eq!(
169            self.items.len(),
170            self.index
171                .iter()
172                .filter(|slot| {
173                    match **slot {
174                        SlabSlot::Full { .. } => true,
175                        SlabSlot::Empty { .. } => false,
176                    }
177                })
178                .count()
179        );
180
181        for (i, &(_, j)) in self.items.iter().enumerate() {
182            let index = match self.index[j] {
183                SlabSlot::Full { value } => value,
184                SlabSlot::Empty { .. } => panic!(),
185            };
186            if index != i {
187                panic!(
188                    "self.index[j] != i : i={} j={} self.index[j]={}",
189                    i, j, index
190                );
191            }
192        }
193
194        for (i, &(ref item, _)) in self.items.iter().enumerate() {
195            if i > 0 {
196                assert!(*item >= self.items[(i - 1) / 2].0, "bad at index: {}", i);
197            }
198            if let Some(left) = self.items.get(2 * i + 1) {
199                assert!(*item <= left.0, "bad left at index: {}", i);
200            }
201            if let Some(right) = self.items.get(2 * i + 2) {
202                assert!(*item <= right.0, "bad right at index: {}", i);
203            }
204        }
205    }
206}
207
208fn set_index<T>(slab: &mut Vec<SlabSlot<T>>, slab_slot: usize, val: T) {
209    match slab[slab_slot] {
210        SlabSlot::Full { ref mut value } => *value = val,
211        SlabSlot::Empty { .. } => panic!(),
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::Heap;
218
219    #[test]
220    fn simple() {
221        let mut h = Heap::new();
222        h.push(1);
223        h.push(2);
224        h.push(8);
225        h.push(4);
226        assert_eq!(h.pop(), Some(1));
227        assert_eq!(h.pop(), Some(2));
228        assert_eq!(h.pop(), Some(4));
229        assert_eq!(h.pop(), Some(8));
230        assert_eq!(h.pop(), None);
231        assert_eq!(h.pop(), None);
232    }
233
234    #[test]
235    fn simple2() {
236        let mut h = Heap::new();
237        h.push(5);
238        h.push(4);
239        h.push(3);
240        h.push(2);
241        h.push(1);
242        assert_eq!(h.pop(), Some(1));
243        h.push(8);
244        assert_eq!(h.pop(), Some(2));
245        h.push(1);
246        assert_eq!(h.pop(), Some(1));
247        assert_eq!(h.pop(), Some(3));
248        assert_eq!(h.pop(), Some(4));
249        h.push(5);
250        assert_eq!(h.pop(), Some(5));
251        assert_eq!(h.pop(), Some(5));
252        assert_eq!(h.pop(), Some(8));
253    }
254
255    #[test]
256    fn remove() {
257        let mut h = Heap::new();
258        h.push(5);
259        h.push(4);
260        h.push(3);
261        let two = h.push(2);
262        h.push(1);
263        assert_eq!(h.pop(), Some(1));
264        assert_eq!(h.remove(two), 2);
265        h.push(1);
266        assert_eq!(h.pop(), Some(1));
267        assert_eq!(h.pop(), Some(3));
268    }
269
270    fn vec2heap<T: Ord>(v: Vec<T>) -> Heap<T> {
271        let mut h = Heap::new();
272        for t in v {
273            h.push(t);
274        }
275        return h;
276    }
277
278    #[test]
279    fn test_peek_and_pop() {
280        let data = vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1];
281        let mut sorted = data.clone();
282        sorted.sort();
283        let mut heap = vec2heap(data);
284        while heap.peek().is_some() {
285            assert_eq!(heap.peek().unwrap(), sorted.first().unwrap());
286            assert_eq!(heap.pop().unwrap(), sorted.remove(0));
287        }
288    }
289
290    #[test]
291    fn test_push() {
292        let mut heap = Heap::new();
293        heap.push(-2);
294        heap.push(-4);
295        heap.push(-9);
296        assert!(*heap.peek().unwrap() == -9);
297        heap.push(-11);
298        assert!(*heap.peek().unwrap() == -11);
299        heap.push(-5);
300        assert!(*heap.peek().unwrap() == -11);
301        heap.push(-27);
302        assert!(*heap.peek().unwrap() == -27);
303        heap.push(-3);
304        assert!(*heap.peek().unwrap() == -27);
305        heap.push(-103);
306        assert!(*heap.peek().unwrap() == -103);
307    }
308
309    fn check_to_vec(mut data: Vec<i32>) {
310        let mut heap = Heap::new();
311        for data in data.iter() {
312            heap.push(*data);
313        }
314        data.sort();
315        let mut v = Vec::new();
316        while let Some(i) = heap.pop() {
317            v.push(i);
318        }
319        assert_eq!(v, data);
320    }
321
322    #[test]
323    fn test_to_vec() {
324        check_to_vec(vec![]);
325        check_to_vec(vec![5]);
326        check_to_vec(vec![3, 2]);
327        check_to_vec(vec![2, 3]);
328        check_to_vec(vec![5, 1, 2]);
329        check_to_vec(vec![1, 100, 2, 3]);
330        check_to_vec(vec![1, 3, 5, 7, 9, 2, 4, 6, 8, 0]);
331        check_to_vec(vec![2, 4, 6, 2, 1, 8, 10, 3, 5, 7, 0, 9, 1]);
332        check_to_vec(vec![9, 11, 9, 9, 9, 9, 11, 2, 3, 4, 11, 9, 0, 0, 0, 0]);
333        check_to_vec(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
334        check_to_vec(vec![10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]);
335        check_to_vec(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, 0, 1, 2]);
336        check_to_vec(vec![5, 4, 3, 2, 1, 5, 4, 3, 2, 1, 5, 4, 3, 2, 1]);
337    }
338
339    #[test]
340    fn test_empty_pop() {
341        let mut heap = Heap::<i32>::new();
342        assert!(heap.pop().is_none());
343    }
344
345    #[test]
346    fn test_empty_peek() {
347        let empty = Heap::<i32>::new();
348        assert!(empty.peek().is_none());
349    }
350}
futures_timer/native/heap.rs

futures_timer/native/
heap.rs
