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