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