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 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_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, void *element,
 350                               GNUNET_CONTAINER_HeapCostType cost)
 351 {
 352   struct GNUNET_CONTAINER_HeapNode *node;
 353
 354   node = GNUNET_malloc (sizeof (struct GNUNET_CONTAINER_HeapNode));
 355   node->heap = heap;
 356   node->element = element;
 357   node->cost = cost;
 358   heap->size++;
 359   if (NULL == heap->root)
 360     heap->root = node;
 361   else
 362     insert_node (heap, heap->root, node);
 363   GNUNET_assert (heap->size == heap->root->tree_size + 1);
 364   CHECK (heap->root);
 365   return node;
 366 }
 367
 368
 369 /**
 370  * Remove root of the heap.
 371  *
 372  * @param heap heap to modify
 373  * @return element data stored at the root node, NULL if heap is empty
 374  */
 375 void *
 376 GNUNET_CONTAINER_heap_remove_root (struct GNUNET_CONTAINER_Heap *heap)
 377 {
 378   void *ret;
 379   struct GNUNET_CONTAINER_HeapNode *root;
 380
 381   if (NULL == (root = heap->root))
 382     return NULL;
 383   heap->size--;
 384   ret = root->element;
 385   if (root->left_child == NULL)
 386   {
 387     heap->root = root->right_child;
 388     if (root->right_child != NULL)
 389       root->right_child->parent = NULL;
 390   }
 391   else if (root->right_child == NULL)
 392   {
 393     heap->root = root->left_child;
 394     root->left_child->parent = NULL;
 395   }
 396   else
 397   {
 398     root->left_child->parent = NULL;
 399     root->right_child->parent = NULL;
 400     heap->root = root->left_child;
 401     insert_node (heap, heap->root, root->right_child);
 402   }
 403   if (heap->walk_pos == root)
 404     heap->walk_pos = heap->root;
 405   GNUNET_free (root);
 406 #if EXTRA_CHECKS
 407   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
 408                  (heap->size == heap->root->tree_size + 1));
 409   CHECK (heap->root);
 410 #endif
 411   return ret;
 412 }
 413
 414
 415 /**
 416  * Remove the given node 'node' from the tree and update
 417  * the 'tree_size' fields accordingly.  Preserves the
 418  * children of 'node' and does NOT change the overall
 419  * 'size' field of the tree.
 420  */
 421 static void
 422 remove_node (struct GNUNET_CONTAINER_HeapNode *node)
 423 {
 424   struct GNUNET_CONTAINER_HeapNode *ancestor;
 425   struct GNUNET_CONTAINER_Heap *heap = node->heap;
 426
 427   /* update 'size' of the ancestors */
 428   ancestor = node;
 429   while (NULL != (ancestor = ancestor->parent))
 430     ancestor->tree_size--;
 431
 432   /* update 'size' of node itself */
 433   if (node->left_child != NULL)
 434     node->tree_size -= (1 + node->left_child->tree_size);
 435   if (node->right_child != NULL)
 436     node->tree_size -= (1 + node->right_child->tree_size);
 437
 438   /* unlink 'node' itself and insert children in its place */
 439   if (node->parent == NULL)
 440   {
 441     if (node->left_child != NULL)
 442     {
 443       heap->root = node->left_child;
 444       node->left_child->parent = NULL;
 445       if (node->right_child != NULL)
 446       {
 447         node->right_child->parent = NULL;
 448         insert_node (heap, heap->root, node->right_child);
 449       }
 450     }
 451     else
 452     {
 453       heap->root = node->right_child;
 454       if (node->right_child != NULL)
 455         node->right_child->parent = NULL;
 456     }
 457   }
 458   else
 459   {
 460     if (node->parent->left_child == node)
 461       node->parent->left_child = NULL;
 462     else
 463       node->parent->right_child = NULL;
 464     if (node->left_child != NULL)
 465     {
 466       node->left_child->parent = NULL;
 467       node->parent->tree_size -= (1 + node->left_child->tree_size);
 468       insert_node (heap, node->parent, node->left_child);
 469     }
 470     if (node->right_child != NULL)
 471     {
 472       node->right_child->parent = NULL;
 473       node->parent->tree_size -= (1 + node->right_child->tree_size);
 474       insert_node (heap, node->parent, node->right_child);
 475     }
 476   }
 477   node->parent = NULL;
 478   node->left_child = NULL;
 479   node->right_child = NULL;
 480   GNUNET_assert (node->tree_size == 0);
 481   CHECK (heap->root);
 482 }
 483
 484
 485 /**
 486  * Removes a node from the heap.
 487  *
 488  * @param node node to remove
 489  * @return element data stored at the node
 490  */
 491 void *
 492 GNUNET_CONTAINER_heap_remove_node (struct GNUNET_CONTAINER_HeapNode *node)
 493 {
 494   void *ret;
 495   struct GNUNET_CONTAINER_Heap *heap;
 496
 497   heap = node->heap;
 498   CHECK (heap->root);
 499   if (heap->walk_pos == node)
 500     (void) GNUNET_CONTAINER_heap_walk_get_next (heap);
 501   remove_node (node);
 502   heap->size--;
 503   ret = node->element;
 504   if (heap->walk_pos == node)
 505     heap->walk_pos = NULL;
 506   GNUNET_free (node);
 507 #if EXTRA_CHECKS
 508   CHECK (heap->root);
 509   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
 510                  (heap->size == heap->root->tree_size + 1));
 511 #endif
 512   return ret;
 513 }
 514
 515
 516 /**
 517  * Updates the cost of any node in the tree
 518  *
 519  * @param heap heap to modify
 520  * @param node node for which the cost is to be changed
 521  * @param new_cost new cost for the node
 522  */
 523 void
 524 GNUNET_CONTAINER_heap_update_cost (struct GNUNET_CONTAINER_Heap *heap,
 525                                    struct GNUNET_CONTAINER_HeapNode *node,
 526                                    GNUNET_CONTAINER_HeapCostType new_cost)
 527 {
 528 #if EXTRA_CHECKS
 529   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
 530                  (heap->size == heap->root->tree_size + 1));
 531   CHECK (heap->root);
 532 #endif
 533   remove_node (node);
 534 #if EXTRA_CHECKS
 535   CHECK (heap->root);
 536   GNUNET_assert (((heap->size == 1) && (heap->root == NULL)) ||
 537                  (heap->size == heap->root->tree_size + 2));
 538 #endif
 539   node->cost = new_cost;
 540   if (heap->root == NULL)
 541     heap->root = node;
 542   else
 543     insert_node (heap, heap->root, node);
 544 #if EXTRA_CHECKS
 545   CHECK (heap->root);
 546   GNUNET_assert (((heap->size == 0) && (heap->root == NULL)) ||
 547                  (heap->size == heap->root->tree_size + 1));
 548 #endif
 549 }
 550
 551
 552 /* end of heap.c */