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