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
   6   it under the terms of the GNU General Public License as published
   7   by the Free Software Foundation; either version 3, 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., 51 Franklin Street, Fifth Floor,
  18   Boston, MA 02110-1301, 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_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 */