coregrind/m_oset.c

   1
   2 /*--------------------------------------------------------------------*/
   3 /*--- An ordered set implemented using an AVL tree.       m_oset.c ---*/
   4 /*--------------------------------------------------------------------*/
   5
   6 /*
   7    This file is part of Valgrind, a dynamic binary instrumentation
   8    framework.
   9
  10    Copyright (C) 2005-2013 Nicholas Nethercote
  11       njn@valgrind.org
  12
  13    This program is free software; you can redistribute it and/or
  14    modify it under the terms of the GNU General Public License as
  15    published by the Free Software Foundation; either version 2 of the
  16    License, or (at your option) any later version.
  17
  18    This program is distributed in the hope that it will be useful, but
  19    WITHOUT ANY WARRANTY; without even the implied warranty of
  20    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  21    General Public License for more details.
  22
  23    You should have received a copy of the GNU General Public License
  24    along with this program; if not, write to the Free Software
  25    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
  26    02111-1307, USA.
  27
  28    The GNU General Public License is contained in the file COPYING.
  29 */
  30
  31 //----------------------------------------------------------------------
  32 // This file is based on:
  33 //
  34 //   ANSI C Library for maintainance of AVL Balanced Trees
  35 //   (C) 2000 Daniel Nagy, Budapest University of Technology and Economics
  36 //   Released under GNU General Public License (GPL) version 2
  37 //----------------------------------------------------------------------
  38
  39 // This file implements a generic ordered set using an AVL tree.
  40 //
  41 // Each node in the tree has two parts.
  42 // - First is the AVL metadata, which is three words: a left pointer, a
  43 //   right pointer, and a word containing balancing information and a
  44 //   "magic" value which provides some checking that the user has not
  45 //   corrupted the metadata.  So the overhead is 12 bytes on 32-bit
  46 //   platforms and 24 bytes on 64-bit platforms.
  47 // - Second is the user's data.  This can be anything.  Note that because it
  48 //   comes after the metadata, it will only be word-aligned, even if the
  49 //   user data is a struct that would normally be doubleword-aligned.
  50 //
  51 // AvlNode* node -> +---------------+  V
  52 //                  | struct        |
  53 //                  |   AvlNode     |
  54 // void* element -> +---------------+  ^
  55 //                  | element       |  |
  56 //      keyOff ->   | key           | elemSize
  57 //                  +---------------+  v
  58 //
  59 // Users have to allocate AvlNodes with OSetGen_AllocNode(), which allocates
  60 // space for the metadata.
  61 //
  62 // The terminology used throughout this file:
  63 // - a "node", usually called "n", is a pointer to the metadata.
  64 // - an "element", usually called "e", is a pointer to the user data.
  65 // - a "key", usually called "k", is a pointer to a key.
  66 //
  67 // The helper functions elem_of_node and node_of_elem do the pointer
  68 // arithmetic to switch between the node and the element.  The node magic is
  69 // checked after each operation to make sure that we're really operating on
  70 // an AvlNode.
  71 //
  72 // Each tree also has an iterator.  Note that we cannot use the iterator
  73 // internally within this file (eg. we could implement OSetGen_Size() by
  74 // stepping through with the iterator and counting nodes) because it's
  75 // non-reentrant -- the user might be using it themselves, and the
  76 // concurrent uses would screw things up.
  77
  78 #include "pub_core_basics.h"
  79 #include "pub_core_libcbase.h"
  80 #include "pub_core_libcassert.h"
  81 #include "pub_core_libcprint.h"
  82 #include "pub_core_oset.h"
  83 #include "pub_core_poolalloc.h"
  84
  85 /*--------------------------------------------------------------------*/
  86 /*--- Types and constants                                          ---*/
  87 /*--------------------------------------------------------------------*/
  88
  89 typedef struct _OSetNode OSetNode;
  90
  91 // Internal names for the OSet types.
  92 typedef OSet     AvlTree;
  93 typedef OSetNode AvlNode;
  94
  95 // The padding ensures that magic is right at the end of the node,
  96 // regardless of the machine's word size, so that any overwrites will be
  97 // detected earlier.
  98 struct _OSetNode {
  99    AvlNode* left;
 100    AvlNode* right;
 101    Char     balance;
 102    Char     padding[sizeof(void*)-sizeof(Char)-sizeof(Short)];
 103    Short    magic;
 104 };
 105
 106 #define STACK_MAX    32    // At most 2**32 entries can be iterated over
 107 #define OSET_MAGIC   0x5b1f
 108
 109 // An OSet (AVL tree).  If cmp is NULL, the key must be a UWord, and must
 110 // be the first word in the element.  If cmp is set, arbitrary keys in
 111 // arbitrary positions can be used.
 112 struct _OSet {
 113    SizeT       keyOff;     // key offset
 114    OSetCmp_t   cmp;        // compare a key and an element, or NULL
 115    OSetAlloc_t alloc_fn;   // allocator
 116    const HChar* cc;        // cost centre for allocator
 117    OSetFree_t  free_fn;    // deallocator
 118    PoolAlloc*  node_pa;    // (optional) pool allocator for nodes.
 119    SizeT       maxEltSize; // for node_pa, must be > 0. Otherwise unused.
 120    Word        nElems;     // number of elements in the tree
 121    AvlNode*    root;       // root node
 122
 123    AvlNode*    nodeStack[STACK_MAX];   // Iterator node stack
 124    Int          numStack[STACK_MAX];   // Iterator num stack
 125    Int         stackTop;               // Iterator stack pointer, one past end
 126 };
 127
 128 /*--------------------------------------------------------------------*/
 129 /*--- Helper operations                                            ---*/
 130 /*--------------------------------------------------------------------*/
 131
 132 // Given a pointer to the node's element, return the pointer to the AvlNode
 133 // structure.  If the node has a bad magic number, it will die with an
 134 // assertion failure.
 135 static inline
 136 AvlNode* node_of_elem(const void *elem)
 137 {
 138    AvlNode* n = (AvlNode*)((Addr)elem - sizeof(AvlNode));
 139    vg_assert2(n->magic == OSET_MAGIC,
 140               "bad magic on node %p = %x (expected %x)\n"
 141               "possible causes:\n"
 142               " - node not allocated with VG_(OSetGen_AllocNode)()?\n"
 143               " - node metadata corrupted by underwriting start of element?\n",
 144               n, n->magic, OSET_MAGIC);
 145    return n;
 146 }
 147
 148 // Given an AvlNode, return the pointer to the element.
 149 static inline
 150 void* elem_of_node(const AvlNode *n)
 151 {
 152    vg_assert2(n->magic == OSET_MAGIC,
 153               "bad magic on node %p = %x (expected %x)\n"
 154               "possible causes:\n"
 155               " - node metadata corrupted by overwriting end of element?\n",
 156               n, n->magic, OSET_MAGIC);
 157    return (void*)((Addr)n + sizeof(AvlNode));
 158 }
 159
 160 // Like elem_of_node, but no magic checking.
 161 static inline
 162 void* elem_of_node_no_check(const AvlNode *n)
 163 {
 164    return (void*)((Addr)n + sizeof(AvlNode));
 165 }
 166
 167 static inline
 168 void* slow_key_of_node(const AvlTree* t, const AvlNode* n)
 169 {
 170    return (void*)((Addr)elem_of_node(n) + t->keyOff);
 171 }
 172
 173 static inline
 174 void* fast_key_of_node(const AvlNode* n)
 175 {
 176    return elem_of_node(n);
 177 }
 178
 179 // Compare the first word of each element.  Inlining is *crucial*.
 180 static inline Word fast_cmp(const void* k, const AvlNode* n)
 181 {
 182    UWord w1 = *(const UWord*)k;
 183    UWord w2 = *(const UWord*)elem_of_node(n);
 184    // In previous versions, we tried to do this faster by doing
 185    // "return w1 - w2".  But it didn't work reliably, because the
 186    // complete result of subtracting two N-bit numbers is an N+1-bit
 187    // number, and what the caller is interested in is the sign of
 188    // the complete N+1-bit result.  The branching version is slightly
 189    // slower, but safer and easier to understand.
 190    if (w1 > w2) return  1;
 191    if (w1 < w2) return -1;
 192    return 0;
 193 }
 194
 195 // Compare a key and an element.  Inlining is *crucial*.
 196 static
 197 inline Word slow_cmp(const AvlTree* t, const void* k, const AvlNode* n)
 198 {
 199    return t->cmp(k, elem_of_node(n));
 200 }
 201
 202
 203 // Swing to the left.   Warning: no balance maintainance.
 204 static void avl_swl ( AvlNode** root )
 205 {
 206    AvlNode* a = *root;
 207    AvlNode* b = a->right;
 208    *root    = b;
 209    a->right = b->left;
 210    b->left  = a;
 211 }
 212
 213 // Swing to the right.  Warning: no balance maintainance.
 214 static void avl_swr ( AvlNode** root )
 215 {
 216    AvlNode* a = *root;
 217    AvlNode* b = a->left;
 218    *root    = b;
 219    a->left  = b->right;
 220    b->right = a;
 221 }
 222
 223 // Balance maintainance after especially nasty swings.
 224 static void avl_nasty ( AvlNode* root )
 225 {
 226    switch (root->balance) {
 227    case -1:
 228       root->left->balance  = 0;
 229       root->right->balance = 1;
 230       break;
 231    case 1:
 232       root->left->balance  =-1;
 233       root->right->balance = 0;
 234       break;
 235    case 0:
 236       root->left->balance  = 0;
 237       root->right->balance = 0;
 238    }
 239    root->balance = 0;
 240 }
 241
 242
 243 // Clear the iterator stack.
 244 static void stackClear(AvlTree* t)
 245 {
 246    Int i;
 247    vg_assert(t);
 248    for (i = 0; i < STACK_MAX; i++) {
 249       t->nodeStack[i] = NULL;
 250       t->numStack[i]  = 0;
 251    }
 252    t->stackTop = 0;
 253 }
 254
 255 // Push onto the iterator stack.
 256 static inline void stackPush(AvlTree* t, AvlNode* n, Int i)
 257 {
 258    vg_assert(t->stackTop < STACK_MAX);
 259    vg_assert(1 <= i && i <= 3);
 260    t->nodeStack[t->stackTop] = n;
 261    t-> numStack[t->stackTop] = i;
 262    t->stackTop++;
 263 }
 264
 265 // Pop from the iterator stack.
 266 static inline Bool stackPop(AvlTree* t, AvlNode** n, Int* i)
 267 {
 268    vg_assert(t->stackTop <= STACK_MAX);
 269
 270    if (t->stackTop > 0) {
 271       t->stackTop--;
 272       *n = t->nodeStack[t->stackTop];
 273       *i = t-> numStack[t->stackTop];
 274       vg_assert(1 <= *i && *i <= 3);
 275       t->nodeStack[t->stackTop] = NULL;
 276       t-> numStack[t->stackTop] = 0;
 277       return True;
 278    } else {
 279       return False;
 280    }
 281 }
 282
 283 /*--------------------------------------------------------------------*/
 284 /*--- Creating and destroying AvlTrees and AvlNodes                ---*/
 285 /*--------------------------------------------------------------------*/
 286
 287 // The underscores avoid GCC complaints about overshadowing global names.
 288 AvlTree* VG_(OSetGen_Create)(PtrdiffT keyOff, OSetCmp_t cmp,
 289                              OSetAlloc_t alloc_fn, const HChar* cc,
 290                              OSetFree_t free_fn)
 291 {
 292    AvlTree* t;
 293
 294    // Check the padding is right and the AvlNode is the expected size.
 295    vg_assert(sizeof(AvlNode) == 3*sizeof(void*));
 296
 297    // Sanity check args
 298    vg_assert(alloc_fn);
 299    vg_assert(free_fn);
 300    if (!cmp) vg_assert(0 == keyOff);    // If no cmp, offset must be zero
 301
 302    t           = alloc_fn(cc, sizeof(AvlTree));
 303    t->keyOff   = keyOff;
 304    t->cmp      = cmp;
 305    t->alloc_fn = alloc_fn;
 306    t->cc       = cc;
 307    t->free_fn  = free_fn;
 308    t->node_pa  = NULL;
 309    t->maxEltSize = 0; // Just in case it would be wrongly used.
 310    t->nElems   = 0;
 311    t->root     = NULL;
 312    stackClear(t);
 313
 314    return t;
 315 }
 316
 317 AvlTree* VG_(OSetGen_Create_With_Pool)(PtrdiffT keyOff, OSetCmp_t cmp,
 318                                        OSetAlloc_t alloc_fn, const HChar* cc,
 319                                        OSetFree_t free_fn,
 320                                        SizeT poolSize,
 321                                        SizeT maxEltSize)
 322 {
 323    AvlTree* t;
 324
 325    t = VG_(OSetGen_Create) (keyOff, cmp, alloc_fn, cc, free_fn);
 326
 327    vg_assert (poolSize > 0);
 328    vg_assert (maxEltSize > 0);
 329    t->maxEltSize = maxEltSize;
 330    t->node_pa = VG_(newPA)(sizeof(AvlNode)
 331                            + VG_ROUNDUP(maxEltSize, sizeof(void*)),
 332                            poolSize,
 333                            t->alloc_fn,
 334                            cc,
 335                            t->free_fn);
 336    VG_(addRefPA) (t->node_pa);
 337
 338    return t;
 339 }
 340
 341 AvlTree* VG_(OSetGen_EmptyClone) (const AvlTree* os)
 342 {
 343    AvlTree* t;
 344
 345    vg_assert(os);
 346
 347    t           = os->alloc_fn(os->cc, sizeof(AvlTree));
 348    t->keyOff   = os->keyOff;
 349    t->cmp      = os->cmp;
 350    t->alloc_fn = os->alloc_fn;
 351    t->cc       = os->cc;
 352    t->free_fn  = os->free_fn;
 353    t->node_pa  = os->node_pa;
 354    if (t->node_pa)
 355       VG_(addRefPA) (t->node_pa);
 356    t->maxEltSize = os->maxEltSize;
 357    t->nElems   = 0;
 358    t->root     = NULL;
 359    stackClear(t);
 360
 361    return t;
 362 }
 363
 364 AvlTree* VG_(OSetWord_Create)(OSetAlloc_t alloc_fn, const HChar* cc,
 365                               OSetFree_t free_fn)
 366 {
 367    return VG_(OSetGen_Create)(/*keyOff*/0, /*cmp*/NULL, alloc_fn, cc, free_fn);
 368 }
 369
 370 // Destructor, frees up all memory held by remaining nodes.
 371 void VG_(OSetGen_Destroy)(AvlTree* t)
 372 {
 373    Bool has_node_pa;
 374    vg_assert(t);
 375
 376    has_node_pa = t->node_pa != NULL;
 377
 378    /*
 379     * If we are the only remaining user of this pool allocator, release all
 380     * the elements by deleting the pool allocator. That's more efficient than
 381     * deleting tree nodes one by one.
 382     */
 383    if (!has_node_pa || VG_(releasePA)(t->node_pa) > 0) {
 384       AvlNode* n = NULL;
 385       Int i = 0;
 386       Word sz = 0;
 387
 388       stackClear(t);
 389       if (t->root)
 390          stackPush(t, t->root, 1);
 391
 392       /* Free all the AvlNodes.  This is a post-order traversal, because we */
 393       /* must free all children of a node before the node itself. */
 394       while (stackPop(t, &n, &i)) {
 395          switch (i) {
 396          case 1:
 397             stackPush(t, n, 2);
 398             if (n->left)  stackPush(t, n->left, 1);
 399             break;
 400          case 2:
 401             stackPush(t, n, 3);
 402             if (n->right) stackPush(t, n->right, 1);
 403             break;
 404          case 3:
 405             if (has_node_pa)
 406                VG_(freeEltPA) (t->node_pa, n);
 407             else
 408                t->free_fn(n);
 409             sz++;
 410             break;
 411          }
 412       }
 413       vg_assert(sz == t->nElems);
 414    }
 415
 416    /* Free the AvlTree itself. */
 417    t->free_fn(t);
 418 }
 419
 420 void VG_(OSetWord_Destroy)(AvlTree* t)
 421 {
 422    VG_(OSetGen_Destroy)(t);
 423 }
 424
 425 // Allocate and initialise a new node.
 426 void* VG_(OSetGen_AllocNode)(const AvlTree* t, SizeT elemSize)
 427 {
 428    AvlNode* n;
 429    Int nodeSize = sizeof(AvlNode) + elemSize;
 430    vg_assert(elemSize > 0);
 431    if (t->node_pa) {
 432       vg_assert(elemSize <= t->maxEltSize);
 433       n = VG_(allocEltPA) (t->node_pa);
 434    } else {
 435       n = t->alloc_fn( t->cc, nodeSize );
 436    }
 437    VG_(memset)(n, 0, nodeSize);
 438    n->magic = OSET_MAGIC;
 439    return elem_of_node(n);
 440 }
 441
 442 void VG_(OSetGen_FreeNode)(const AvlTree* t, void* e)
 443 {
 444    if (t->node_pa)
 445       VG_(freeEltPA) (t->node_pa, node_of_elem (e));
 446    else
 447       t->free_fn( node_of_elem(e) );
 448 }
 449
 450 /*--------------------------------------------------------------------*/
 451 /*--- Insertion                                                    ---*/
 452 /*--------------------------------------------------------------------*/
 453
 454 static inline Word cmp_key_root(const AvlTree* t, const AvlNode* n)
 455 {
 456    return t->cmp
 457           ? slow_cmp(t, slow_key_of_node(t, n), t->root)
 458           : fast_cmp(   fast_key_of_node(   n), t->root);
 459 }
 460
 461 // Insert element e into the non-empty AVL tree t.
 462 // Returns True if the depth of the tree has grown.
 463 static Bool avl_insert(AvlTree* t, AvlNode* n)
 464 {
 465    Word cmpres = cmp_key_root(t, n);
 466
 467    if (cmpres < 0) {
 468       // Insert into the left subtree.
 469       if (t->root->left) {
 470          // Only need to set the used fields in the subtree.
 471          AvlTree left_subtree;
 472          left_subtree.root   = t->root->left;
 473          left_subtree.cmp    = t->cmp;
 474          left_subtree.keyOff = t->keyOff;
 475          if (avl_insert(&left_subtree, n)) {
 476              switch (t->root->balance--) {
 477              case 1: return False;
 478              case 0: return True;
 479              }
 480              if (t->root->left->balance < 0) {
 481                 avl_swr(&(t->root));
 482                 t->root->balance = 0;
 483                 t->root->right->balance = 0;
 484              } else {
 485                 avl_swl(&(t->root->left));
 486                 avl_swr(&(t->root));
 487                 avl_nasty(t->root);
 488              }
 489          } else {
 490             t->root->left=left_subtree.root;
 491          }
 492          return False;
 493       } else {
 494          t->root->left = n;
 495          if (t->root->balance--) return False;
 496          return True;
 497       }
 498
 499    } else if (cmpres > 0) {
 500       // Insert into the right subtree
 501       if (t->root->right) {
 502          // Only need to set the used fields in the subtree.
 503          AvlTree right_subtree;
 504          right_subtree.root   = t->root->right;
 505          right_subtree.cmp    = t->cmp;
 506          right_subtree.keyOff = t->keyOff;
 507          if (avl_insert(&right_subtree, n)) {
 508             switch (t->root->balance++) {
 509             case -1: return False;
 510             case  0: return True;
 511             }
 512             if (t->root->right->balance > 0) {
 513                avl_swl(&(t->root));
 514                t->root->balance = 0;
 515                t->root->left->balance = 0;
 516             } else {
 517                avl_swr(&(t->root->right));
 518                avl_swl(&(t->root));
 519                avl_nasty(t->root);
 520             }
 521          } else {
 522             t->root->right=right_subtree.root;
 523          }
 524          return False;
 525       } else {
 526          t->root->right = n;
 527          if (t->root->balance++) return False;
 528          return True;
 529       }
 530
 531    } else {
 532       vg_assert2(0, "OSet{Word,Gen}_Insert: duplicate element added");
 533    }
 534 }
 535
 536 // Insert element e into the AVL tree t.  This is just a wrapper for
 537 // avl_insert() which doesn't return a Bool.
 538 void VG_(OSetGen_Insert)(AvlTree* t, void* e)
 539 {
 540    AvlNode* n;
 541
 542    vg_assert(t);
 543
 544    // Initialise.  Even though OSetGen_AllocNode zeroes these fields,
 545    // we should do it again in case a node is removed and then
 546    // re-added to the tree.
 547    n          = node_of_elem(e);
 548    n->left    = 0;
 549    n->right   = 0;
 550    n->balance = 0;
 551
 552    // Insert into an empty tree
 553    if (!t->root) {
 554       t->root = n;
 555    } else {
 556       avl_insert(t, n);
 557    }
 558
 559    t->nElems++;
 560    t->stackTop = 0;  // So the iterator can't get out of sync
 561 }
 562
 563 void VG_(OSetWord_Insert)(AvlTree* t, UWord val)
 564 {
 565    Word* node = VG_(OSetGen_AllocNode)(t, sizeof(UWord));
 566    *node = val;
 567    VG_(OSetGen_Insert)(t, node);
 568 }
 569
 570 /*--------------------------------------------------------------------*/
 571 /*--- Lookup                                                       ---*/
 572 /*--------------------------------------------------------------------*/
 573
 574 // Find the *node* in t matching k, or NULL if not found.
 575 static AvlNode* avl_lookup(const AvlTree* t, const void* k)
 576 {
 577    Word     cmpres;
 578    AvlNode* curr = t->root;
 579
 580    if (t->cmp) {
 581       // General case
 582       while (True) {
 583          if (curr == NULL) return NULL;
 584          cmpres = slow_cmp(t, k, curr);
 585          if      (cmpres < 0) curr = curr->left;
 586          else if (cmpres > 0) curr = curr->right;
 587          else return curr;
 588       }
 589    } else {
 590       // Fast-track special case.  We use the no-check version of
 591       // elem_of_node because it saves about 10% on lookup time.  This
 592       // shouldn't be very dangerous because each node will have been
 593       // checked on insertion.
 594       UWord w1 = *(const UWord*)k;
 595       UWord w2;
 596       while (True) {
 597          if (curr == NULL) return NULL;
 598          w2 = *(UWord*)elem_of_node_no_check(curr);
 599          if      (w1 < w2) curr = curr->left;
 600          else if (w1 > w2) curr = curr->right;
 601          else return curr;
 602       }
 603    }
 604 }
 605
 606 // Find the *element* in t matching k, or NULL if not found.
 607 void* VG_(OSetGen_Lookup)(const AvlTree* t, const void* k)
 608 {
 609    AvlNode* n;
 610    vg_assert(t);
 611    n = avl_lookup(t, k);
 612    return ( n ? elem_of_node(n) : NULL );
 613 }
 614
 615 // Find the *element* in t matching k, or NULL if not found;  use the given
 616 // comparison function rather than the standard one.
 617 void* VG_(OSetGen_LookupWithCmp)(AvlTree* t, const void* k, OSetCmp_t cmp)
 618 {
 619    // Save the normal one to the side, then restore once we're done.
 620    void* e;
 621    OSetCmp_t tmpcmp;
 622    vg_assert(t);
 623    tmpcmp = t->cmp;
 624    t->cmp = cmp;
 625    e = VG_(OSetGen_Lookup)(t, k);
 626    t->cmp = tmpcmp;
 627    return e;
 628 }
 629
 630 // Is there an element matching k?
 631 Bool VG_(OSetGen_Contains)(const AvlTree* t, const void* k)
 632 {
 633    return (NULL != VG_(OSetGen_Lookup)(t, k));
 634 }
 635
 636 Bool VG_(OSetWord_Contains)(const AvlTree* t, UWord val)
 637 {
 638    return (NULL != VG_(OSetGen_Lookup)(t, &val));
 639 }
 640
 641 /*--------------------------------------------------------------------*/
 642 /*--- Deletion                                                     ---*/
 643 /*--------------------------------------------------------------------*/
 644
 645 static Bool avl_removeroot(AvlTree* t);
 646
 647 // Remove an already-selected node n from the AVL tree t.
 648 // Returns True if the depth of the tree has shrunk.
 649 static Bool avl_remove(AvlTree* t, const AvlNode* n)
 650 {
 651    Bool ch;
 652    Word cmpres = cmp_key_root(t, n);
 653
 654    if (cmpres < 0) {
 655       AvlTree left_subtree;
 656       // Remove from the left subtree
 657       vg_assert(t->root->left);
 658       // Only need to set the used fields in the subtree.
 659       left_subtree.root   = t->root->left;
 660       left_subtree.cmp    = t->cmp;
 661       left_subtree.keyOff = t->keyOff;
 662       ch = avl_remove(&left_subtree, n);
 663       t->root->left = left_subtree.root;
 664       if (ch) {
 665          switch (t->root->balance++) {
 666          case -1: return True;
 667          case  0: return False;
 668          }
 669          switch (t->root->right->balance) {
 670          case 0:
 671             avl_swl(&(t->root));
 672             t->root->balance = -1;
 673             t->root->left->balance = 1;
 674             return False;
 675          case 1:
 676             avl_swl(&(t->root));
 677             t->root->balance = 0;
 678             t->root->left->balance = 0;
 679             return True;
 680          }
 681          avl_swr(&(t->root->right));
 682          avl_swl(&(t->root));
 683          avl_nasty(t->root);
 684          return True;
 685       } else {
 686          return False;
 687       }
 688
 689    } else if (cmpres > 0) {
 690       // Remove from the right subtree
 691       AvlTree right_subtree;
 692       vg_assert(t->root->right);
 693       // Only need to set the used fields in the subtree.
 694       right_subtree.root   = t->root->right;
 695       right_subtree.cmp    = t->cmp;
 696       right_subtree.keyOff = t->keyOff;
 697       ch = avl_remove(&right_subtree, n);
 698       t->root->right = right_subtree.root;
 699       if (ch) {
 700          switch (t->root->balance--) {
 701          case 1: return True;
 702          case 0: return False;
 703          }
 704          switch (t->root->left->balance) {
 705          case 0:
 706             avl_swr(&(t->root));
 707             t->root->balance = 1;
 708             t->root->right->balance = -1;
 709             return False;
 710          case -1:
 711             avl_swr(&(t->root));
 712             t->root->balance = 0;
 713             t->root->right->balance = 0;
 714             return True;
 715          }
 716          avl_swl(&(t->root->left));
 717          avl_swr(&(t->root));
 718          avl_nasty(t->root);
 719          return True;
 720       } else {
 721          return False;
 722       }
 723
 724    } else {
 725       // Found the node to be removed.
 726       vg_assert(t->root == n);
 727       return avl_removeroot(t);
 728    }
 729 }
 730
 731 // Remove the root of the AVL tree t.
 732 // Returns True if the depth of the tree has shrunk.
 733 static Bool avl_removeroot(AvlTree* t)
 734 {
 735    Bool     ch;
 736    AvlNode* n;
 737
 738    if (!t->root->left) {
 739       if (!t->root->right) {
 740          t->root = NULL;
 741          return True;
 742       }
 743       t->root = t->root->right;
 744       return True;
 745    }
 746    if (!t->root->right) {
 747       t->root = t->root->left;
 748       return True;
 749    }
 750    if (t->root->balance < 0) {
 751       // Remove from the left subtree
 752       n = t->root->left;
 753       while (n->right) n = n->right;
 754    } else {
 755       // Remove from the right subtree
 756       n = t->root->right;
 757       while (n->left) n = n->left;
 758    }
 759    ch = avl_remove(t, n);
 760    n->left    = t->root->left;
 761    n->right   = t->root->right;
 762    n->balance = t->root->balance;
 763    t->root    = n;
 764    if (n->balance == 0) return ch;
 765    return False;
 766 }
 767
 768 // Remove and return the element matching the key 'k', or NULL
 769 // if not present.
 770 void* VG_(OSetGen_Remove)(AvlTree* t, const void* k)
 771 {
 772    // Have to find the node first, then remove it.
 773    AvlNode* n = avl_lookup(t, k);
 774    if (n) {
 775       avl_remove(t, n);
 776       t->nElems--;
 777       t->stackTop = 0;     // So the iterator can't get out of sync
 778       return elem_of_node(n);
 779    } else {
 780       return NULL;
 781    }
 782 }
 783
 784 Bool VG_(OSetWord_Remove)(AvlTree* t, UWord val)
 785 {
 786    void* n = VG_(OSetGen_Remove)(t, &val);
 787    if (n) {
 788       VG_(OSetGen_FreeNode)(t, n);
 789       return True;
 790    } else {
 791       return False;
 792    }
 793 }
 794
 795 /*--------------------------------------------------------------------*/
 796 /*--- Iterator                                                     ---*/
 797 /*--------------------------------------------------------------------*/
 798
 799 // The iterator is implemented using in-order traversal with an explicit
 800 // stack, which lets us do the traversal one step at a time and remember
 801 // where we are between each call to OSetGen_Next().
 802
 803 void VG_(OSetGen_ResetIter)(AvlTree* t)
 804 {
 805    vg_assert(t);
 806    stackClear(t);
 807    if (t->root)
 808       stackPush(t, t->root, 1);
 809 }
 810
 811 void VG_(OSetWord_ResetIter)(AvlTree* t)
 812 {
 813    VG_(OSetGen_ResetIter)(t);
 814 }
 815
 816 void* VG_(OSetGen_Next)(AvlTree* t)
 817 {
 818    Int i = 0;
 819    OSetNode* n = NULL;
 820
 821    vg_assert(t);
 822
 823    // This in-order traversal requires each node to be pushed and popped
 824    // three times.  These could be avoided by updating nodes in-situ on the
 825    // top of the stack, but the push/pop cost is so small that it's worth
 826    // keeping this loop in this simpler form.
 827    while (stackPop(t, &n, &i)) {
 828       switch (i) {
 829       case 1: case_1:
 830          stackPush(t, n, 2);
 831          /* if (n->left)  stackPush(t, n->left, 1); */
 832          if (n->left) { n = n->left; goto case_1; }
 833          break;
 834       case 2:
 835          stackPush(t, n, 3);
 836          return elem_of_node(n);
 837       case 3:
 838          /* if (n->right) stackPush(t, n->right, 1); */
 839          if (n->right) { n = n->right; goto case_1; }
 840          break;
 841       }
 842    }
 843
 844    // Stack empty, iterator is exhausted, return NULL
 845    return NULL;
 846 }
 847
 848 Bool VG_(OSetWord_Next)(AvlTree* t, UWord* val)
 849 {
 850    UWord* n = VG_(OSetGen_Next)(t);
 851    if (n) {
 852       *val = *n;
 853       return True;
 854    } else {
 855       return False;
 856    }
 857 }
 858
 859 // set up 'oset' for iteration so that the first key subsequently
 860 // produced VG_(OSetGen_Next) is the smallest key in the map
 861 // >= start_at.  Naturally ">=" is defined by the comparison
 862 // function supplied to VG_(OSetGen_Create).
 863 void VG_(OSetGen_ResetIterAt)(AvlTree* oset, const void* k)
 864 {
 865    AvlNode *t;
 866    Word    cmpresS; /* signed */
 867    UWord   cmpresU; /* unsigned */
 868
 869    vg_assert(oset);
 870    stackClear(oset);
 871
 872    if (!oset->root)
 873       return;
 874
 875    // We need to do regular search and fill in the stack.
 876    t = oset->root;
 877
 878    while (True) {
 879       if (t == NULL) return;
 880
 881       if (oset->cmp) {
 882          cmpresS = (Word)slow_cmp(oset, k, t);
 883       } else {
 884          cmpresS = fast_cmp(k, t);
 885       }
 886
 887       /* Switch the sense of the comparison, since the comparison
 888          order of args (k vs t) above is opposite to that of the
 889          corresponding code in hg_wordfm.c. */
 890       if (cmpresS < 0) { cmpresS = 1; }
 891       else if (cmpresS > 0) { cmpresS = -1; }
 892
 893       if (cmpresS == 0) {
 894          // We found the exact key -- we are done.
 895          // The iteration should start with this node.
 896          stackPush(oset, t, 2);
 897          // The stack now looks like {2, 2, ... ,2, 2}
 898          return;
 899       }
 900       cmpresU = (UWord)cmpresS;
 901       cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
 902       vg_assert(cmpresU == 0 || cmpresU == 1);
 903       if (!cmpresU) {
 904          // Push this node only if we go to the left child.
 905          stackPush(oset, t, 2);
 906       }
 907       t = cmpresU==0 ? t->left : t->right;
 908    }
 909 }
 910
 911 /*--------------------------------------------------------------------*/
 912 /*--- Miscellaneous operations                                     ---*/
 913 /*--------------------------------------------------------------------*/
 914
 915 Word VG_(OSetGen_Size)(const AvlTree* t)
 916 {
 917    vg_assert(t);
 918    return t->nElems;
 919 }
 920
 921 Word VG_(OSetWord_Size)(const AvlTree* t)
 922 {
 923    return VG_(OSetGen_Size)(t);
 924 }
 925
 926 static void OSet_Print2( const AvlTree* t, const AvlNode* n,
 927                          const HChar*(*strElem)(const void *), Int p )
 928 {
 929    // This is a recursive in-order traversal.
 930    Int q = p;
 931    if (NULL == n) return;
 932    if (n->right) OSet_Print2(t, n->right, strElem, p+1);
 933    while (q--) VG_(printf)(".. ");
 934    VG_(printf)("%s\n", strElem(elem_of_node(n)));
 935    if (n->left) OSet_Print2(t, n->left, strElem, p+1);
 936 }
 937
 938 __attribute__((unused))
 939 static void OSet_Print( const AvlTree* t, const HChar *where,
 940                         const HChar*(*strElem)(const void *) )
 941 {
 942    VG_(printf)("-- start %s ----------------\n", where);
 943    OSet_Print2(t, t->root, strElem, 0);
 944    VG_(printf)("-- end   %s ----------------\n", where);
 945 }
 946
 947 /*--------------------------------------------------------------------*/
 948 /*--- end                                                          ---*/
 949 /*--------------------------------------------------------------------*/