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
88typedef struct _OSetNode OSetNode;
89
90// Internal names for the OSet types.
91typedef OSet     AvlTree;
92typedef 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.
97struct _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.
111struct _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.
132static inline
133AvlNode* 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.
146static inline
147void* 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.
158static inline
159void* elem_of_node_no_check(const AvlNode *n)
160{
161   return (void*)((Addr)n + sizeof(AvlNode));
162}
163
164static inline
165void* slow_key_of_node(AvlTree* t, AvlNode* n)
166{
167   return (void*)((Addr)elem_of_node(n) + t->keyOff);
168}
169
170static inline
171void* 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*.
177static 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*.
193static
194inline 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.
201static 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.
211static 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.
221static 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.
241static 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.
253static 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.
263static 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.
285AvlTree* 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
312AvlTree* 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.
319void 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
354void VG_(OSetWord_Destroy)(AvlTree* t)
355{
356   VG_(OSetGen_Destroy)(t);
357}
358
359// Allocate and initialise a new node.
360void* 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
370void VG_(OSetGen_FreeNode)(AvlTree* t, void* e)
371{
372   t->free( node_of_elem(e) );
373}
374
375/*--------------------------------------------------------------------*/
376/*--- Insertion                                                    ---*/
377/*--------------------------------------------------------------------*/
378
379static 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.
388static 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.
463void 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
488void 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.
500static 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.
532void* 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.
542void* 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?
556Bool VG_(OSetGen_Contains)(const AvlTree* t, const void* k)
557{
558   return (NULL != VG_(OSetGen_Lookup)(t, k));
559}
560
561Bool VG_(OSetWord_Contains)(AvlTree* t, UWord val)
562{
563   return (NULL != VG_(OSetGen_Lookup)(t, &val));
564}
565
566/*--------------------------------------------------------------------*/
567/*--- Deletion                                                     ---*/
568/*--------------------------------------------------------------------*/
569
570static 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.
574static 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.
658static 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.
695void* 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
709Bool 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
728void 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
736void VG_(OSetWord_ResetIter)(AvlTree* t)
737{
738   VG_(OSetGen_ResetIter)(t);
739}
740
741void* 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
773Bool 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).
788void 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
849Word VG_(OSetGen_Size)(const AvlTree* t)
850{
851   vg_assert(t);
852   return t->nElems;
853}
854
855Word VG_(OSetWord_Size)(AvlTree* t)
856{
857   return VG_(OSetGen_Size)(t);
858}
859
860static 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))
873static 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