482039a52fbba691dc3c61f27787a461b7e74451
[polintos/scott/priv.git] / include / c++ / util / heap.h
1 // util/heap.h -- A heap-based priority queue implementation
2 //
3 // This software is copyright (c) 2006 Scott Wood <scott@buserror.net>.
4 // 
5 // Permission is hereby granted, free of charge, to any person obtaining a copy of
6 // this software and associated documentation files (the "Software"), to deal with
7 // the Software without restriction, including without limitation the rights to
8 // use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
9 // of the Software, and to permit persons to whom the Software is furnished to do
10 // so, subject to the following conditions:
11 // 
12 //     * Redistributions of source code must retain the above copyright notice,
13 //       this list of conditions and the following disclaimers.
14 // 
15 //     * Redistributions in binary form must reproduce the above copyright notice,
16 //       this list of conditions and the following disclaimers in the
17 //       documentation and/or other materials provided with the distribution.
18 // 
19 //     * The names of the Software's authors and/or contributors
20 //       may not be used to endorse or promote products derived from
21 //       this Software without specific prior written permission.
22 // 
23 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
24 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
25 // FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
26 // CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
27 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
28 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH THE
29 // SOFTWARE.
30
31 #ifndef _UTIL_HEAP_H
32 #define _UTIL_HEAP_H
33
34 #include <assert.h>
35 #include <stddef.h>
36 #include <stdint.h>
37 #include <lowlevel/bitops.h>
38
39 namespace Util {
40         // T must support operator < for key ordering, and
41         // have a Heap<T>::Node member called heap_node.
42
43         template <typename T>
44         class Heap {
45         public:
46                 struct Node {
47                         Node *parent, *left, *right;
48                         
49                         // A pointer to the parent's left or right pointer
50                         // (whichever side this node is on).  This is also
51                         // used for sanity checks (it is NULL for nodes not
52                         // on a queue).
53                         
54                         Node **parentlink;
55                         
56                         Node()
57                         {
58                                 parent = left = right = NULL;
59                                 parentlink = NULL;
60                         }
61                         
62                         bool is_on_heap()
63                         {
64                                 return parentlink != NULL;
65                         }
66                 };
67         
68         private:
69                 Node *top;
70                 size_t num_nodes;
71         
72                 T *node_to_type(Node *n)
73                 {
74                         return reinterpret_cast<T *>(reinterpret_cast<uintptr_t>(n) - 
75                                                      offsetof(T, heap_node));
76                 }
77                 
78                 // Swap n with its parent
79                 void swap(Node *n)
80                 {
81                         Node *p = n->parent;
82                         *p->parentlink = n;
83                         
84                         n->parent = p->parent;
85                         p->parent = n;
86                         
87                         Node **plink = p->parentlink;
88                         
89                         if (n->parentlink == &p->left) {
90                                 assert(p->left == n);
91                         
92                                 Node *tmp = p->right;
93                                 p->right = n->right;
94                                 n->right = tmp;
95                                 
96                                 p->left = n->left;
97                                 n->left = p;
98                                 
99                                 p->parentlink = &n->left;
100                                 
101                                 if (n->right) {
102                                         n->right->parent = n;
103                                         n->right->parentlink = &n->right;
104                                 }
105                         } else {
106                                 assert(n->parentlink == &p->right);
107                                 assert(p->right == n);
108
109                                 Node *tmp = p->left;
110                                 p->left = n->left;
111                                 n->left = tmp;
112                                 
113                                 p->right = n->right;
114                                 n->right = p;
115
116                                 p->parentlink = &n->right;
117                                 
118                                 if (n->left) {
119                                         n->left->parent = n;
120                                         n->left->parentlink = &n->left;
121                                 }
122                         }
123
124                         if (p->left) {
125                                 p->left->parent = p;
126                                 p->left->parentlink = &p->left;
127         
128                                 if (p->right) {
129                                         p->right->parent = p;
130                                         p->right->parentlink = &p->right;
131                                 }
132                         } else {
133                                 assert(!p->right);
134                         }
135                         
136                         n->parentlink = plink;
137                 }
138                 
139                 // Restore heap condition by moving n towards the root,
140                 // if required.
141
142                 void up(Node *n)
143                 {
144                         T *n_type = node_to_type(n);
145
146                         assert(!n->left || !(*node_to_type(n->left) < *n_type));
147                         assert(!n->right || !(*node_to_type(n->right) < *n_type));
148                         
149                         while (n->parent && *n_type < *node_to_type(n->parent)) {
150                                 assert(!n->left || !(*node_to_type(n->left) < *n_type));
151                                 assert(!n->right || !(*node_to_type(n->right) < *n_type));
152                 
153                                 swap(n);
154                         }
155
156                         assert(!n->left || !(*node_to_type(n->left) < *n_type));
157                         assert(!n->right || !(*node_to_type(n->right) < *n_type));
158                 }
159                 
160                 // Restore heap condition by moving n towards the leaves,
161                 // if required.
162
163                 void down(Node *n)
164                 {
165                         T *n_type = node_to_type(n);
166
167                         assert(!n->parent || !(*n_type < *node_to_type(n->parent)));
168                         
169                         // If n has a right, it'll have a left.
170                         while (n->left) {
171                                 assert(!n->parent || !(*n_type < *node_to_type(n->parent)));
172                                 Node *child = n->left;
173                                 
174                                 if (n->right && 
175                                     *node_to_type(n->right) < *node_to_type(n->left))
176                                         child = n->right;
177                         
178                                 if (*node_to_type(child) < *n_type)
179                                         swap(child);
180                                 else
181                                         break;
182                         }
183
184                         assert(!n->parent || !(*n_type < *node_to_type(n->parent)));
185                 }
186                 
187                 Node *get_by_index(size_t num, Node **parent, Node ***link)
188                 {
189                         assert(num_nodes == 0 || top);
190                         assert(num <= num_nodes);
191                         
192                         // It's easier to use 1-based indices, so that the high set bit
193                         // perfectly indicates which level the node is on.
194                         
195                         num += 1;
196                 
197                         Node *n = NULL;
198                         size_t bit = 1 << ll_get_order_round_down(num);
199                         *parent = NULL;
200                         Node **nextlink = &top;
201                         
202                         do {
203                                 *link = nextlink;
204                                 *parent = n;
205                                 n = *nextlink;
206                                 
207                                 bit /= 2;
208                         
209                                 if (num & bit)
210                                         nextlink = &n->right;
211                                 else
212                                         nextlink = &n->left;
213                         } while (bit);
214                         
215                         return n;
216                 }
217
218         public:
219                 Heap()
220                 {
221                         top = NULL;
222                         num_nodes = 0;
223                 }
224                 
225                 bool empty()
226                 {
227                         return num_nodes == 0;
228                 }
229                 
230                 T *get_top()
231                 {
232                         if (!top)
233                                 return NULL;
234                         
235                         return node_to_type(top);
236                 }
237         
238                 void add(T *t)
239                 {
240                         Node *n = &t->heap_node;
241                         assert(!n->is_on_heap());
242                         
243                         Node *ret = get_by_index(num_nodes, &n->parent, &n->parentlink);
244
245                         assert(ret == NULL);
246                         assert(num_nodes == 0 || n->parent);
247                         
248                         num_nodes++;
249                         *n->parentlink = n;
250                         n->left = n->right = NULL;
251                         
252                         up(n);
253                 }
254                 
255                 void del(T *t)
256                 {
257                         Node *n = &t->heap_node;
258                         assert(*n->parentlink == n);
259                         
260                         if (--num_nodes == 0) {
261                                 top = NULL;
262                                 n->parentlink = NULL;
263                                 return;
264                         }
265                 
266                         Node *parent, **link;
267                         Node *last = get_by_index(num_nodes, &parent, &link);
268
269                         assert(last->parent);
270                         assert(last->is_on_heap());
271                         assert(!last->left);
272                         assert(!last->right);
273
274                         *link = NULL;
275                         *n->parentlink = last;
276                         *last = *n;
277                         n->parentlink = NULL;
278
279                         if (last->left) {
280                                 last->left->parent = last;
281                                 last->left->parentlink = &last->left;
282                         }
283                                 
284                         if (last->right) {
285                                 last->right->parent = last;
286                                 last->right->parentlink = &last->right;
287                         }
288                         
289                         down(last);
290                 }
291                 
292                 void requeue(T *t) {
293                         Node *n = &t->heap_node;
294                         assert(*n->parentlink == n);
295                 
296                         if (!n->parent || *node_to_type(n->parent) < *t)
297                                 down(n);
298                         else
299                                 up(n);
300                 }
301         };
302 }
303
304 #endif