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