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