• 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-2010 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 
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
node_of_elem(const void * elem)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
elem_of_node(const AvlNode * n)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
elem_of_node_no_check(const AvlNode * n)159 void* elem_of_node_no_check(const AvlNode *n)
160 {
161    return (void*)((Addr)n + sizeof(AvlNode));
162 }
163 
164 static inline
slow_key_of_node(AvlTree * t,AvlNode * n)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
fast_key_of_node(AvlNode * n)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*.
fast_cmp(const void * k,const AvlNode * n)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
slow_cmp(const AvlTree * t,const void * k,const AvlNode * n)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.
avl_swl(AvlNode ** root)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.
avl_swr(AvlNode ** root)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.
avl_nasty(AvlNode * root)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.
stackClear(AvlTree * t)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.
stackPush(AvlTree * t,AvlNode * n,Int i)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.
stackPop(AvlTree * t,AvlNode ** n,Int * i)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.
VG_(OSetGen_Create)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 
VG_(OSetWord_Create)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.
VG_(OSetGen_Destroy)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 
VG_(OSetWord_Destroy)354 void VG_(OSetWord_Destroy)(AvlTree* t)
355 {
356    VG_(OSetGen_Destroy)(t);
357 }
358 
359 // Allocate and initialise a new node.
VG_(OSetGen_AllocNode)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 
VG_(OSetGen_FreeNode)370 void VG_(OSetGen_FreeNode)(AvlTree* t, void* e)
371 {
372    t->free( node_of_elem(e) );
373 }
374 
375 /*--------------------------------------------------------------------*/
376 /*--- Insertion                                                    ---*/
377 /*--------------------------------------------------------------------*/
378 
cmp_key_root(AvlTree * t,AvlNode * n)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.
avl_insert(AvlTree * t,AvlNode * n)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.
VG_(OSetGen_Insert)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 
VG_(OSetWord_Insert)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.
avl_lookup(const AvlTree * t,const void * k)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.
VG_(OSetGen_Lookup)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.
VG_(OSetGen_LookupWithCmp)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?
VG_(OSetGen_Contains)556 Bool VG_(OSetGen_Contains)(const AvlTree* t, const void* k)
557 {
558    return (NULL != VG_(OSetGen_Lookup)(t, k));
559 }
560 
VG_(OSetWord_Contains)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.
avl_remove(AvlTree * t,AvlNode * n)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.
avl_removeroot(AvlTree * t)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.
VG_(OSetGen_Remove)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 
VG_(OSetWord_Remove)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 
VG_(OSetGen_ResetIter)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 
VG_(OSetWord_ResetIter)736 void VG_(OSetWord_ResetIter)(AvlTree* t)
737 {
738    VG_(OSetGen_ResetIter)(t);
739 }
740 
VG_(OSetGen_Next)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 
VG_(OSetWord_Next)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).
VG_(OSetGen_ResetIterAt)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 
VG_(OSetGen_Size)849 Word VG_(OSetGen_Size)(const AvlTree* t)
850 {
851    vg_assert(t);
852    return t->nElems;
853 }
854 
VG_(OSetWord_Size)855 Word VG_(OSetWord_Size)(AvlTree* t)
856 {
857    return VG_(OSetGen_Size)(t);
858 }
859 
OSet_Print2(AvlTree * t,AvlNode * n,Char * (* strElem)(void *),Int p)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))
OSet_Print(AvlTree * t,const HChar * where,Char * (* strElem)(void *))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