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
  16 /**
  17  * @file util/container_heap.c
  18  * @brief Implementation of a heap
  19  * @author Nathan Evans
  20  * @author Christian Grothoff
  21  */
  22
  23 #include "platform.h"
  24 #include "gnunet_container_lib.h"
  25
  26 #define LOG(kind,...) GNUNET_log_from (kind, "util-container-heap", __VA_ARGS__)
  27
  28 #define EXTRA_CHECKS 0
  29
  30 /**
  31  * Node in the heap.
  32  */
  33 struct GNUNET_CONTAINER_HeapNode
  34 {
  35   /**
  36    * Heap this node belongs to.
  37    */
  38   struct GNUNET_CONTAINER_Heap *heap;
  39
  40   /**
  41    * Parent node.
  42    */
  43   struct GNUNET_CONTAINER_HeapNode *parent;
  44
  45   /**
  46    * Left child.
  47    */
  48   struct GNUNET_CONTAINER_HeapNode *left_child;
  49
  50   /**
  51    * Right child.
  52    */
  53   struct GNUNET_CONTAINER_HeapNode *right_child;
  54
  55   /**
  56    * Our element.
  57    */
  58   void *element;
  59
  60   /**
  61    * Cost for this element.
  62    */
  63   GNUNET_CONTAINER_HeapCostType cost;
  64
  65   /**
  66    * Number of elements below this node in the heap
  67    * (excluding this node itself).
  68    */
  69   unsigned int tree_size;
  70
  71 };
  72
  73 /**
  74  * Handle to a node in a heap.
  75  */
  76 struct GNUNET_CONTAINER_Heap
  77 {
  78
  79   /**
  80    * Root of the heap.
  81    */
  82   struct GNUNET_CONTAINER_HeapNode *root;
  83
  84   /**
  85    * Current position of our random walk.
  86    */
  87   struct GNUNET_CONTAINER_HeapNode *walk_pos;
  88
  89   /**
  90    * Number of elements in the heap.
  91    */
  92   unsigned int size;
  93
  94   /**
  95    * How is the heap sorted?
  96    */
  97   enum GNUNET_CONTAINER_HeapOrder order;
  98
  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 */