implement decode vector
[Mitarbeiter/Tim-Zeitz/stud-rust-base.git] / src / index_heap.rs
1 use std;
2 use std::cmp::min;
3 use std::mem::swap;
4 use std::ptr;
5
6 pub trait Indexing {
7     fn as_index(&self) -> usize;
8 }
9
10 // A priority queue where the elements are IDs from 0 to id_count-1 where id_count is a number that is set in the constructor.
11 // The elements are sorted by integer keys.
12 #[derive(Debug)]
13 pub struct IndexdMinHeap<T: Ord + Indexing> {
14     positions: Vec<usize>,
15     data: Vec<T>
16 }
17
18 const TREE_ARITY: usize = 4;
19 const INVALID_POSITION: usize = std::usize::MAX;
20
21 impl<T: Ord + Indexing> IndexdMinHeap<T> {
22     pub fn new(max_id: usize) -> IndexdMinHeap<T> {
23         IndexdMinHeap {
24             positions: vec![INVALID_POSITION; max_id],
25             data: Vec::new()
26         }
27     }
28
29     pub fn len(&self) -> usize {
30         self.data.len()
31     }
32
33     pub fn is_empty(&self) -> bool {
34         self.len() == 0
35     }
36
37     pub fn contains_index(&self, id: usize) -> bool {
38         self.positions[id] != INVALID_POSITION
39     }
40
41     pub fn clear(&mut self) {
42         for element in &self.data {
43             self.positions[element.as_index()] = INVALID_POSITION;
44         }
45         self.data.clear();
46     }
47
48     pub fn peek(&self) -> Option<&T> {
49         self.data.first()
50     }
51
52     pub fn pop(&mut self) -> Option<T> {
53         self.data.pop().map(|mut item| {
54             self.positions[item.as_index()] = INVALID_POSITION;
55             if !self.is_empty() {
56                 self.positions[item.as_index()] = 0;
57                 self.positions[self.data[0].as_index()] = INVALID_POSITION;
58                 swap(&mut item, &mut self.data[0]);
59                 self.move_down_in_tree(0);
60             }
61             item
62         })
63     }
64
65     pub fn push(&mut self, element: T) {
66         assert!(!self.contains_index(element.as_index()));
67         let insert_position = self.len();
68         self.positions[element.as_index()] = insert_position;
69         self.data.push(element);
70         self.move_up_in_tree(insert_position);
71     }
72
73     // Updates the key of an element if the new key is smaller than the old key.
74     // Does nothing if the new key is larger.
75     // Undefined if the element is not part of the queue.
76     pub fn decrease_key(&mut self, element: T) {
77         let position = self.positions[element.as_index()];
78         self.data[position] = element;
79         self.move_up_in_tree(position);
80     }
81
82     // Updates the key of an element if the new key is larger than the old key.
83     // Does nothing if the new key is smaller.
84     // Undefined if the element is not part of the queue.
85     pub fn increase_key(&mut self, element: T) {
86         let position = self.positions[element.as_index()];
87         self.data[position] = element;
88         self.move_down_in_tree(position);
89     }
90
91     fn move_up_in_tree(&mut self, position: usize) {
92         unsafe {
93             let mut position = position;
94             let mut hole = Hole::new(&mut self.data, position);
95
96             while position > 0 {
97                 let parent = (position - 1) / TREE_ARITY;
98
99                 if hole.get(parent) < hole.element() {
100                     break;
101                 }
102
103                 self.positions[hole.get(parent).as_index()] = position;
104                 hole.move_to(parent);
105                 position = parent;
106             }
107
108             self.positions[hole.element().as_index()] = position;
109         }
110     }
111
112     fn move_down_in_tree(&mut self, position: usize) {
113         unsafe {
114             let mut position = position;
115             let heap_size = self.len();
116             let mut hole = Hole::new(&mut self.data, position);
117
118             loop {
119                 if let Some(smallest_child) = IndexdMinHeap::<T>::children_index_range(position, heap_size).min_by_key(|&child_index| hole.get(child_index)) {
120                     if hole.get(smallest_child) >= hole.element() {
121                         self.positions[hole.element().as_index()] = position;
122                         return; // no child is smaller
123                     }
124
125                     self.positions[hole.get(smallest_child).as_index()] = position;
126                     hole.move_to(smallest_child);
127                     position = smallest_child;
128                 } else {
129                     self.positions[hole.element().as_index()] = position;
130                     return; // no children at all
131                 }
132             }
133         }
134     }
135
136     fn children_index_range(parent_index: usize, heap_size: usize) -> std::ops::Range<usize> {
137         let first_child = TREE_ARITY * parent_index + 1;
138         let last_child = min(TREE_ARITY * parent_index + TREE_ARITY + 1, heap_size);
139         first_child..last_child
140     }
141 }
142
143
144 // This is an optimization copied straight from the rust stdlib binary heap
145 // it allows to avoid always swapping elements pairwise and rather
146 // move each element only once.
147
148
149 /// Hole represents a hole in a slice i.e. an index without valid value
150 /// (because it was moved from or duplicated).
151 /// In drop, `Hole` will restore the slice by filling the hole
152 /// position with the value that was originally removed.
153 struct Hole<'a, T: 'a> {
154     data: &'a mut [T],
155     /// `elt` is always `Some` from new until drop.
156     elt: Option<T>,
157     pos: usize,
158 }
159
160 impl<'a, T> Hole<'a, T> {
161     /// Create a new Hole at index `pos`.
162     ///
163     /// Unsafe because pos must be within the data slice.
164     #[inline]
165     unsafe fn new(data: &'a mut [T], pos: usize) -> Self {
166         debug_assert!(pos < data.len());
167         let elt = ptr::read(&data[pos]);
168         Hole {
169             data,
170             elt: Some(elt),
171             pos,
172         }
173     }
174
175     /// Returns a reference to the element removed.
176     #[inline]
177     fn element(&self) -> &T {
178         self.elt.as_ref().unwrap()
179     }
180
181     /// Returns a reference to the element at `index`.
182     ///
183     /// Unsafe because index must be within the data slice and not equal to pos.
184     #[inline]
185     unsafe fn get(&self, index: usize) -> &T {
186         debug_assert!(index != self.pos);
187         debug_assert!(index < self.data.len());
188         self.data.get_unchecked(index)
189     }
190
191     /// Move hole to new location
192     ///
193     /// Unsafe because index must be within the data slice and not equal to pos.
194     #[inline]
195     unsafe fn move_to(&mut self, index: usize) {
196         debug_assert!(index != self.pos);
197         debug_assert!(index < self.data.len());
198         let index_ptr: *const _ = self.data.get_unchecked(index);
199         let hole_ptr = self.data.get_unchecked_mut(self.pos);
200         ptr::copy_nonoverlapping(index_ptr, hole_ptr, 1);
201         self.pos = index;
202     }
203 }
204
205 impl<'a, T> Drop for Hole<'a, T> {
206     #[inline]
207     fn drop(&mut self) {
208         // fill the hole again
209         unsafe {
210             let pos = self.pos;
211             ptr::write(self.data.get_unchecked_mut(pos), self.elt.take().unwrap());
212         }
213     }
214 }