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