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 inline 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    vg_assert(t);
 608    if (LIKELY(t->root == NULL))
 609       return NULL;
 610    AvlNode* n = avl_lookup(t, k);
 611    return ( n ? elem_of_node(n) : NULL );
 612 }
 613
 614 // Find the *element* in t matching k, or NULL if not found;  use the given
 615 // comparison function rather than the standard one.
 616 void* VG_(OSetGen_LookupWithCmp)(AvlTree* t, const void* k, OSetCmp_t cmp)
 617 {
 618    // Save the normal one to the side, then restore once we're done.
 619    void* e;
 620    OSetCmp_t tmpcmp;
 621    vg_assert(t);
 622    tmpcmp = t->cmp;
 623    t->cmp = cmp;
 624    e = VG_(OSetGen_Lookup)(t, k);
 625    t->cmp = tmpcmp;
 626    return e;
 627 }
 628
 629 // Is there an element matching k?
 630 Bool VG_(OSetGen_Contains)(const AvlTree* t, const void* k)
 631 {
 632    return (NULL != VG_(OSetGen_Lookup)(t, k));
 633 }
 634
 635 Bool VG_(OSetWord_Contains)(const AvlTree* t, UWord val)
 636 {
 637    return (NULL != VG_(OSetGen_Lookup)(t, &val));
 638 }
 639
 640 /*--------------------------------------------------------------------*/
 641 /*--- Deletion                                                     ---*/
 642 /*--------------------------------------------------------------------*/
 643
 644 static Bool avl_removeroot(AvlTree* t);
 645
 646 // Remove an already-selected node n from the AVL tree t.
 647 // Returns True if the depth of the tree has shrunk.
 648 static Bool avl_remove(AvlTree* t, const AvlNode* n)
 649 {
 650    Bool ch;
 651    Word cmpres = cmp_key_root(t, n);
 652
 653    if (cmpres < 0) {
 654       AvlTree left_subtree;
 655       // Remove from the left subtree
 656       vg_assert(t->root->left);
 657       // Only need to set the used fields in the subtree.
 658       left_subtree.root   = t->root->left;
 659       left_subtree.cmp    = t->cmp;
 660       left_subtree.keyOff = t->keyOff;
 661       ch = avl_remove(&left_subtree, n);
 662       t->root->left = left_subtree.root;
 663       if (ch) {
 664          switch (t->root->balance++) {
 665          case -1: return True;
 666          case  0: return False;
 667          }
 668          switch (t->root->right->balance) {
 669          case 0:
 670             avl_swl(&(t->root));
 671             t->root->balance = -1;
 672             t->root->left->balance = 1;
 673             return False;
 674          case 1:
 675             avl_swl(&(t->root));
 676             t->root->balance = 0;
 677             t->root->left->balance = 0;
 678             return True;
 679          }
 680          avl_swr(&(t->root->right));
 681          avl_swl(&(t->root));
 682          avl_nasty(t->root);
 683          return True;
 684       } else {
 685          return False;
 686       }
 687
 688    } else if (cmpres > 0) {
 689       // Remove from the right subtree
 690       AvlTree right_subtree;
 691       vg_assert(t->root->right);
 692       // Only need to set the used fields in the subtree.
 693       right_subtree.root   = t->root->right;
 694       right_subtree.cmp    = t->cmp;
 695       right_subtree.keyOff = t->keyOff;
 696       ch = avl_remove(&right_subtree, n);
 697       t->root->right = right_subtree.root;
 698       if (ch) {
 699          switch (t->root->balance--) {
 700          case 1: return True;
 701          case 0: return False;
 702          }
 703          switch (t->root->left->balance) {
 704          case 0:
 705             avl_swr(&(t->root));
 706             t->root->balance = 1;
 707             t->root->right->balance = -1;
 708             return False;
 709          case -1:
 710             avl_swr(&(t->root));
 711             t->root->balance = 0;
 712             t->root->right->balance = 0;
 713             return True;
 714          }
 715          avl_swl(&(t->root->left));
 716          avl_swr(&(t->root));
 717          avl_nasty(t->root);
 718          return True;
 719       } else {
 720          return False;
 721       }
 722
 723    } else {
 724       // Found the node to be removed.
 725       vg_assert(t->root == n);
 726       return avl_removeroot(t);
 727    }
 728 }
 729
 730 // Remove the root of the AVL tree t.
 731 // Returns True if the depth of the tree has shrunk.
 732 static Bool avl_removeroot(AvlTree* t)
 733 {
 734    Bool     ch;
 735    AvlNode* n;
 736
 737    if (!t->root->left) {
 738       if (!t->root->right) {
 739          t->root = NULL;
 740          return True;
 741       }
 742       t->root = t->root->right;
 743       return True;
 744    }
 745    if (!t->root->right) {
 746       t->root = t->root->left;
 747       return True;
 748    }
 749    if (t->root->balance < 0) {
 750       // Remove from the left subtree
 751       n = t->root->left;
 752       while (n->right) n = n->right;
 753    } else {
 754       // Remove from the right subtree
 755       n = t->root->right;
 756       while (n->left) n = n->left;
 757    }
 758    ch = avl_remove(t, n);
 759    n->left    = t->root->left;
 760    n->right   = t->root->right;
 761    n->balance = t->root->balance;
 762    t->root    = n;
 763    if (n->balance == 0) return ch;
 764    return False;
 765 }
 766
 767 // Remove and return the element matching the key 'k', or NULL
 768 // if not present.
 769 void* VG_(OSetGen_Remove)(AvlTree* t, const void* k)
 770 {
 771    if (LIKELY(t->root == NULL))
 772       return NULL;
 773    // Have to find the node first, then remove it.
 774    AvlNode* n = avl_lookup(t, k);
 775    if (n) {
 776       avl_remove(t, n);
 777       t->nElems--;
 778       t->stackTop = 0;     // So the iterator can't get out of sync
 779       return elem_of_node(n);
 780    } else {
 781       return NULL;
 782    }
 783 }
 784
 785 Bool VG_(OSetWord_Remove)(AvlTree* t, UWord val)
 786 {
 787    void* n = VG_(OSetGen_Remove)(t, &val);
 788    if (n) {
 789       VG_(OSetGen_FreeNode)(t, n);
 790       return True;
 791    } else {
 792       return False;
 793    }
 794 }
 795
 796 /*--------------------------------------------------------------------*/
 797 /*--- Iterator                                                     ---*/
 798 /*--------------------------------------------------------------------*/
 799
 800 // The iterator is implemented using in-order traversal with an explicit
 801 // stack, which lets us do the traversal one step at a time and remember
 802 // where we are between each call to OSetGen_Next().
 803
 804 void VG_(OSetGen_ResetIter)(AvlTree* t)
 805 {
 806    vg_assert(t);
 807    stackClear(t);
 808    if (t->root)
 809       stackPush(t, t->root, 1);
 810 }
 811
 812 void VG_(OSetWord_ResetIter)(AvlTree* t)
 813 {
 814    VG_(OSetGen_ResetIter)(t);
 815 }
 816
 817 void* VG_(OSetGen_Next)(AvlTree* t)
 818 {
 819    Int i = 0;
 820    OSetNode* n = NULL;
 821
 822    vg_assert(t);
 823
 824    // This in-order traversal requires each node to be pushed and popped
 825    // three times.  These could be avoided by updating nodes in-situ on the
 826    // top of the stack, but the push/pop cost is so small that it's worth
 827    // keeping this loop in this simpler form.
 828    while (stackPop(t, &n, &i)) {
 829       switch (i) {
 830       case 1: case_1:
 831          stackPush(t, n, 2);
 832          /* if (n->left)  stackPush(t, n->left, 1); */
 833          if (n->left) { n = n->left; goto case_1; }
 834          break;
 835       case 2:
 836          stackPush(t, n, 3);
 837          return elem_of_node(n);
 838       case 3:
 839          /* if (n->right) stackPush(t, n->right, 1); */
 840          if (n->right) { n = n->right; goto case_1; }
 841          break;
 842       }
 843    }
 844
 845    // Stack empty, iterator is exhausted, return NULL
 846    return NULL;
 847 }
 848
 849 Bool VG_(OSetWord_Next)(AvlTree* t, UWord* val)
 850 {
 851    UWord* n = VG_(OSetGen_Next)(t);
 852    if (n) {
 853       *val = *n;
 854       return True;
 855    } else {
 856       return False;
 857    }
 858 }
 859
 860 // set up 'oset' for iteration so that the first key subsequently
 861 // produced VG_(OSetGen_Next) is the smallest key in the map
 862 // >= start_at.  Naturally ">=" is defined by the comparison
 863 // function supplied to VG_(OSetGen_Create).
 864 void VG_(OSetGen_ResetIterAt)(AvlTree* oset, const void* k)
 865 {
 866    AvlNode *t;
 867    Word    cmpresS; /* signed */
 868    UWord   cmpresU; /* unsigned */
 869
 870    vg_assert(oset);
 871    stackClear(oset);
 872
 873    if (!oset->root)
 874       return;
 875
 876    // We need to do regular search and fill in the stack.
 877    t = oset->root;
 878
 879    while (True) {
 880       if (t == NULL) return;
 881
 882       if (oset->cmp) {
 883          cmpresS = (Word)slow_cmp(oset, k, t);
 884       } else {
 885          cmpresS = fast_cmp(k, t);
 886       }
 887
 888       /* Switch the sense of the comparison, since the comparison
 889          order of args (k vs t) above is opposite to that of the
 890          corresponding code in hg_wordfm.c. */
 891       if (cmpresS < 0) { cmpresS = 1; }
 892       else if (cmpresS > 0) { cmpresS = -1; }
 893
 894       if (cmpresS == 0) {
 895          // We found the exact key -- we are done.
 896          // The iteration should start with this node.
 897          stackPush(oset, t, 2);
 898          // The stack now looks like {2, 2, ... ,2, 2}
 899          return;
 900       }
 901       cmpresU = (UWord)cmpresS;
 902       cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
 903       vg_assert(cmpresU == 0 || cmpresU == 1);
 904       if (!cmpresU) {
 905          // Push this node only if we go to the left child.
 906          stackPush(oset, t, 2);
 907       }
 908       t = cmpresU==0 ? t->left : t->right;
 909    }
 910 }
 911
 912 /*--------------------------------------------------------------------*/
 913 /*--- Miscellaneous operations                                     ---*/
 914 /*--------------------------------------------------------------------*/
 915
 916 UInt VG_(OSetGen_Size)(const AvlTree* t)
 917 {
 918    vg_assert(t);
 919    return t->nElems;
 920 }
 921
 922 Word VG_(OSetWord_Size)(const AvlTree* t)
 923 {
 924    return VG_(OSetGen_Size)(t);
 925 }
 926
 927 static void OSet_Print2( const AvlTree* t, const AvlNode* n,
 928                          const HChar*(*strElem)(const void *), Int p )
 929 {
 930    // This is a recursive in-order traversal.
 931    Int q = p;
 932    if (NULL == n) return;
 933    if (n->right) OSet_Print2(t, n->right, strElem, p+1);
 934    while (q--) VG_(printf)(".. ");
 935    VG_(printf)("%s\n", strElem(elem_of_node(n)));
 936    if (n->left) OSet_Print2(t, n->left, strElem, p+1);
 937 }
 938
 939 __attribute__((unused))
 940 static void OSet_Print( const AvlTree* t, const HChar *where,
 941                         const HChar*(*strElem)(const void *) )
 942 {
 943    VG_(printf)("-- start %s ----------------\n", where);
 944    OSet_Print2(t, t->root, strElem, 0);
 945    VG_(printf)("-- end   %s ----------------\n", where);
 946 }
 947
 948 /*--------------------------------------------------------------------*/
 949 /*--- end                                                          ---*/
 950 /*--------------------------------------------------------------------*/