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