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