src/util/container_heap.c

   1 /*
   2   This file is part of GNUnet.
   3   (C) 2008, 2009 Christian Grothoff (and other contributing authors)
   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 2, 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., 59 Temple Place - Suite 330,
  18   Boston, MA 02111-1307, 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_util_lib.h"
  30
  31 #define LOG(kind,...) GNUNET_log_from (kind, "util", __VA_ARGS__)
  32
  33 #define DEBUG 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 DEBUG
 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_malloc (sizeof (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 root of heap.
 167  *
 168  * @param heap heap to inspect
 169  * @return 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 the current size of the heap
 182  *
 183  * @param heap the heap to get the size of
 184  * @return number of elements stored
 185  */
 186 unsigned int
 187 GNUNET_CONTAINER_heap_get_size (const struct GNUNET_CONTAINER_Heap *heap)
 188 {
 189   return heap->size;
 190 }
 191
 192
 193 /**
 194  * Get the current cost of the node
 195  *
 196  * @param node the node to get the cost of
 197  * @return cost of the node
 198  */
 199 GNUNET_CONTAINER_HeapCostType
 200 GNUNET_CONTAINER_heap_node_get_cost (const struct GNUNET_CONTAINER_HeapNode
 201                                      * node)
 202 {
 203   return node->cost;
 204 }
 205
 206 /**
 207  * Iterate over the children of the given node.
 208  *
 209  * @param heap argument to give to iterator
 210  * @param node node to iterate over
 211  * @param iterator function to call on each node
 212  * @param iterator_cls closure for iterator
 213  * @return GNUNET_YES to continue to iterate
 214  */
 215 static int
 216 node_iterator (const struct GNUNET_CONTAINER_Heap *heap,
 217                struct GNUNET_CONTAINER_HeapNode *node,
 218                GNUNET_CONTAINER_HeapIterator iterator, void *iterator_cls)
 219 {
 220   if (node == NULL)
 221     return GNUNET_YES;
 222   if (GNUNET_YES !=
 223       node_iterator (heap, node->left_child, iterator, iterator_cls))
 224     return GNUNET_NO;
 225   if (GNUNET_YES !=
 226       node_iterator (heap, node->right_child, iterator, iterator_cls))
 227     return GNUNET_NO;
 228   return iterator (iterator_cls, node, node->element, node->cost);
 229 }
 230
 231
 232 /**
 233  * Iterate over all entries in the heap.
 234  *
 235  * @param heap the heap
 236  * @param iterator function to call on each entry
 237  * @param iterator_cls closure for iterator
 238  */
 239 void
 240 GNUNET_CONTAINER_heap_iterate (const struct GNUNET_CONTAINER_Heap *heap,
 241                                GNUNET_CONTAINER_HeapIterator iterator,
 242                                void *iterator_cls)
 243 {
 244   (void) node_iterator (heap, heap->root, iterator, iterator_cls);
 245 }
 246
 247
 248 /**
 249  * Perform a random walk of the tree.  The walk is biased
 250  * towards elements closer to the root of the tree (since
 251  * each walk starts at the root and ends at a random leaf).
 252  * The heap internally tracks the current position of the
 253  * walk.
 254  *
 255  * @param heap heap to walk
 256  * @return data stored at the next random node in the walk;
 257  *         NULL if the tree is empty.
 258  */
 259 void *
 260 GNUNET_CONTAINER_heap_walk_get_next (struct GNUNET_CONTAINER_Heap *heap)
 261 {
 262   struct GNUNET_CONTAINER_HeapNode *pos;
 263   void *element;
 264
 265   if (heap->root == NULL)
 266     return NULL;
 267   pos = heap->walk_pos;
 268   if (pos == NULL)
 269     pos = heap->root;
 270   element = pos->element;
 271   heap->walk_pos =
 272     (0 ==
 273      GNUNET_CRYPTO_random_u32 (GNUNET_CRYPTO_QUALITY_WEAK,
 274                                2)) ? pos->right_child : pos->left_child;
 275   return element;
 276 }
 277
 278
 279 /**
 280  * Insert the given node 'node' into the subtree starting
 281  * at 'pos' (while keeping the tree somewhat balanced).
 282  *
 283  * @param heap heap to modify
 284  * @param pos existing tree
 285  * @param node node to insert (which may be a subtree itself)
 286  */
 287 static void
 288 insert_node (struct GNUNET_CONTAINER_Heap *heap,
 289              struct GNUNET_CONTAINER_HeapNode *pos,
 290              struct GNUNET_CONTAINER_HeapNode *node)
 291 {
 292   struct GNUNET_CONTAINER_HeapNode *parent;
 293
 294   GNUNET_assert (node->parent == NULL);
 295   while ((heap->order == GNUNET_CONTAINER_HEAP_ORDER_MAX) ? (pos->cost >=
 296                                                              node->cost)
 297          : (pos->cost <= node->cost))
 298     {
 299       /* node is descendent of pos */
 300       pos->tree_size += (1 + node->tree_size);
 301       if (pos->left_child == NULL)
 302         {
 303           pos->left_child = node;
 304           node->parent = pos;
 305           return;
 306         }
 307       if (pos->right_child == NULL)
 308         {
 309           pos->right_child = node;
 310           node->parent = pos;
 311           return;
 312         }
 313       /* keep it balanced by descending into smaller subtree */
 314       if (pos->left_child->tree_size < pos->right_child->tree_size)
 315         pos = pos->left_child;
 316       else
 317         pos = pos->right_child;
 318     }
 319   /* make 'node' parent of 'pos' */
 320   parent = pos->parent;
 321   pos->parent = NULL;
 322   node->parent = parent;
 323   if (NULL == parent)
 324     {
 325       heap->root = node;
 326     }
 327   else
 328     {
 329       if (parent->left_child == pos)
 330         parent->left_child = node;
 331       else
 332         parent->right_child = node;
 333     }
 334   /* insert 'pos' below 'node' */
 335   insert_node (heap, node, pos);
 336   CHECK (pos);
 337 }
 338
 339
 340 /**
 341  * Inserts a new element into the heap.
 342  *
 343  * @param heap heap to modify
 344  * @param element element to insert
 345  * @param cost cost for the element
 346  * @return node for the new element
 347  */
 348 struct GNUNET_CONTAINER_HeapNode *
 349 GNUNET_CONTAINER_heap_insert (struct GNUNET_CONTAINER_Heap *heap,
 350                               void *element,
 351                               GNUNET_CONTAINER_HeapCostType cost)
 352 {
 353   struct GNUNET_CONTAINER_HeapNode *node;
 354
 355   node = GNUNET_malloc (sizeof (struct GNUNET_CONTAINER_HeapNode));
 356   node->heap = heap;
 357   node->element = element;
 358   node->cost = cost;
 359   heap->size++;
 360   if (NULL == heap->root)
 361     heap->root = node;
 362   else
 363     insert_node (heap, heap->root, node);
 364   GNUNET_assert (heap->size == heap->root->tree_size + 1);
 365   CHECK (heap->root);
 366   return node;
 367 }
 368
 369
 370 /**
 371  * Remove root of the heap.
 372  *
 373  * @param heap heap to modify
 374  * @return element data stored at the root node, NULL if heap is empty
 375  */
 376 void *
 377 GNUNET_CONTAINER_heap_remove_root (struct GNUNET_CONTAINER_Heap *heap)
 378 {
 379   void *ret;
 380   struct GNUNET_CONTAINER_HeapNode *root;
 381
 382   if (NULL == (root = heap->root))
 383     return NULL;
 384   heap->size--;
 385   ret = root->element;
 386   if (root->left_child == NULL)
 387     {
 388       heap->root = root->right_child;
 389       if (root->right_child != NULL)
 390         root->right_child->parent = NULL;
 391     }
 392   else if (root->right_child == NULL)
 393     {
 394       heap->root = root->left_child;
 395       root->left_child->parent = NULL;
 396     }
 397   else
 398     {
 399       root->left_child->parent = NULL;
 400       root->right_child->parent = NULL;
 401       heap->root = root->left_child;
 402       insert_node (heap, heap->root, root->right_child);
 403     }
 404   GNUNET_free (root);
 405 #if DEBUG
 406   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
 407                  (heap->size == heap->root->tree_size + 1));
 408   CHECK (heap->root);
 409 #endif
 410   return ret;
 411 }
 412
 413
 414 /**
 415  * Remove the given node 'node' from the tree and update
 416  * the 'tree_size' fields accordingly.  Preserves the
 417  * children of 'node' and does NOT change the overall
 418  * 'size' field of the tree.
 419  */
 420 static void
 421 remove_node (struct GNUNET_CONTAINER_HeapNode *node)
 422 {
 423   struct GNUNET_CONTAINER_HeapNode *ancestor;
 424   struct GNUNET_CONTAINER_Heap *heap = node->heap;
 425
 426   /* update 'size' of the ancestors */
 427   ancestor = node;
 428   while (NULL != (ancestor = ancestor->parent))
 429     ancestor->tree_size--;
 430
 431   /* update 'size' of node itself */
 432   if (node->left_child != NULL)
 433     node->tree_size -= (1 + node->left_child->tree_size);
 434   if (node->right_child != NULL)
 435     node->tree_size -= (1 + node->right_child->tree_size);
 436
 437   /* unlink 'node' itself and insert children in its place */
 438   if (node->parent == NULL)
 439     {
 440       if (node->left_child != NULL)
 441         {
 442           heap->root = node->left_child;
 443           node->left_child->parent = NULL;
 444           if (node->right_child != NULL)
 445             {
 446               node->right_child->parent = NULL;
 447               insert_node (heap, heap->root, node->right_child);
 448             }
 449         }
 450       else
 451         {
 452           heap->root = node->right_child;
 453           if (node->right_child != NULL)
 454             node->right_child->parent = NULL;
 455         }
 456     }
 457   else
 458     {
 459       if (node->parent->left_child == node)
 460         node->parent->left_child = NULL;
 461       else
 462         node->parent->right_child = NULL;
 463       if (node->left_child != NULL)
 464         {
 465           node->left_child->parent = NULL;
 466           node->parent->tree_size -= (1 + node->left_child->tree_size);
 467           insert_node (heap, node->parent, node->left_child);
 468         }
 469       if (node->right_child != NULL)
 470         {
 471           node->right_child->parent = NULL;
 472           node->parent->tree_size -= (1 + node->right_child->tree_size);
 473           insert_node (heap, node->parent, node->right_child);
 474         }
 475     }
 476   node->parent = NULL;
 477   node->left_child = NULL;
 478   node->right_child = NULL;
 479   GNUNET_assert (node->tree_size == 0);
 480   CHECK (heap->root);
 481 }
 482
 483
 484 /**
 485  * Removes a node from the heap.
 486  *
 487  * @param node node to remove
 488  * @return element data stored at the node
 489  */
 490 void *
 491 GNUNET_CONTAINER_heap_remove_node (struct GNUNET_CONTAINER_HeapNode *node)
 492 {
 493   void *ret;
 494   struct GNUNET_CONTAINER_Heap *heap;
 495
 496   heap = node->heap;
 497   CHECK (heap->root);
 498   if (heap->walk_pos == node)
 499     (void) GNUNET_CONTAINER_heap_walk_get_next (heap);
 500   remove_node (node);
 501   heap->size--;
 502   ret = node->element;
 503   if (heap->walk_pos == node)
 504     heap->walk_pos = NULL;
 505   GNUNET_free (node);
 506 #if DEBUG
 507   CHECK (heap->root);
 508   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
 509                  (heap->size == heap->root->tree_size + 1));
 510 #endif
 511   return ret;
 512 }
 513
 514
 515 /**
 516  * Updates the cost of any node in the tree
 517  *
 518  * @param heap heap to modify
 519  * @param node node for which the cost is to be changed
 520  * @param new_cost new cost for the node
 521  */
 522 void
 523 GNUNET_CONTAINER_heap_update_cost (struct GNUNET_CONTAINER_Heap *heap,
 524                                    struct GNUNET_CONTAINER_HeapNode *node,
 525                                    GNUNET_CONTAINER_HeapCostType new_cost)
 526 {
 527 #if DEBUG
 528   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
 529                  (heap->size == heap->root->tree_size + 1));
 530   CHECK (heap->root);
 531 #endif
 532   remove_node (node);
 533 #if DEBUG
 534   CHECK (heap->root);
 535   GNUNET_assert (((heap->size == 1) && (heap->root == NULL)) ||
 536                  (heap->size == heap->root->tree_size + 2));
 537 #endif
 538   node->cost = new_cost;
 539   if (heap->root == NULL)
 540     heap->root = node;
 541   else
 542     insert_node (heap, heap->root, node);
 543 #if DEBUG
 544   CHECK (heap->root);
 545   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
 546                  (heap->size == heap->root->tree_size + 1));
 547 #endif
 548 }
 549
 550
 551 /* end of heap.c */