• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 
2 /*--------------------------------------------------------------------*/
3 /*--- An ordered set implemented using an AVL tree.       m_oset.c ---*/
4 /*--------------------------------------------------------------------*/
5 
6 /*
7    This file is part of Valgrind, a dynamic binary instrumentation
8    framework.
9 
10    Copyright (C) 2005-2013 Nicholas Nethercote
11       njn@valgrind.org
12 
13    This program is free software; you can redistribute it and/or
14    modify it under the terms of the GNU General Public License as
15    published by the Free Software Foundation; either version 2 of the
16    License, or (at your option) any later version.
17 
18    This program is distributed in the hope that it will be useful, but
19    WITHOUT ANY WARRANTY; without even the implied warranty of
20    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
21    General Public License for more details.
22 
23    You should have received a copy of the GNU General Public License
24    along with this program; if not, write to the Free Software
25    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
26    02111-1307, USA.
27 
28    The GNU General Public License is contained in the file COPYING.
29 */
30 
31 //----------------------------------------------------------------------
32 // This file is based on:
33 //
34 //   ANSI C Library for maintainance of AVL Balanced Trees
35 //   (C) 2000 Daniel Nagy, Budapest University of Technology and Economics
36 //   Released under GNU General Public License (GPL) version 2
37 //----------------------------------------------------------------------
38 
39 // This file implements a generic ordered set using an AVL tree.
40 //
41 // Each node in the tree has two parts.
42 // - First is the AVL metadata, which is three words: a left pointer, a
43 //   right pointer, and a word containing balancing information and a
44 //   "magic" value which provides some checking that the user has not
45 //   corrupted the metadata.  So the overhead is 12 bytes on 32-bit
46 //   platforms and 24 bytes on 64-bit platforms.
47 // - Second is the user's data.  This can be anything.  Note that because it
48 //   comes after the metadata, it will only be word-aligned, even if the
49 //   user data is a struct that would normally be doubleword-aligned.
50 //
51 // AvlNode* node -> +---------------+  V
52 //                  | struct        |
53 //                  |   AvlNode     |
54 // void* element -> +---------------+  ^
55 //                  | element       |  |
56 //      keyOff ->   | key           | elemSize
57 //                  +---------------+  v
58 //
59 // Users have to allocate AvlNodes with OSetGen_AllocNode(), which allocates
60 // space for the metadata.
61 //
62 // The terminology used throughout this file:
63 // - a "node", usually called "n", is a pointer to the metadata.
64 // - an "element", usually called "e", is a pointer to the user data.
65 // - a "key", usually called "k", is a pointer to a key.
66 //
67 // The helper functions elem_of_node and node_of_elem do the pointer
68 // arithmetic to switch between the node and the element.  The node magic is
69 // checked after each operation to make sure that we're really operating on
70 // an AvlNode.
71 //
72 // Each tree also has an iterator.  Note that we cannot use the iterator
73 // internally within this file (eg. we could implement OSetGen_Size() by
74 // stepping through with the iterator and counting nodes) because it's
75 // non-reentrant -- the user might be using it themselves, and the
76 // concurrent uses would screw things up.
77 
78 #include "pub_core_basics.h"
79 #include "pub_core_libcbase.h"
80 #include "pub_core_libcassert.h"
81 #include "pub_core_libcprint.h"
82 #include "pub_core_oset.h"
83 #include "pub_core_poolalloc.h"
84 
85 /*--------------------------------------------------------------------*/
86 /*--- Types and constants                                          ---*/
87 /*--------------------------------------------------------------------*/
88 
89 typedef struct _OSetNode OSetNode;
90 
91 // Internal names for the OSet types.
92 typedef OSet     AvlTree;
93 typedef OSetNode AvlNode;
94 
95 // The padding ensures that magic is right at the end of the node,
96 // regardless of the machine's word size, so that any overwrites will be
97 // detected earlier.
98 struct _OSetNode {
99    AvlNode* left;
100    AvlNode* right;
101    Char     balance;
102    Char     padding[sizeof(void*)-sizeof(Char)-sizeof(Short)];
103    Short    magic;
104 };
105 
106 #define STACK_MAX    32    // At most 2**32 entries can be iterated over
107 #define OSET_MAGIC   0x5b1f
108 
109 // An OSet (AVL tree).  If cmp is NULL, the key must be a UWord, and must
110 // be the first word in the element.  If cmp is set, arbitrary keys in
111 // arbitrary positions can be used.
112 struct _OSet {
113    SizeT       keyOff;     // key offset
114    OSetCmp_t   cmp;        // compare a key and an element, or NULL
115    OSetAlloc_t alloc;      // allocator
116    const 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
node_of_elem(const void * elem)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
elem_of_node(const AvlNode * n)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
elem_of_node_no_check(const AvlNode * n)162 void* elem_of_node_no_check(const AvlNode *n)
163 {
164    return (void*)((Addr)n + sizeof(AvlNode));
165 }
166 
167 static inline
slow_key_of_node(AvlTree * t,AvlNode * n)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
fast_key_of_node(AvlNode * n)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*.
fast_cmp(const void * k,const AvlNode * n)180 static inline Word fast_cmp(const void* k, const AvlNode* n)
181 {
182    UWord w1 = *(const UWord*)k;
183    UWord w2 = *(const UWord*)elem_of_node(n);
184    // In previous versions, we tried to do this faster by doing
185    // "return w1 - w2".  But it didn't work reliably, because the
186    // complete result of subtracting two N-bit numbers is an N+1-bit
187    // number, and what the caller is interested in is the sign of
188    // the complete N+1-bit result.  The branching version is slightly
189    // slower, but safer and easier to understand.
190    if (w1 > w2) return  1;
191    if (w1 < w2) return -1;
192    return 0;
193 }
194 
195 // Compare a key and an element.  Inlining is *crucial*.
196 static
slow_cmp(const AvlTree * t,const void * k,const AvlNode * n)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.
avl_swl(AvlNode ** root)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.
avl_swr(AvlNode ** root)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.
avl_nasty(AvlNode * root)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.
stackClear(AvlTree * t)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.
stackPush(AvlTree * t,AvlNode * n,Int i)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.
stackPop(AvlTree * t,AvlNode ** n,Int * i)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.
VG_(OSetGen_Create)288 AvlTree* VG_(OSetGen_Create)(PtrdiffT _keyOff, OSetCmp_t _cmp,
289                              OSetAlloc_t _alloc, const 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 
VG_(OSetGen_Create_With_Pool)317 AvlTree* VG_(OSetGen_Create_With_Pool)(PtrdiffT _keyOff, OSetCmp_t _cmp,
318                                        OSetAlloc_t _alloc, const 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 
VG_(OSetGen_EmptyClone)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 
VG_(OSetWord_Create)366 AvlTree* VG_(OSetWord_Create)(OSetAlloc_t _alloc, const 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.
VG_(OSetGen_Destroy)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 
VG_(OSetWord_Destroy)422 void VG_(OSetWord_Destroy)(AvlTree* t)
423 {
424    VG_(OSetGen_Destroy)(t);
425 }
426 
427 // Allocate and initialise a new node.
VG_(OSetGen_AllocNode)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 
VG_(OSetGen_FreeNode)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 
cmp_key_root(AvlTree * t,AvlNode * n)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.
avl_insert(AvlTree * t,AvlNode * n)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.
VG_(OSetGen_Insert)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 
VG_(OSetWord_Insert)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.
avl_lookup(const AvlTree * t,const void * k)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 = *(const 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.
VG_(OSetGen_Lookup)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.
VG_(OSetGen_LookupWithCmp)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?
VG_(OSetGen_Contains)633 Bool VG_(OSetGen_Contains)(const AvlTree* t, const void* k)
634 {
635    return (NULL != VG_(OSetGen_Lookup)(t, k));
636 }
637 
VG_(OSetWord_Contains)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.
avl_remove(AvlTree * t,AvlNode * n)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.
avl_removeroot(AvlTree * t)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.
VG_(OSetGen_Remove)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 
VG_(OSetWord_Remove)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 
VG_(OSetGen_ResetIter)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 
VG_(OSetWord_ResetIter)813 void VG_(OSetWord_ResetIter)(AvlTree* t)
814 {
815    VG_(OSetGen_ResetIter)(t);
816 }
817 
VG_(OSetGen_Next)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 
VG_(OSetWord_Next)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).
VG_(OSetGen_ResetIterAt)865 void VG_(OSetGen_ResetIterAt)(AvlTree* oset, const void* k)
866 {
867    AvlNode *t;
868    Word    cmpresS; /* signed */
869    UWord   cmpresU; /* unsigned */
870 
871    vg_assert(oset);
872    stackClear(oset);
873 
874    if (!oset->root)
875       return;
876 
877    // We need to do regular search and fill in the stack.
878    t = oset->root;
879 
880    while (True) {
881       if (t == NULL) return;
882 
883       if (oset->cmp) {
884          cmpresS = (Word)slow_cmp(oset, k, t);
885       } else {
886          cmpresS = fast_cmp(k, t);
887       }
888 
889       /* Switch the sense of the comparison, since the comparison
890          order of args (k vs t) above is opposite to that of the
891          corresponding code in hg_wordfm.c. */
892       if (cmpresS < 0) { cmpresS = 1; }
893       else if (cmpresS > 0) { cmpresS = -1; }
894 
895       if (cmpresS == 0) {
896          // We found the exact key -- we are done.
897          // The iteration should start with this node.
898          stackPush(oset, t, 2);
899          // The stack now looks like {2, 2, ... ,2, 2}
900          return;
901       }
902       cmpresU = (UWord)cmpresS;
903       cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
904       vg_assert(cmpresU == 0 || cmpresU == 1);
905       if (!cmpresU) {
906          // Push this node only if we go to the left child.
907          stackPush(oset, t, 2);
908       }
909       t = cmpresU==0 ? t->left : t->right;
910    }
911 }
912 
913 /*--------------------------------------------------------------------*/
914 /*--- Miscellaneous operations                                     ---*/
915 /*--------------------------------------------------------------------*/
916 
VG_(OSetGen_Size)917 Word VG_(OSetGen_Size)(const AvlTree* t)
918 {
919    vg_assert(t);
920    return t->nElems;
921 }
922 
VG_(OSetWord_Size)923 Word VG_(OSetWord_Size)(AvlTree* t)
924 {
925    return VG_(OSetGen_Size)(t);
926 }
927 
OSet_Print2(AvlTree * t,AvlNode * n,HChar * (* strElem)(void *),Int p)928 static void OSet_Print2( AvlTree* t, AvlNode* n,
929                          HChar*(*strElem)(void *), Int p )
930 {
931    // This is a recursive in-order traversal.
932    Int q = p;
933    if (NULL == n) return;
934    if (n->right) OSet_Print2(t, n->right, strElem, p+1);
935    while (q--) VG_(printf)(".. ");
936    VG_(printf)("%s\n", strElem(elem_of_node(n)));
937    if (n->left) OSet_Print2(t, n->left, strElem, p+1);
938 }
939 
940 __attribute__((unused))
OSet_Print(AvlTree * t,const HChar * where,HChar * (* strElem)(void *))941 static void OSet_Print( AvlTree* t, const HChar *where,
942                         HChar*(*strElem)(void *) )
943 {
944    VG_(printf)("-- start %s ----------------\n", where);
945    OSet_Print2(t, t->root, strElem, 0);
946    VG_(printf)("-- end   %s ----------------\n", where);
947 }
948 
949 /*--------------------------------------------------------------------*/
950 /*--- end                                                          ---*/
951 /*--------------------------------------------------------------------*/
952