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