cd4d7909aac6ef1f1c32e8b3c1da6e33ce59e520
[oweals/gnunet.git] / src / util / container_heap.c
1 /*
2   This file is part of GNUnet.
3   (C) 2008, 2009 Christian Grothoff (and other contributing authors)
4
5   GNUnet is free software; you can redistribute it and/or modify
6   it under the terms of the GNU General Public License as published
7   by the Free Software Foundation; either version 2, or (at your
8   option) any later version.
9
10   GNUnet is distributed in the hope that it will be useful, but
11   WITHOUT ANY WARRANTY; without even the implied warranty of
12   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
13   General Public License for more details.
14
15   You should have received a copy of the GNU General Public License
16   along with GNUnet; see the file COPYING.  If not, write to the
17   Free Software Foundation, Inc., 59 Temple Place - Suite 330,
18   Boston, MA 02111-1307, USA.
19 */
20
21 /**
22  * @file util/container_heap.c
23  * @brief Implementation of a heap
24  * @author Nathan Evans
25  * @author Christian Grothoff
26  */
27
28 #include "platform.h"
29 #include "gnunet_util_lib.h"
30
31
32 #define DEBUG 0
33
34 /**
35  * Node in the heap.
36  */
37 struct GNUNET_CONTAINER_HeapNode
38 {
39   /**
40    * Heap this node belongs to.
41    */
42   struct GNUNET_CONTAINER_Heap *heap;
43
44   /**
45    * Parent node.
46    */
47   struct GNUNET_CONTAINER_HeapNode *parent;
48
49   /**
50    * Left child.
51    */
52   struct GNUNET_CONTAINER_HeapNode *left_child;
53
54   /**
55    * Right child.
56    */
57   struct GNUNET_CONTAINER_HeapNode *right_child;
58
59   /**
60    * Our element.
61    */
62   void *element;
63
64   /**
65    * Cost for this element.
66    */
67   GNUNET_CONTAINER_HeapCostType cost;
68
69   /**
70    * Number of elements below this node in the heap
71    * (excluding this node itself).
72    */
73   unsigned int tree_size;
74
75 };
76
77 /**
78  * Handle to a node in a heap.
79  */
80 struct GNUNET_CONTAINER_Heap
81 {
82
83   /**
84    * Root of the heap.
85    */
86   struct GNUNET_CONTAINER_HeapNode *root;
87
88   /**
89    * Current position of our random walk.
90    */
91   struct GNUNET_CONTAINER_HeapNode *walk_pos;
92
93   /**
94    * Number of elements in the heap.
95    */
96   unsigned int size;
97
98   /**
99    * How is the heap sorted?
100    */
101   enum GNUNET_CONTAINER_HeapOrder order;
102
103 };
104
105
106 #if DEBUG
107 /**
108  * Check if internal invariants hold for the given node.
109  *
110  * @param node subtree to check
111  */
112 static void
113 check (const struct GNUNET_CONTAINER_HeapNode *node)
114 {
115   if (NULL == node)
116     return;
117   GNUNET_assert (node->tree_size ==
118                  ((node->left_child ==
119                    NULL) ? 0 : 1 + node->left_child->tree_size) +
120                  ((node->right_child ==
121                    NULL) ? 0 : 1 + node->right_child->tree_size));
122   check (node->left_child);
123   check (node->right_child);
124 }
125
126
127 #define CHECK(n) check(n)
128 #else
129 #define CHECK(n) do {} while (0)
130 #endif
131
132
133 /**
134  * Create a new heap.
135  *
136  * @param order how should the heap be sorted?
137  * @return handle to the heap
138  */
139 struct GNUNET_CONTAINER_Heap *
140 GNUNET_CONTAINER_heap_create (enum GNUNET_CONTAINER_HeapOrder order)
141 {
142   struct GNUNET_CONTAINER_Heap *heap;
143
144   heap = GNUNET_malloc (sizeof (struct GNUNET_CONTAINER_Heap));
145   heap->order = order;
146   return heap;
147 }
148
149
150 /**
151  * Destroys the heap.  Only call on a heap that
152  * is already empty.
153  *
154  * @param heap heap to destroy
155  */
156 void
157 GNUNET_CONTAINER_heap_destroy (struct GNUNET_CONTAINER_Heap *heap)
158 {
159   GNUNET_break (heap->size == 0);
160   GNUNET_free (heap);
161 }
162
163
164 /**
165  * Get element stored at root of heap.
166  *
167  * @param heap heap to inspect
168  * @return NULL if heap is empty
169  */
170 void *
171 GNUNET_CONTAINER_heap_peek (const struct GNUNET_CONTAINER_Heap *heap)
172 {
173   if (heap->root == NULL)
174     return NULL;
175   return heap->root->element;
176 }
177
178
179 /**
180  * Get the current size of the heap
181  *
182  * @param heap the heap to get the size of
183  * @return number of elements stored
184  */
185 unsigned int
186 GNUNET_CONTAINER_heap_get_size (const struct GNUNET_CONTAINER_Heap *heap)
187 {
188   return heap->size;
189 }
190
191
192 /**
193  * Get the current cost of the node
194  *
195  * @param node the node to get the cost of
196  * @return cost of the node
197  */
198 GNUNET_CONTAINER_HeapCostType
199 GNUNET_CONTAINER_heap_node_get_cost (const struct GNUNET_CONTAINER_HeapNode
200                                      *node)
201 {
202   return node->cost;
203 }
204
205 /**
206  * Iterate over the children of the given node.
207  *
208  * @param heap argument to give to iterator
209  * @param node node to iterate over
210  * @param iterator function to call on each node
211  * @param iterator_cls closure for iterator
212  * @return GNUNET_YES to continue to iterate
213  */
214 static int
215 node_iterator (const struct GNUNET_CONTAINER_Heap *heap,
216                struct GNUNET_CONTAINER_HeapNode *node,
217                GNUNET_CONTAINER_HeapIterator iterator, void *iterator_cls)
218 {
219   if (node == NULL)
220     return GNUNET_YES;
221   if (GNUNET_YES !=
222       node_iterator (heap, node->left_child, iterator, iterator_cls))
223     return GNUNET_NO;
224   if (GNUNET_YES !=
225       node_iterator (heap, node->right_child, iterator, iterator_cls))
226     return GNUNET_NO;
227   return iterator (iterator_cls, node, node->element, node->cost);
228 }
229
230
231 /**
232  * Iterate over all entries in the heap.
233  *
234  * @param heap the heap
235  * @param iterator function to call on each entry
236  * @param iterator_cls closure for iterator
237  */
238 void
239 GNUNET_CONTAINER_heap_iterate (const struct GNUNET_CONTAINER_Heap *heap,
240                                GNUNET_CONTAINER_HeapIterator iterator,
241                                void *iterator_cls)
242 {
243   (void) node_iterator (heap, heap->root, iterator, iterator_cls);
244 }
245
246
247 /**
248  * Perform a random walk of the tree.  The walk is biased
249  * towards elements closer to the root of the tree (since
250  * each walk starts at the root and ends at a random leaf).
251  * The heap internally tracks the current position of the
252  * walk.
253  *
254  * @param heap heap to walk
255  * @return data stored at the next random node in the walk;
256  *         NULL if the tree is empty.
257  */
258 void *
259 GNUNET_CONTAINER_heap_walk_get_next (struct GNUNET_CONTAINER_Heap *heap)
260 {
261   struct GNUNET_CONTAINER_HeapNode *pos;
262   void *element;
263
264   if (heap->root == NULL)
265     return NULL;
266   pos = heap->walk_pos;
267   if (pos == NULL)
268     pos = heap->root;
269   element = pos->element;
270   heap->walk_pos =
271       (0 ==
272        GNUNET_CRYPTO_random_u32 (GNUNET_CRYPTO_QUALITY_WEAK,
273                                  2)) ? pos->right_child : pos->left_child;
274   return element;
275 }
276
277
278 /**
279  * Insert the given node 'node' into the subtree starting
280  * at 'pos' (while keeping the tree somewhat balanced).
281  *
282  * @param heap heap to modify
283  * @param pos existing tree
284  * @param node node to insert (which may be a subtree itself)
285  */
286 static void
287 insert_node (struct GNUNET_CONTAINER_Heap *heap,
288              struct GNUNET_CONTAINER_HeapNode *pos,
289              struct GNUNET_CONTAINER_HeapNode *node)
290 {
291   struct GNUNET_CONTAINER_HeapNode *parent;
292
293   GNUNET_assert (node->parent == NULL);
294   while ((heap->order == GNUNET_CONTAINER_HEAP_ORDER_MAX) ? (pos->cost >=
295                                                              node->cost)
296          : (pos->cost <= node->cost))
297   {
298     /* node is descendent of pos */
299     pos->tree_size += (1 + node->tree_size);
300     if (pos->left_child == NULL)
301     {
302       pos->left_child = node;
303       node->parent = pos;
304       return;
305     }
306     if (pos->right_child == NULL)
307     {
308       pos->right_child = node;
309       node->parent = pos;
310       return;
311     }
312     /* keep it balanced by descending into smaller subtree */
313     if (pos->left_child->tree_size < pos->right_child->tree_size)
314       pos = pos->left_child;
315     else
316       pos = pos->right_child;
317   }
318   /* make 'node' parent of 'pos' */
319   parent = pos->parent;
320   pos->parent = NULL;
321   node->parent = parent;
322   if (NULL == parent)
323   {
324     heap->root = node;
325   }
326   else
327   {
328     if (parent->left_child == pos)
329       parent->left_child = node;
330     else
331       parent->right_child = node;
332   }
333   /* insert 'pos' below 'node' */
334   insert_node (heap, node, pos);
335   CHECK (pos);
336 }
337
338
339 /**
340  * Inserts a new element into the heap.
341  *
342  * @param heap heap to modify
343  * @param element element to insert
344  * @param cost cost for the element
345  * @return node for the new element
346  */
347 struct GNUNET_CONTAINER_HeapNode *
348 GNUNET_CONTAINER_heap_insert (struct GNUNET_CONTAINER_Heap *heap, void *element,
349                               GNUNET_CONTAINER_HeapCostType cost)
350 {
351   struct GNUNET_CONTAINER_HeapNode *node;
352
353   node = GNUNET_malloc (sizeof (struct GNUNET_CONTAINER_HeapNode));
354   node->heap = heap;
355   node->element = element;
356   node->cost = cost;
357   heap->size++;
358   if (NULL == heap->root)
359     heap->root = node;
360   else
361     insert_node (heap, heap->root, node);
362   GNUNET_assert (heap->size == heap->root->tree_size + 1);
363   CHECK (heap->root);
364   return node;
365 }
366
367
368 /**
369  * Remove root of the heap.
370  *
371  * @param heap heap to modify
372  * @return element data stored at the root node, NULL if heap is empty
373  */
374 void *
375 GNUNET_CONTAINER_heap_remove_root (struct GNUNET_CONTAINER_Heap *heap)
376 {
377   void *ret;
378   struct GNUNET_CONTAINER_HeapNode *root;
379
380   if (NULL == (root = heap->root))
381     return NULL;
382   heap->size--;
383   ret = root->element;
384   if (root->left_child == NULL)
385   {
386     heap->root = root->right_child;
387     if (root->right_child != NULL)
388       root->right_child->parent = NULL;
389   }
390   else if (root->right_child == NULL)
391   {
392     heap->root = root->left_child;
393     root->left_child->parent = NULL;
394   }
395   else
396   {
397     root->left_child->parent = NULL;
398     root->right_child->parent = NULL;
399     heap->root = root->left_child;
400     insert_node (heap, heap->root, root->right_child);
401   }
402   GNUNET_free (root);
403 #if DEBUG
404   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
405                  (heap->size == heap->root->tree_size + 1));
406   CHECK (heap->root);
407 #endif
408   return ret;
409 }
410
411
412 /**
413  * Remove the given node 'node' from the tree and update
414  * the 'tree_size' fields accordingly.  Preserves the
415  * children of 'node' and does NOT change the overall
416  * 'size' field of the tree.
417  */
418 static void
419 remove_node (struct GNUNET_CONTAINER_HeapNode *node)
420 {
421   struct GNUNET_CONTAINER_HeapNode *ancestor;
422   struct GNUNET_CONTAINER_Heap *heap = node->heap;
423
424   /* update 'size' of the ancestors */
425   ancestor = node;
426   while (NULL != (ancestor = ancestor->parent))
427     ancestor->tree_size--;
428
429   /* update 'size' of node itself */
430   if (node->left_child != NULL)
431     node->tree_size -= (1 + node->left_child->tree_size);
432   if (node->right_child != NULL)
433     node->tree_size -= (1 + node->right_child->tree_size);
434
435   /* unlink 'node' itself and insert children in its place */
436   if (node->parent == NULL)
437   {
438     if (node->left_child != NULL)
439     {
440       heap->root = node->left_child;
441       node->left_child->parent = NULL;
442       if (node->right_child != NULL)
443       {
444         node->right_child->parent = NULL;
445         insert_node (heap, heap->root, node->right_child);
446       }
447     }
448     else
449     {
450       heap->root = node->right_child;
451       if (node->right_child != NULL)
452         node->right_child->parent = NULL;
453     }
454   }
455   else
456   {
457     if (node->parent->left_child == node)
458       node->parent->left_child = NULL;
459     else
460       node->parent->right_child = NULL;
461     if (node->left_child != NULL)
462     {
463       node->left_child->parent = NULL;
464       node->parent->tree_size -= (1 + node->left_child->tree_size);
465       insert_node (heap, node->parent, node->left_child);
466     }
467     if (node->right_child != NULL)
468     {
469       node->right_child->parent = NULL;
470       node->parent->tree_size -= (1 + node->right_child->tree_size);
471       insert_node (heap, node->parent, node->right_child);
472     }
473   }
474   node->parent = NULL;
475   node->left_child = NULL;
476   node->right_child = NULL;
477   GNUNET_assert (node->tree_size == 0);
478   CHECK (heap->root);
479 }
480
481
482 /**
483  * Removes a node from the heap.
484  *
485  * @param node node to remove
486  * @return element data stored at the node
487  */
488 void *
489 GNUNET_CONTAINER_heap_remove_node (struct GNUNET_CONTAINER_HeapNode *node)
490 {
491   void *ret;
492   struct GNUNET_CONTAINER_Heap *heap;
493
494   heap = node->heap;
495   CHECK (heap->root);
496   if (heap->walk_pos == node)
497     (void) GNUNET_CONTAINER_heap_walk_get_next (heap);
498   remove_node (node);
499   heap->size--;
500   ret = node->element;
501   if (heap->walk_pos == node)
502     heap->walk_pos = NULL;
503   GNUNET_free (node);
504 #if DEBUG
505   CHECK (heap->root);
506   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
507                  (heap->size == heap->root->tree_size + 1));
508 #endif
509   return ret;
510 }
511
512
513 /**
514  * Updates the cost of any node in the tree
515  *
516  * @param heap heap to modify
517  * @param node node for which the cost is to be changed
518  * @param new_cost new cost for the node
519  */
520 void
521 GNUNET_CONTAINER_heap_update_cost (struct GNUNET_CONTAINER_Heap *heap,
522                                    struct GNUNET_CONTAINER_HeapNode *node,
523                                    GNUNET_CONTAINER_HeapCostType new_cost)
524 {
525 #if DEBUG
526   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
527                  (heap->size == heap->root->tree_size + 1));
528   CHECK (heap->root);
529 #endif
530   remove_node (node);
531 #if DEBUG
532   CHECK (heap->root);
533   GNUNET_assert (((heap->size == 1) && (heap->root == NULL)) ||
534                  (heap->size == heap->root->tree_size + 2));
535 #endif
536   node->cost = new_cost;
537   if (heap->root == NULL)
538     heap->root = node;
539   else
540     insert_node (heap, heap->root, node);
541 #if DEBUG
542   CHECK (heap->root);
543   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
544                  (heap->size == heap->root->tree_size + 1));
545 #endif
546 }
547
548
549 /* end of heap.c */