Home | History | Annotate | Download | only in coregrind
      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-2012 Nicholas Nethercote
     11       njn (at) 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_tool_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;      // allocator
    116    HChar* cc;              // cc for allocator
    117    OSetFree_t  free;       // 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(AvlTree* t, AvlNode* n)
    169 {
    170    return (void*)((Addr)elem_of_node(n) + t->keyOff);
    171 }
    172 
    173 static inline
    174 void* fast_key_of_node(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 = *(UWord*)k;
    183    UWord w2 = *(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, HChar* _cc,
    290                              OSetFree_t _free)
    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);
    299    vg_assert(_free);
    300    if (!_cmp) vg_assert(0 == _keyOff);    // If no cmp, offset must be zero
    301 
    302    t           = _alloc(_cc, sizeof(AvlTree));
    303    t->keyOff   = _keyOff;
    304    t->cmp      = _cmp;
    305    t->alloc    = _alloc;
    306    t->cc       = _cc;
    307    t->free     = _free;
    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, HChar* _cc,
    319                                        OSetFree_t _free,
    320                                        SizeT _poolSize,
    321                                        SizeT _maxEltSize)
    322 {
    323    AvlTree* t;
    324 
    325    t = VG_(OSetGen_Create) (_keyOff, _cmp,
    326                             _alloc, _cc,
    327                             _free);
    328 
    329    vg_assert (_poolSize > 0);
    330    vg_assert (_maxEltSize > 0);
    331    t->maxEltSize = _maxEltSize;
    332    t->node_pa = VG_(newPA)(sizeof(AvlNode)
    333                            + VG_ROUNDUP(_maxEltSize, sizeof(void*)),
    334                            _poolSize,
    335                            t->alloc,
    336                            _cc,
    337                            t->free);
    338    VG_(addRefPA) (t->node_pa);
    339 
    340    return t;
    341 }
    342 
    343 AvlTree* VG_(OSetGen_EmptyClone) (AvlTree* os)
    344 {
    345    AvlTree* t;
    346 
    347    vg_assert(os);
    348 
    349    t           = os->alloc(os->cc, sizeof(AvlTree));
    350    t->keyOff   = os->keyOff;
    351    t->cmp      = os->cmp;
    352    t->alloc    = os->alloc;
    353    t->cc       = os->cc;
    354    t->free     = os->free;
    355    t->node_pa  = os->node_pa;
    356    if (t->node_pa)
    357       VG_(addRefPA) (t->node_pa);
    358    t->maxEltSize = os->maxEltSize;
    359    t->nElems   = 0;
    360    t->root     = NULL;
    361    stackClear(t);
    362 
    363    return t;
    364 }
    365 
    366 AvlTree* VG_(OSetWord_Create)(OSetAlloc_t _alloc, HChar* _cc,
    367                               OSetFree_t _free)
    368 {
    369    return VG_(OSetGen_Create)(/*keyOff*/0, /*cmp*/NULL, _alloc, _cc, _free);
    370 }
    371 
    372 // Destructor, frees up all memory held by remaining nodes.
    373 void VG_(OSetGen_Destroy)(AvlTree* t)
    374 {
    375    Bool has_node_pa;
    376    vg_assert(t);
    377 
    378    has_node_pa = t->node_pa != NULL;
    379 
    380    /*
    381     * If we are the only remaining user of this pool allocator, release all
    382     * the elements by deleting the pool allocator. That's more efficient than
    383     * deleting tree nodes one by one.
    384     */
    385    if (!has_node_pa || VG_(releasePA)(t->node_pa) > 0) {
    386       AvlNode* n = NULL;
    387       Int i = 0;
    388       Word sz = 0;
    389 
    390       stackClear(t);
    391       if (t->root)
    392          stackPush(t, t->root, 1);
    393 
    394       /* Free all the AvlNodes.  This is a post-order traversal, because we */
    395       /* must free all children of a node before the node itself. */
    396       while (stackPop(t, &n, &i)) {
    397          switch (i) {
    398          case 1:
    399             stackPush(t, n, 2);
    400             if (n->left)  stackPush(t, n->left, 1);
    401             break;
    402          case 2:
    403             stackPush(t, n, 3);
    404             if (n->right) stackPush(t, n->right, 1);
    405             break;
    406          case 3:
    407             if (has_node_pa)
    408                VG_(freeEltPA) (t->node_pa, n);
    409             else
    410                t->free(n);
    411             sz++;
    412             break;
    413          }
    414       }
    415       vg_assert(sz == t->nElems);
    416    }
    417 
    418    /* Free the AvlTree itself. */
    419    t->free(t);
    420 }
    421 
    422 void VG_(OSetWord_Destroy)(AvlTree* t)
    423 {
    424    VG_(OSetGen_Destroy)(t);
    425 }
    426 
    427 // Allocate and initialise a new node.
    428 void* VG_(OSetGen_AllocNode)(AvlTree* t, SizeT elemSize)
    429 {
    430    AvlNode* n;
    431    Int nodeSize = sizeof(AvlNode) + elemSize;
    432    vg_assert(elemSize > 0);
    433    if (t->node_pa) {
    434       vg_assert(elemSize <= t->maxEltSize);
    435       n = VG_(allocEltPA) (t->node_pa);
    436    } else {
    437       n = t->alloc( t->cc, nodeSize );
    438    }
    439    VG_(memset)(n, 0, nodeSize);
    440    n->magic = OSET_MAGIC;
    441    return elem_of_node(n);
    442 }
    443 
    444 void VG_(OSetGen_FreeNode)(AvlTree* t, void* e)
    445 {
    446    if (t->node_pa)
    447       VG_(freeEltPA) (t->node_pa, node_of_elem (e));
    448    else
    449       t->free( node_of_elem(e) );
    450 }
    451 
    452 /*--------------------------------------------------------------------*/
    453 /*--- Insertion                                                    ---*/
    454 /*--------------------------------------------------------------------*/
    455 
    456 static inline Word cmp_key_root(AvlTree* t, AvlNode* n)
    457 {
    458    return t->cmp
    459           ? slow_cmp(t, slow_key_of_node(t, n), t->root)
    460           : fast_cmp(   fast_key_of_node(   n), t->root);
    461 }
    462 
    463 // Insert element e into the non-empty AVL tree t.
    464 // Returns True if the depth of the tree has grown.
    465 static Bool avl_insert(AvlTree* t, AvlNode* n)
    466 {
    467    Word cmpres = cmp_key_root(t, n);
    468 
    469    if (cmpres < 0) {
    470       // Insert into the left subtree.
    471       if (t->root->left) {
    472          // Only need to set the used fields in the subtree.
    473          AvlTree left_subtree;
    474          left_subtree.root   = t->root->left;
    475          left_subtree.cmp    = t->cmp;
    476          left_subtree.keyOff = t->keyOff;
    477          if (avl_insert(&left_subtree, n)) {
    478              switch (t->root->balance--) {
    479              case 1: return False;
    480              case 0: return True;
    481              }
    482              if (t->root->left->balance < 0) {
    483                 avl_swr(&(t->root));
    484                 t->root->balance = 0;
    485                 t->root->right->balance = 0;
    486              } else {
    487                 avl_swl(&(t->root->left));
    488                 avl_swr(&(t->root));
    489                 avl_nasty(t->root);
    490              }
    491          } else {
    492             t->root->left=left_subtree.root;
    493          }
    494          return False;
    495       } else {
    496          t->root->left = n;
    497          if (t->root->balance--) return False;
    498          return True;
    499       }
    500 
    501    } else if (cmpres > 0) {
    502       // Insert into the right subtree
    503       if (t->root->right) {
    504          // Only need to set the used fields in the subtree.
    505          AvlTree right_subtree;
    506          right_subtree.root   = t->root->right;
    507          right_subtree.cmp    = t->cmp;
    508          right_subtree.keyOff = t->keyOff;
    509          if (avl_insert(&right_subtree, n)) {
    510             switch (t->root->balance++) {
    511             case -1: return False;
    512             case  0: return True;
    513             }
    514             if (t->root->right->balance > 0) {
    515                avl_swl(&(t->root));
    516                t->root->balance = 0;
    517                t->root->left->balance = 0;
    518             } else {
    519                avl_swr(&(t->root->right));
    520                avl_swl(&(t->root));
    521                avl_nasty(t->root);
    522             }
    523          } else {
    524             t->root->right=right_subtree.root;
    525          }
    526          return False;
    527       } else {
    528          t->root->right = n;
    529          if (t->root->balance++) return False;
    530          return True;
    531       }
    532 
    533    } else {
    534       vg_assert2(0, "OSet{Word,Gen}_Insert: duplicate element added");
    535    }
    536 }
    537 
    538 // Insert element e into the AVL tree t.  This is just a wrapper for
    539 // avl_insert() which doesn't return a Bool.
    540 void VG_(OSetGen_Insert)(AvlTree* t, void* e)
    541 {
    542    AvlNode* n;
    543 
    544    vg_assert(t);
    545 
    546    // Initialise.  Even though OSetGen_AllocNode zeroes these fields,
    547    // we should do it again in case a node is removed and then
    548    // re-added to the tree.
    549    n          = node_of_elem(e);
    550    n->left    = 0;
    551    n->right   = 0;
    552    n->balance = 0;
    553 
    554    // Insert into an empty tree
    555    if (!t->root) {
    556       t->root = n;
    557    } else {
    558       avl_insert(t, n);
    559    }
    560 
    561    t->nElems++;
    562    t->stackTop = 0;  // So the iterator can't get out of sync
    563 }
    564 
    565 void VG_(OSetWord_Insert)(AvlTree* t, UWord val)
    566 {
    567    Word* node = VG_(OSetGen_AllocNode)(t, sizeof(UWord));
    568    *node = val;
    569    VG_(OSetGen_Insert)(t, node);
    570 }
    571 
    572 /*--------------------------------------------------------------------*/
    573 /*--- Lookup                                                       ---*/
    574 /*--------------------------------------------------------------------*/
    575 
    576 // Find the *node* in t matching k, or NULL if not found.
    577 static AvlNode* avl_lookup(const AvlTree* t, const void* k)
    578 {
    579    Word     cmpres;
    580    AvlNode* curr = t->root;
    581 
    582    if (t->cmp) {
    583       // General case
    584       while (True) {
    585          if (curr == NULL) return NULL;
    586          cmpres = slow_cmp(t, k, curr);
    587          if      (cmpres < 0) curr = curr->left;
    588          else if (cmpres > 0) curr = curr->right;
    589          else return curr;
    590       }
    591    } else {
    592       // Fast-track special case.  We use the no-check version of
    593       // elem_of_node because it saves about 10% on lookup time.  This
    594       // shouldn't be very dangerous because each node will have been
    595       // checked on insertion.
    596       UWord w1 = *(UWord*)k;
    597       UWord w2;
    598       while (True) {
    599          if (curr == NULL) return NULL;
    600          w2 = *(UWord*)elem_of_node_no_check(curr);
    601          if      (w1 < w2) curr = curr->left;
    602          else if (w1 > w2) curr = curr->right;
    603          else return curr;
    604       }
    605    }
    606 }
    607 
    608 // Find the *element* in t matching k, or NULL if not found.
    609 void* VG_(OSetGen_Lookup)(const AvlTree* t, const void* k)
    610 {
    611    AvlNode* n;
    612    vg_assert(t);
    613    n = avl_lookup(t, k);
    614    return ( n ? elem_of_node(n) : NULL );
    615 }
    616 
    617 // Find the *element* in t matching k, or NULL if not found;  use the given
    618 // comparison function rather than the standard one.
    619 void* VG_(OSetGen_LookupWithCmp)(AvlTree* t, const void* k, OSetCmp_t cmp)
    620 {
    621    // Save the normal one to the side, then restore once we're done.
    622    void* e;
    623    OSetCmp_t tmpcmp;
    624    vg_assert(t);
    625    tmpcmp = t->cmp;
    626    t->cmp = cmp;
    627    e = VG_(OSetGen_Lookup)(t, k);
    628    t->cmp = tmpcmp;
    629    return e;
    630 }
    631 
    632 // Is there an element matching k?
    633 Bool VG_(OSetGen_Contains)(const AvlTree* t, const void* k)
    634 {
    635    return (NULL != VG_(OSetGen_Lookup)(t, k));
    636 }
    637 
    638 Bool VG_(OSetWord_Contains)(AvlTree* t, UWord val)
    639 {
    640    return (NULL != VG_(OSetGen_Lookup)(t, &val));
    641 }
    642 
    643 /*--------------------------------------------------------------------*/
    644 /*--- Deletion                                                     ---*/
    645 /*--------------------------------------------------------------------*/
    646 
    647 static Bool avl_removeroot(AvlTree* t);
    648 
    649 // Remove an already-selected node n from the AVL tree t.
    650 // Returns True if the depth of the tree has shrunk.
    651 static Bool avl_remove(AvlTree* t, AvlNode* n)
    652 {
    653    Bool ch;
    654    Word cmpres = cmp_key_root(t, n);
    655 
    656    if (cmpres < 0) {
    657       AvlTree left_subtree;
    658       // Remove from the left subtree
    659       vg_assert(t->root->left);
    660       // Only need to set the used fields in the subtree.
    661       left_subtree.root   = t->root->left;
    662       left_subtree.cmp    = t->cmp;
    663       left_subtree.keyOff = t->keyOff;
    664       ch = avl_remove(&left_subtree, n);
    665       t->root->left = left_subtree.root;
    666       if (ch) {
    667          switch (t->root->balance++) {
    668          case -1: return True;
    669          case  0: return False;
    670          }
    671          switch (t->root->right->balance) {
    672          case 0:
    673             avl_swl(&(t->root));
    674             t->root->balance = -1;
    675             t->root->left->balance = 1;
    676             return False;
    677          case 1:
    678             avl_swl(&(t->root));
    679             t->root->balance = 0;
    680             t->root->left->balance = 0;
    681             return True;
    682          }
    683          avl_swr(&(t->root->right));
    684          avl_swl(&(t->root));
    685          avl_nasty(t->root);
    686          return True;
    687       } else {
    688          return False;
    689       }
    690 
    691    } else if (cmpres > 0) {
    692       // Remove from the right subtree
    693       AvlTree right_subtree;
    694       vg_assert(t->root->right);
    695       // Only need to set the used fields in the subtree.
    696       right_subtree.root   = t->root->right;
    697       right_subtree.cmp    = t->cmp;
    698       right_subtree.keyOff = t->keyOff;
    699       ch = avl_remove(&right_subtree, n);
    700       t->root->right = right_subtree.root;
    701       if (ch) {
    702          switch (t->root->balance--) {
    703          case 1: return True;
    704          case 0: return False;
    705          }
    706          switch (t->root->left->balance) {
    707          case 0:
    708             avl_swr(&(t->root));
    709             t->root->balance = 1;
    710             t->root->right->balance = -1;
    711             return False;
    712          case -1:
    713             avl_swr(&(t->root));
    714             t->root->balance = 0;
    715             t->root->right->balance = 0;
    716             return True;
    717          }
    718          avl_swl(&(t->root->left));
    719          avl_swr(&(t->root));
    720          avl_nasty(t->root);
    721          return True;
    722       } else {
    723          return False;
    724       }
    725 
    726    } else {
    727       // Found the node to be removed.
    728       vg_assert(t->root == n);
    729       return avl_removeroot(t);
    730    }
    731 }
    732 
    733 // Remove the root of the AVL tree t.
    734 // Returns True if the depth of the tree has shrunk.
    735 static Bool avl_removeroot(AvlTree* t)
    736 {
    737    Bool     ch;
    738    AvlNode* n;
    739 
    740    if (!t->root->left) {
    741       if (!t->root->right) {
    742          t->root = NULL;
    743          return True;
    744       }
    745       t->root = t->root->right;
    746       return True;
    747    }
    748    if (!t->root->right) {
    749       t->root = t->root->left;
    750       return True;
    751    }
    752    if (t->root->balance < 0) {
    753       // Remove from the left subtree
    754       n = t->root->left;
    755       while (n->right) n = n->right;
    756    } else {
    757       // Remove from the right subtree
    758       n = t->root->right;
    759       while (n->left) n = n->left;
    760    }
    761    ch = avl_remove(t, n);
    762    n->left    = t->root->left;
    763    n->right   = t->root->right;
    764    n->balance = t->root->balance;
    765    t->root    = n;
    766    if (n->balance == 0) return ch;
    767    return False;
    768 }
    769 
    770 // Remove and return the element matching the key 'k', or NULL
    771 // if not present.
    772 void* VG_(OSetGen_Remove)(AvlTree* t, const void* k)
    773 {
    774    // Have to find the node first, then remove it.
    775    AvlNode* n = avl_lookup(t, k);
    776    if (n) {
    777       avl_remove(t, n);
    778       t->nElems--;
    779       t->stackTop = 0;     // So the iterator can't get out of sync
    780       return elem_of_node(n);
    781    } else {
    782       return NULL;
    783    }
    784 }
    785 
    786 Bool VG_(OSetWord_Remove)(AvlTree* t, UWord val)
    787 {
    788    void* n = VG_(OSetGen_Remove)(t, &val);
    789    if (n) {
    790       VG_(OSetGen_FreeNode)(t, n);
    791       return True;
    792    } else {
    793       return False;
    794    }
    795 }
    796 
    797 /*--------------------------------------------------------------------*/
    798 /*--- Iterator                                                     ---*/
    799 /*--------------------------------------------------------------------*/
    800 
    801 // The iterator is implemented using in-order traversal with an explicit
    802 // stack, which lets us do the traversal one step at a time and remember
    803 // where we are between each call to OSetGen_Next().
    804 
    805 void VG_(OSetGen_ResetIter)(AvlTree* t)
    806 {
    807    vg_assert(t);
    808    stackClear(t);
    809    if (t->root)
    810       stackPush(t, t->root, 1);
    811 }
    812 
    813 void VG_(OSetWord_ResetIter)(AvlTree* t)
    814 {
    815    VG_(OSetGen_ResetIter)(t);
    816 }
    817 
    818 void* VG_(OSetGen_Next)(AvlTree* t)
    819 {
    820    Int i = 0;
    821    OSetNode* n = NULL;
    822 
    823    vg_assert(t);
    824 
    825    // This in-order traversal requires each node to be pushed and popped
    826    // three times.  These could be avoided by updating nodes in-situ on the
    827    // top of the stack, but the push/pop cost is so small that it's worth
    828    // keeping this loop in this simpler form.
    829    while (stackPop(t, &n, &i)) {
    830       switch (i) {
    831       case 1: case_1:
    832          stackPush(t, n, 2);
    833          /* if (n->left)  stackPush(t, n->left, 1); */
    834          if (n->left) { n = n->left; goto case_1; }
    835          break;
    836       case 2:
    837          stackPush(t, n, 3);
    838          return elem_of_node(n);
    839       case 3:
    840          /* if (n->right) stackPush(t, n->right, 1); */
    841          if (n->right) { n = n->right; goto case_1; }
    842          break;
    843       }
    844    }
    845 
    846    // Stack empty, iterator is exhausted, return NULL
    847    return NULL;
    848 }
    849 
    850 Bool VG_(OSetWord_Next)(AvlTree* t, UWord* val)
    851 {
    852    UWord* n = VG_(OSetGen_Next)(t);
    853    if (n) {
    854       *val = *n;
    855       return True;
    856    } else {
    857       return False;
    858    }
    859 }
    860 
    861 // set up 'oset' for iteration so that the first key subsequently
    862 // produced VG_(OSetGen_Next) is the smallest key in the map
    863 // >= start_at.  Naturally ">=" is defined by the comparison
    864 // function supplied to VG_(OSetGen_Create).
    865 void VG_(OSetGen_ResetIterAt)(AvlTree* oset, const void* k)
    866 {
    867    Int     i;
    868    AvlNode *n, *t;
    869    Word    cmpresS; /* signed */
    870    UWord   cmpresU; /* unsigned */
    871 
    872    vg_assert(oset);
    873    stackClear(oset);
    874 
    875    if (!oset->root)
    876       return;
    877 
    878    n = NULL;
    879    // We need to do regular search and fill in the stack.
    880    t = oset->root;
    881 
    882    while (True) {
    883       if (t == NULL) return;
    884 
    885       if (oset->cmp) {
    886          cmpresS = (Word)slow_cmp(oset, k, t);
    887       } else {
    888          cmpresS = fast_cmp(k, t);
    889       }
    890 
    891       /* Switch the sense of the comparison, since the comparison
    892          order of args (k vs t) above is opposite to that of the
    893          corresponding code in hg_wordfm.c. */
    894       if (cmpresS < 0) { cmpresS = 1; }
    895       else if (cmpresS > 0) { cmpresS = -1; }
    896 
    897       if (cmpresS == 0) {
    898          // We found the exact key -- we are done.
    899          // The iteration should start with this node.
    900          stackPush(oset, t, 2);
    901          // The stack now looks like {2, 2, ... ,2, 2}
    902          return;
    903       }
    904       cmpresU = (UWord)cmpresS;
    905       cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
    906       vg_assert(cmpresU == 0 || cmpresU == 1);
    907       if (!cmpresU) {
    908          // Push this node only if we go to the left child.
    909          stackPush(oset, t, 2);
    910       }
    911       t = cmpresU==0 ? t->left : t->right;
    912    }
    913    if (stackPop(oset, &n, &i)) {
    914       // If we've pushed something to stack and did not find the exact key,
    915       // we must fix the top element of stack.
    916       vg_assert(i == 2);
    917       stackPush(oset, n, 3);
    918       // the stack looks like {2, 2, ..., 2, 3}
    919    }
    920 }
    921 
    922 /*--------------------------------------------------------------------*/
    923 /*--- Miscellaneous operations                                     ---*/
    924 /*--------------------------------------------------------------------*/
    925 
    926 Word VG_(OSetGen_Size)(const AvlTree* t)
    927 {
    928    vg_assert(t);
    929    return t->nElems;
    930 }
    931 
    932 Word VG_(OSetWord_Size)(AvlTree* t)
    933 {
    934    return VG_(OSetGen_Size)(t);
    935 }
    936 
    937 static void OSet_Print2( AvlTree* t, AvlNode* n,
    938                          Char*(*strElem)(void *), Int p )
    939 {
    940    // This is a recursive in-order traversal.
    941    Int q = p;
    942    if (NULL == n) return;
    943    if (n->right) OSet_Print2(t, n->right, strElem, p+1);
    944    while (q--) VG_(printf)(".. ");
    945    VG_(printf)("%s\n", strElem(elem_of_node(n)));
    946    if (n->left) OSet_Print2(t, n->left, strElem, p+1);
    947 }
    948 
    949 __attribute__((unused))
    950 static void OSet_Print( AvlTree* t, const HChar *where, Char*(*strElem)(void *) )
    951 {
    952    VG_(printf)("-- start %s ----------------\n", where);
    953    OSet_Print2(t, t->root, strElem, 0);
    954    VG_(printf)("-- end   %s ----------------\n", where);
    955 }
    956 
    957 /*--------------------------------------------------------------------*/
    958 /*--- end                                                          ---*/
    959 /*--------------------------------------------------------------------*/
    960