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