1
2/*--------------------------------------------------------------------*/
3/*--- An ordered set implemented using an AVL tree.       m_oset.c ---*/
4/*--------------------------------------------------------------------*/
5
6/*
7   This file is part of Valgrind, a dynamic binary instrumentation
8   framework.
9
10   Copyright (C) 2005-2012 Nicholas Nethercote
11      njn@valgrind.org
12
13   This program is free software; you can redistribute it and/or
14   modify it under the terms of the GNU General Public License as
15   published by the Free Software Foundation; either version 2 of the
16   License, or (at your option) any later version.
17
18   This program is distributed in the hope that it will be useful, but
19   WITHOUT ANY WARRANTY; without even the implied warranty of
20   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
21   General Public License for more details.
22
23   You should have received a copy of the GNU General Public License
24   along with this program; if not, write to the Free Software
25   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
26   02111-1307, USA.
27
28   The GNU General Public License is contained in the file COPYING.
29*/
30
31//----------------------------------------------------------------------
32// This file is based on:
33//
34//   ANSI C Library for maintainance of AVL Balanced Trees
35//   (C) 2000 Daniel Nagy, Budapest University of Technology and Economics
36//   Released under GNU General Public License (GPL) version 2
37//----------------------------------------------------------------------
38
39// This file implements a generic ordered set using an AVL tree.
40//
41// Each node in the tree has two parts.
42// - First is the AVL metadata, which is three words: a left pointer, a
43//   right pointer, and a word containing balancing information and a
44//   "magic" value which provides some checking that the user has not
45//   corrupted the metadata.  So the overhead is 12 bytes on 32-bit
46//   platforms and 24 bytes on 64-bit platforms.
47// - Second is the user's data.  This can be anything.  Note that because it
48//   comes after the metadata, it will only be word-aligned, even if the
49//   user data is a struct that would normally be doubleword-aligned.
50//
51// AvlNode* node -> +---------------+  V
52//                  | struct        |
53//                  |   AvlNode     |
54// void* element -> +---------------+  ^
55//                  | element       |  |
56//      keyOff ->   | key           | elemSize
57//                  +---------------+  v
58//
59// Users have to allocate AvlNodes with OSetGen_AllocNode(), which allocates
60// space for the metadata.
61//
62// The terminology used throughout this file:
63// - a "node", usually called "n", is a pointer to the metadata.
64// - an "element", usually called "e", is a pointer to the user data.
65// - a "key", usually called "k", is a pointer to a key.
66//
67// The helper functions elem_of_node and node_of_elem do the pointer
68// arithmetic to switch between the node and the element.  The node magic is
69// checked after each operation to make sure that we're really operating on
70// an AvlNode.
71//
72// Each tree also has an iterator.  Note that we cannot use the iterator
73// internally within this file (eg. we could implement OSetGen_Size() by
74// stepping through with the iterator and counting nodes) because it's
75// non-reentrant -- the user might be using it themselves, and the
76// concurrent uses would screw things up.
77
78#include "pub_core_basics.h"
79#include "pub_core_libcbase.h"
80#include "pub_core_libcassert.h"
81#include "pub_core_libcprint.h"
82#include "pub_core_oset.h"
83#include "pub_tool_poolalloc.h"
84
85/*--------------------------------------------------------------------*/
86/*--- Types and constants                                          ---*/
87/*--------------------------------------------------------------------*/
88
89typedef struct _OSetNode OSetNode;
90
91// Internal names for the OSet types.
92typedef OSet     AvlTree;
93typedef 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.
98struct _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.
112struct _OSet {
113   SizeT       keyOff;     // key offset
114   OSetCmp_t   cmp;        // compare a key and an element, or NULL
115   OSetAlloc_t alloc;      // allocator
116   HChar* cc;              // cc for allocator
117   OSetFree_t  free;       // deallocator
118   PoolAlloc*  node_pa;    // (optional) pool allocator for nodes.
119   SizeT       maxEltSize; // for node_pa, must be > 0. Otherwise unused.
120   Word        nElems;     // number of elements in the tree
121   AvlNode*    root;       // root node
122
123   AvlNode*    nodeStack[STACK_MAX];   // Iterator node stack
124   Int          numStack[STACK_MAX];   // Iterator num stack
125   Int         stackTop;               // Iterator stack pointer, one past end
126};
127
128/*--------------------------------------------------------------------*/
129/*--- Helper operations                                            ---*/
130/*--------------------------------------------------------------------*/
131
132// Given a pointer to the node's element, return the pointer to the AvlNode
133// structure.  If the node has a bad magic number, it will die with an
134// assertion failure.
135static inline
136AvlNode* 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.
149static inline
150void* 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.
161static inline
162void* elem_of_node_no_check(const AvlNode *n)
163{
164   return (void*)((Addr)n + sizeof(AvlNode));
165}
166
167static inline
168void* slow_key_of_node(AvlTree* t, AvlNode* n)
169{
170   return (void*)((Addr)elem_of_node(n) + t->keyOff);
171}
172
173static inline
174void* 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*.
180static inline Word fast_cmp(const void* k, const AvlNode* n)
181{
182   UWord w1 = *(UWord*)k;
183   UWord w2 = *(UWord*)elem_of_node(n);
184   // In previous versions, we tried to do this faster by doing
185   // "return w1 - w2".  But it didn't work reliably, because the
186   // complete result of subtracting two N-bit numbers is an N+1-bit
187   // number, and what the caller is interested in is the sign of
188   // the complete N+1-bit result.  The branching version is slightly
189   // slower, but safer and easier to understand.
190   if (w1 > w2) return  1;
191   if (w1 < w2) return -1;
192   return 0;
193}
194
195// Compare a key and an element.  Inlining is *crucial*.
196static
197inline 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.
204static 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.
214static 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.
224static 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.
244static 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.
256static 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.
266static 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.
288AvlTree* VG_(OSetGen_Create)(PtrdiffT _keyOff, OSetCmp_t _cmp,
289                             OSetAlloc_t _alloc, HChar* _cc,
290                             OSetFree_t _free)
291{
292   AvlTree* t;
293
294   // Check the padding is right and the AvlNode is the expected size.
295   vg_assert(sizeof(AvlNode) == 3*sizeof(void*));
296
297   // Sanity check args
298   vg_assert(_alloc);
299   vg_assert(_free);
300   if (!_cmp) vg_assert(0 == _keyOff);    // If no cmp, offset must be zero
301
302   t           = _alloc(_cc, sizeof(AvlTree));
303   t->keyOff   = _keyOff;
304   t->cmp      = _cmp;
305   t->alloc    = _alloc;
306   t->cc       = _cc;
307   t->free     = _free;
308   t->node_pa  = NULL;
309   t->maxEltSize = 0; // Just in case it would be wrongly used.
310   t->nElems   = 0;
311   t->root     = NULL;
312   stackClear(t);
313
314   return t;
315}
316
317AvlTree* VG_(OSetGen_Create_With_Pool)(PtrdiffT _keyOff, OSetCmp_t _cmp,
318                                       OSetAlloc_t _alloc, HChar* _cc,
319                                       OSetFree_t _free,
320                                       SizeT _poolSize,
321                                       SizeT _maxEltSize)
322{
323   AvlTree* t;
324
325   t = VG_(OSetGen_Create) (_keyOff, _cmp,
326                            _alloc, _cc,
327                            _free);
328
329   vg_assert (_poolSize > 0);
330   vg_assert (_maxEltSize > 0);
331   t->maxEltSize = _maxEltSize;
332   t->node_pa = VG_(newPA)(sizeof(AvlNode)
333                           + VG_ROUNDUP(_maxEltSize, sizeof(void*)),
334                           _poolSize,
335                           t->alloc,
336                           _cc,
337                           t->free);
338   VG_(addRefPA) (t->node_pa);
339
340   return t;
341}
342
343AvlTree* 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
366AvlTree* VG_(OSetWord_Create)(OSetAlloc_t _alloc, HChar* _cc,
367                              OSetFree_t _free)
368{
369   return VG_(OSetGen_Create)(/*keyOff*/0, /*cmp*/NULL, _alloc, _cc, _free);
370}
371
372// Destructor, frees up all memory held by remaining nodes.
373void 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
422void VG_(OSetWord_Destroy)(AvlTree* t)
423{
424   VG_(OSetGen_Destroy)(t);
425}
426
427// Allocate and initialise a new node.
428void* 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
444void 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
456static 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.
465static 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.
540void 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
565void 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.
577static AvlNode* avl_lookup(const AvlTree* t, const void* k)
578{
579   Word     cmpres;
580   AvlNode* curr = t->root;
581
582   if (t->cmp) {
583      // General case
584      while (True) {
585         if (curr == NULL) return NULL;
586         cmpres = slow_cmp(t, k, curr);
587         if      (cmpres < 0) curr = curr->left;
588         else if (cmpres > 0) curr = curr->right;
589         else return curr;
590      }
591   } else {
592      // Fast-track special case.  We use the no-check version of
593      // elem_of_node because it saves about 10% on lookup time.  This
594      // shouldn't be very dangerous because each node will have been
595      // checked on insertion.
596      UWord w1 = *(UWord*)k;
597      UWord w2;
598      while (True) {
599         if (curr == NULL) return NULL;
600         w2 = *(UWord*)elem_of_node_no_check(curr);
601         if      (w1 < w2) curr = curr->left;
602         else if (w1 > w2) curr = curr->right;
603         else return curr;
604      }
605   }
606}
607
608// Find the *element* in t matching k, or NULL if not found.
609void* 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.
619void* 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?
633Bool VG_(OSetGen_Contains)(const AvlTree* t, const void* k)
634{
635   return (NULL != VG_(OSetGen_Lookup)(t, k));
636}
637
638Bool VG_(OSetWord_Contains)(AvlTree* t, UWord val)
639{
640   return (NULL != VG_(OSetGen_Lookup)(t, &val));
641}
642
643/*--------------------------------------------------------------------*/
644/*--- Deletion                                                     ---*/
645/*--------------------------------------------------------------------*/
646
647static 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.
651static 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.
735static 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.
772void* 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
786Bool 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
805void 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
813void VG_(OSetWord_ResetIter)(AvlTree* t)
814{
815   VG_(OSetGen_ResetIter)(t);
816}
817
818void* 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
850Bool 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).
865void VG_(OSetGen_ResetIterAt)(AvlTree* oset, const void* k)
866{
867   Int     i;
868   AvlNode *n, *t;
869   Word    cmpresS; /* signed */
870   UWord   cmpresU; /* unsigned */
871
872   vg_assert(oset);
873   stackClear(oset);
874
875   if (!oset->root)
876      return;
877
878   n = NULL;
879   // We need to do regular search and fill in the stack.
880   t = oset->root;
881
882   while (True) {
883      if (t == NULL) return;
884
885      if (oset->cmp) {
886         cmpresS = (Word)slow_cmp(oset, k, t);
887      } else {
888         cmpresS = fast_cmp(k, t);
889      }
890
891      /* Switch the sense of the comparison, since the comparison
892         order of args (k vs t) above is opposite to that of the
893         corresponding code in hg_wordfm.c. */
894      if (cmpresS < 0) { cmpresS = 1; }
895      else if (cmpresS > 0) { cmpresS = -1; }
896
897      if (cmpresS == 0) {
898         // We found the exact key -- we are done.
899         // The iteration should start with this node.
900         stackPush(oset, t, 2);
901         // The stack now looks like {2, 2, ... ,2, 2}
902         return;
903      }
904      cmpresU = (UWord)cmpresS;
905      cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
906      vg_assert(cmpresU == 0 || cmpresU == 1);
907      if (!cmpresU) {
908         // Push this node only if we go to the left child.
909         stackPush(oset, t, 2);
910      }
911      t = cmpresU==0 ? t->left : t->right;
912   }
913   if (stackPop(oset, &n, &i)) {
914      // If we've pushed something to stack and did not find the exact key,
915      // we must fix the top element of stack.
916      vg_assert(i == 2);
917      stackPush(oset, n, 3);
918      // the stack looks like {2, 2, ..., 2, 3}
919   }
920}
921
922/*--------------------------------------------------------------------*/
923/*--- Miscellaneous operations                                     ---*/
924/*--------------------------------------------------------------------*/
925
926Word VG_(OSetGen_Size)(const AvlTree* t)
927{
928   vg_assert(t);
929   return t->nElems;
930}
931
932Word VG_(OSetWord_Size)(AvlTree* t)
933{
934   return VG_(OSetGen_Size)(t);
935}
936
937static void OSet_Print2( AvlTree* t, AvlNode* n,
938                         Char*(*strElem)(void *), Int p )
939{
940   // This is a recursive in-order traversal.
941   Int q = p;
942   if (NULL == n) return;
943   if (n->right) OSet_Print2(t, n->right, strElem, p+1);
944   while (q--) VG_(printf)(".. ");
945   VG_(printf)("%s\n", strElem(elem_of_node(n)));
946   if (n->left) OSet_Print2(t, n->left, strElem, p+1);
947}
948
949__attribute__((unused))
950static void OSet_Print( AvlTree* t, const HChar *where, Char*(*strElem)(void *) )
951{
952   VG_(printf)("-- start %s ----------------\n", where);
953   OSet_Print2(t, t->root, strElem, 0);
954   VG_(printf)("-- end   %s ----------------\n", where);
955}
956
957/*--------------------------------------------------------------------*/
958/*--- end                                                          ---*/
959/*--------------------------------------------------------------------*/
960