1
2/*--------------------------------------------------------------------*/
3/*--- An AVL tree based finite map for word keys and word values.  ---*/
4/*--- Inspired by Haskell's "FiniteMap" library.                   ---*/
5/*---                                                   m_wordfm.c ---*/
6/*--------------------------------------------------------------------*/
7
8/*
9   This file is part of Valgrind, a dynamic binary instrumentation
10   framework.
11
12   Copyright (C) 2007-2013 Julian Seward
13      jseward@acm.org
14
15   This code is based on previous work by Nicholas Nethercote
16   (coregrind/m_oset.c) which is
17
18   Copyright (C) 2005-2013 Nicholas Nethercote
19       njn@valgrind.org
20
21   which in turn was derived partially from:
22
23      AVL C library
24      Copyright (C) 2000,2002  Daniel Nagy
25
26      This program is free software; you can redistribute it and/or
27      modify it under the terms of the GNU General Public License as
28      published by the Free Software Foundation; either version 2 of
29      the License, or (at your option) any later version.
30      [...]
31
32      (taken from libavl-0.4/debian/copyright)
33
34   This program is free software; you can redistribute it and/or
35   modify it under the terms of the GNU General Public License as
36   published by the Free Software Foundation; either version 2 of the
37   License, or (at your option) any later version.
38
39   This program is distributed in the hope that it will be useful, but
40   WITHOUT ANY WARRANTY; without even the implied warranty of
41   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
42   General Public License for more details.
43
44   You should have received a copy of the GNU General Public License
45   along with this program; if not, write to the Free Software
46   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
47   02111-1307, USA.
48
49   The GNU General Public License is contained in the file COPYING.
50*/
51
52#include "pub_core_basics.h"
53#include "pub_core_libcassert.h"
54#include "pub_core_libcbase.h"
55#include "pub_core_wordfm.h"   /* self */
56
57
58//------------------------------------------------------------------//
59//---                           WordFM                           ---//
60//---                       Implementation                       ---//
61//------------------------------------------------------------------//
62
63/* One element of the AVL tree */
64typedef
65   struct _AvlNode {
66      UWord key;
67      UWord val;
68      struct _AvlNode* child[2]; /* [0] is left subtree, [1] is right */
69      Char balance; /* do not make this unsigned */
70   }
71   AvlNode;
72
73typedef
74   struct {
75      UWord w;
76      Bool b;
77   }
78   MaybeWord;
79
80#define WFM_STKMAX    32    // At most 2**32 entries can be iterated over
81
82struct _WordFM {
83   AvlNode* root;
84   void*    (*alloc_nofail)( const HChar*, SizeT );
85   const HChar* cc;
86   void     (*dealloc)(void*);
87   Word     (*kCmp)(UWord,UWord);
88   AvlNode* nodeStack[WFM_STKMAX]; // Iterator node stack
89   Int      numStack[WFM_STKMAX];  // Iterator num stack
90   Int      stackTop;              // Iterator stack pointer, one past end
91};
92
93/* forward */
94static Bool avl_removeroot_wrk(AvlNode** t, Word(*kCmp)(UWord,UWord));
95
96/* Swing to the left.  Warning: no balance maintainance. */
97static void avl_swl ( AvlNode** root )
98{
99   AvlNode* a  = *root;
100   AvlNode* b  = a->child[1];
101   *root       = b;
102   a->child[1] = b->child[0];
103   b->child[0] = a;
104}
105
106/* Swing to the right.  Warning: no balance maintainance. */
107static void avl_swr ( AvlNode** root )
108{
109   AvlNode* a  = *root;
110   AvlNode* b  = a->child[0];
111   *root       = b;
112   a->child[0] = b->child[1];
113   b->child[1] = a;
114}
115
116/* Balance maintainance after especially nasty swings. */
117static void avl_nasty ( AvlNode* root )
118{
119   switch (root->balance) {
120      case -1:
121         root->child[0]->balance = 0;
122         root->child[1]->balance = 1;
123         break;
124      case 1:
125         root->child[0]->balance = -1;
126         root->child[1]->balance = 0;
127         break;
128      case 0:
129         root->child[0]->balance = 0;
130         root->child[1]->balance = 0;
131         break;
132      default:
133         tl_assert(0);
134   }
135   root->balance=0;
136}
137
138/* Find size of a non-NULL tree. */
139static UWord size_avl_nonNull ( AvlNode* nd )
140{
141   return 1 + (nd->child[0] ? size_avl_nonNull(nd->child[0]) : 0)
142            + (nd->child[1] ? size_avl_nonNull(nd->child[1]) : 0);
143}
144
145/* Unsignedly compare w1 and w2.  If w1 < w2, produce a negative
146   number; if w1 > w2 produce a positive number, and if w1 == w2
147   produce zero. */
148static inline Word cmp_unsigned_Words ( UWord w1, UWord w2 ) {
149   if (w1 < w2) return -1;
150   if (w1 > w2) return 1;
151   return 0;
152}
153
154/* Insert element a into the AVL tree t.  Returns True if the depth of
155   the tree has grown.  If element with that key is already present,
156   just copy a->val to existing node, first returning old ->val field
157   of existing node in *oldV, so that the caller can finalize it
158   however it wants.
159*/
160static
161Bool avl_insert_wrk ( AvlNode**         rootp,
162                      /*OUT*/MaybeWord* oldV,
163                      AvlNode*          a,
164                      Word              (*kCmp)(UWord,UWord) )
165{
166   Word cmpres;
167
168   /* initialize */
169   a->child[0] = 0;
170   a->child[1] = 0;
171   a->balance  = 0;
172   oldV->b     = False;
173
174   /* insert into an empty tree? */
175   if (!(*rootp)) {
176      (*rootp) = a;
177      return True;
178   }
179
180   cmpres = kCmp ? /*boxed*/   kCmp( (*rootp)->key, a->key )
181                 : /*unboxed*/ cmp_unsigned_Words( (UWord)(*rootp)->key,
182                                                   (UWord)a->key );
183
184   if (cmpres > 0) {
185      /* insert into the left subtree */
186      if ((*rootp)->child[0]) {
187         AvlNode* left_subtree = (*rootp)->child[0];
188         if (avl_insert_wrk(&left_subtree, oldV, a, kCmp)) {
189            switch ((*rootp)->balance--) {
190               case  1: return False;
191               case  0: return True;
192               case -1: break;
193               default: tl_assert(0);
194            }
195            if ((*rootp)->child[0]->balance < 0) {
196               avl_swr( rootp );
197               (*rootp)->balance = 0;
198               (*rootp)->child[1]->balance = 0;
199            } else {
200               avl_swl( &((*rootp)->child[0]) );
201               avl_swr( rootp );
202               avl_nasty( *rootp );
203            }
204         } else {
205            (*rootp)->child[0] = left_subtree;
206         }
207         return False;
208      } else {
209         (*rootp)->child[0] = a;
210         if ((*rootp)->balance--)
211            return False;
212         return True;
213      }
214      tl_assert(0);/*NOTREACHED*/
215   }
216   else
217   if (cmpres < 0) {
218      /* insert into the right subtree */
219      if ((*rootp)->child[1]) {
220         AvlNode* right_subtree = (*rootp)->child[1];
221         if (avl_insert_wrk(&right_subtree, oldV, a, kCmp)) {
222            switch((*rootp)->balance++){
223               case -1: return False;
224               case  0: return True;
225               case  1: break;
226               default: tl_assert(0);
227            }
228            if ((*rootp)->child[1]->balance > 0) {
229               avl_swl( rootp );
230               (*rootp)->balance = 0;
231               (*rootp)->child[0]->balance = 0;
232            } else {
233               avl_swr( &((*rootp)->child[1]) );
234               avl_swl( rootp );
235               avl_nasty( *rootp );
236            }
237         } else {
238            (*rootp)->child[1] = right_subtree;
239         }
240         return False;
241      } else {
242         (*rootp)->child[1] = a;
243         if ((*rootp)->balance++)
244            return False;
245         return True;
246      }
247      tl_assert(0);/*NOTREACHED*/
248   }
249   else {
250      /* cmpres == 0, a duplicate - replace the val, but don't
251         incorporate the node in the tree */
252      oldV->b = True;
253      oldV->w = (*rootp)->val;
254      (*rootp)->val = a->val;
255      return False;
256   }
257}
258
259/* Remove an element a from the AVL tree t.  a must be part of
260   the tree.  Returns True if the depth of the tree has shrunk.
261*/
262static
263Bool avl_remove_wrk ( AvlNode** rootp,
264                      AvlNode*  a,
265                      Word(*kCmp)(UWord,UWord) )
266{
267   Bool ch;
268   Word cmpres;
269   cmpres = kCmp ? /*boxed*/   kCmp( (*rootp)->key, a->key )
270                 : /*unboxed*/ cmp_unsigned_Words( (UWord)(*rootp)->key,
271                                                   (UWord)a->key );
272
273   if (cmpres > 0){
274      /* remove from the left subtree */
275      AvlNode* left_subtree = (*rootp)->child[0];
276      tl_assert(left_subtree);
277      ch = avl_remove_wrk(&left_subtree, a, kCmp);
278      (*rootp)->child[0]=left_subtree;
279      if (ch) {
280         switch ((*rootp)->balance++) {
281            case -1: return True;
282            case  0: return False;
283            case  1: break;
284            default: tl_assert(0);
285         }
286         switch ((*rootp)->child[1]->balance) {
287            case 0:
288               avl_swl( rootp );
289               (*rootp)->balance = -1;
290               (*rootp)->child[0]->balance = 1;
291               return False;
292            case 1:
293               avl_swl( rootp );
294               (*rootp)->balance = 0;
295               (*rootp)->child[0]->balance = 0;
296               return True;
297            case -1:
298               break;
299            default:
300               tl_assert(0);
301         }
302         avl_swr( &((*rootp)->child[1]) );
303         avl_swl( rootp );
304         avl_nasty( *rootp );
305         return True;
306      }
307   }
308   else
309   if (cmpres < 0) {
310      /* remove from the right subtree */
311      AvlNode* right_subtree = (*rootp)->child[1];
312      tl_assert(right_subtree);
313      ch = avl_remove_wrk(&right_subtree, a, kCmp);
314      (*rootp)->child[1] = right_subtree;
315      if (ch) {
316         switch ((*rootp)->balance--) {
317            case  1: return True;
318            case  0: return False;
319            case -1: break;
320            default: tl_assert(0);
321         }
322         switch ((*rootp)->child[0]->balance) {
323            case 0:
324               avl_swr( rootp );
325               (*rootp)->balance = 1;
326               (*rootp)->child[1]->balance = -1;
327               return False;
328            case -1:
329               avl_swr( rootp );
330               (*rootp)->balance = 0;
331               (*rootp)->child[1]->balance = 0;
332               return True;
333            case 1:
334               break;
335            default:
336               tl_assert(0);
337         }
338         avl_swl( &((*rootp)->child[0]) );
339         avl_swr( rootp );
340         avl_nasty( *rootp );
341         return True;
342      }
343   }
344   else {
345      tl_assert(cmpres == 0);
346      tl_assert((*rootp)==a);
347      return avl_removeroot_wrk(rootp, kCmp);
348   }
349   return 0;
350}
351
352/* Remove the root of the AVL tree *rootp.
353 * Warning: dumps core if *rootp is empty
354 */
355static
356Bool avl_removeroot_wrk ( AvlNode** rootp,
357                          Word(*kCmp)(UWord,UWord) )
358{
359   Bool     ch;
360   AvlNode* a;
361   if (!(*rootp)->child[0]) {
362      if (!(*rootp)->child[1]) {
363         (*rootp) = 0;
364         return True;
365      }
366      (*rootp) = (*rootp)->child[1];
367      return True;
368   }
369   if (!(*rootp)->child[1]) {
370      (*rootp) = (*rootp)->child[0];
371      return True;
372   }
373   if ((*rootp)->balance < 0) {
374      /* remove from the left subtree */
375      a = (*rootp)->child[0];
376      while (a->child[1]) a = a->child[1];
377   } else {
378      /* remove from the right subtree */
379      a = (*rootp)->child[1];
380      while (a->child[0]) a = a->child[0];
381   }
382   ch = avl_remove_wrk(rootp, a, kCmp);
383   a->child[0] = (*rootp)->child[0];
384   a->child[1] = (*rootp)->child[1];
385   a->balance  = (*rootp)->balance;
386   (*rootp)    = a;
387   if(a->balance == 0) return ch;
388   return False;
389}
390
391static
392AvlNode* avl_find_node ( AvlNode* t, Word k, Word(*kCmp)(UWord,UWord) )
393{
394   if (kCmp) {
395      /* Boxed comparisons */
396      Word cmpresS;
397      while (True) {
398         if (t == NULL) return NULL;
399         cmpresS = kCmp(t->key, k);
400         if (cmpresS > 0) t = t->child[0]; else
401         if (cmpresS < 0) t = t->child[1]; else
402         return t;
403      }
404   } else {
405      /* Unboxed comparisons */
406      Word  cmpresS; /* signed */
407      UWord cmpresU; /* unsigned */
408      while (True) {
409         if (t == NULL) return NULL; /* unlikely ==> predictable */
410         cmpresS = cmp_unsigned_Words( (UWord)t->key, (UWord)k );
411         if (cmpresS == 0) return t; /* unlikely ==> predictable */
412         cmpresU = (UWord)cmpresS;
413         cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
414         t = t->child[cmpresU];
415      }
416   }
417}
418
419static
420Bool avl_find_bounds ( AvlNode* t,
421                       /*OUT*/UWord* kMinP, /*OUT*/UWord* vMinP,
422                       /*OUT*/UWord* kMaxP, /*OUT*/UWord* vMaxP,
423                       UWord minKey, UWord minVal,
424                       UWord maxKey, UWord maxVal,
425                       UWord key,
426                       Word(*kCmp)(UWord,UWord) )
427{
428   UWord kLowerBound = minKey;
429   UWord vLowerBound = minVal;
430   UWord kUpperBound = maxKey;
431   UWord vUpperBound = maxVal;
432   while (t) {
433      Word cmpresS = kCmp ? kCmp(t->key, key)
434                          : cmp_unsigned_Words(t->key, key);
435      if (cmpresS < 0) {
436         kLowerBound = t->key;
437         vLowerBound = t->val;
438         t = t->child[1];
439         continue;
440      }
441      if (cmpresS > 0) {
442         kUpperBound = t->key;
443         vUpperBound = t->val;
444         t = t->child[0];
445         continue;
446      }
447      /* We should never get here.  If we do, it means the given key
448         is actually present in the tree, which means the original
449         call was invalid -- an error on the caller's part, and we
450         cannot give any meaningful values for the bounds.  (Well,
451         maybe we could, but we're not gonna.  Ner!) */
452      return False;
453   }
454   if (kMinP) *kMinP = kLowerBound;
455   if (vMinP) *vMinP = vLowerBound;
456   if (kMaxP) *kMaxP = kUpperBound;
457   if (vMaxP) *vMaxP = vUpperBound;
458   return True;
459}
460
461// Clear the iterator stack.
462static void stackClear(WordFM* fm)
463{
464   Int i;
465   tl_assert(fm);
466   for (i = 0; i < WFM_STKMAX; i++) {
467      fm->nodeStack[i] = NULL;
468      fm->numStack[i]  = 0;
469   }
470   fm->stackTop = 0;
471}
472
473// Push onto the iterator stack.
474static inline void stackPush(WordFM* fm, AvlNode* n, Int i)
475{
476   tl_assert(fm->stackTop < WFM_STKMAX);
477   tl_assert(1 <= i && i <= 3);
478   fm->nodeStack[fm->stackTop] = n;
479   fm-> numStack[fm->stackTop] = i;
480   fm->stackTop++;
481}
482
483// Pop from the iterator stack.
484static inline Bool stackPop(WordFM* fm, AvlNode** n, Int* i)
485{
486   tl_assert(fm->stackTop <= WFM_STKMAX);
487
488   if (fm->stackTop > 0) {
489      fm->stackTop--;
490      *n = fm->nodeStack[fm->stackTop];
491      *i = fm-> numStack[fm->stackTop];
492      tl_assert(1 <= *i && *i <= 3);
493      fm->nodeStack[fm->stackTop] = NULL;
494      fm-> numStack[fm->stackTop] = 0;
495      return True;
496   } else {
497      return False;
498   }
499}
500
501static
502AvlNode* avl_dopy ( AvlNode* nd,
503                    UWord(*dopyK)(UWord),
504                    UWord(*dopyV)(UWord),
505                    void*(alloc_nofail)(const HChar*,SizeT),
506                    const HChar* cc )
507{
508   AvlNode* nyu;
509   if (! nd)
510      return NULL;
511   nyu = alloc_nofail(cc, sizeof(AvlNode));
512   tl_assert(nyu);
513
514   nyu->child[0] = nd->child[0];
515   nyu->child[1] = nd->child[1];
516   nyu->balance = nd->balance;
517
518   /* Copy key */
519   if (dopyK) {
520      nyu->key = dopyK( nd->key );
521      if (nd->key != 0 && nyu->key == 0)
522         return NULL; /* oom in key dcopy */
523   } else {
524      /* copying assumedly unboxed keys */
525      nyu->key = nd->key;
526   }
527
528   /* Copy val */
529   if (dopyV) {
530      nyu->val = dopyV( nd->val );
531      if (nd->val != 0 && nyu->val == 0)
532         return NULL; /* oom in val dcopy */
533   } else {
534      /* copying assumedly unboxed vals */
535      nyu->val = nd->val;
536   }
537
538   /* Copy subtrees */
539   if (nyu->child[0]) {
540      nyu->child[0] = avl_dopy( nyu->child[0], dopyK, dopyV,
541                                alloc_nofail, cc );
542      if (! nyu->child[0])
543         return NULL;
544   }
545   if (nyu->child[1]) {
546      nyu->child[1] = avl_dopy( nyu->child[1], dopyK, dopyV,
547                                alloc_nofail, cc );
548      if (! nyu->child[1])
549         return NULL;
550   }
551
552   return nyu;
553}
554
555/* Initialise a WordFM. */
556static void initFM ( WordFM* fm,
557                     void*   (*alloc_nofail)( const HChar*, SizeT ),
558                     const HChar* cc,
559                     void    (*dealloc)(void*),
560                     Word    (*kCmp)(UWord,UWord) )
561{
562   fm->root         = 0;
563   fm->kCmp         = kCmp;
564   fm->alloc_nofail = alloc_nofail;
565   fm->cc           = cc;
566   fm->dealloc      = dealloc;
567   fm->stackTop     = 0;
568}
569
570/* --- Public interface functions --- */
571
572/* Allocate and initialise a WordFM.  If kCmp is non-NULL, elements in
573   the set are ordered according to the ordering specified by kCmp,
574   which becomes obvious if you use VG_(initIterFM),
575   VG_(initIterAtFM), VG_(nextIterFM), VG_(doneIterFM) to iterate over
576   sections of the map, or the whole thing.  If kCmp is NULL then the
577   ordering used is unsigned word ordering (UWord) on the key
578   values. */
579WordFM* VG_(newFM) ( void* (*alloc_nofail)( const HChar*, SizeT ),
580                     const HChar* cc,
581                     void  (*dealloc)(void*),
582                     Word  (*kCmp)(UWord,UWord) )
583{
584   WordFM* fm = alloc_nofail(cc, sizeof(WordFM));
585   tl_assert(fm);
586   initFM(fm, alloc_nofail, cc, dealloc, kCmp);
587   return fm;
588}
589
590static void avl_free ( AvlNode* nd,
591                       void(*kFin)(UWord),
592                       void(*vFin)(UWord),
593                       void(*dealloc)(void*) )
594{
595   if (!nd)
596      return;
597   if (nd->child[0])
598      avl_free(nd->child[0], kFin, vFin, dealloc);
599   if (nd->child[1])
600      avl_free(nd->child[1], kFin, vFin, dealloc);
601   if (kFin)
602      kFin( nd->key );
603   if (vFin)
604      vFin( nd->val );
605   VG_(memset)(nd, 0, sizeof(AvlNode));
606   dealloc(nd);
607}
608
609/* Free up the FM.  If kFin is non-NULL, it is applied to keys
610   before the FM is deleted; ditto with vFin for vals. */
611void VG_(deleteFM) ( WordFM* fm, void(*kFin)(UWord), void(*vFin)(UWord) )
612{
613   void(*dealloc)(void*) = fm->dealloc;
614   avl_free( fm->root, kFin, vFin, dealloc );
615   VG_(memset)(fm, 0, sizeof(WordFM) );
616   dealloc(fm);
617}
618
619/* Add (k,v) to fm. */
620Bool VG_(addToFM) ( WordFM* fm, UWord k, UWord v )
621{
622   MaybeWord oldV;
623   AvlNode* node;
624   node = fm->alloc_nofail( fm->cc, sizeof(AvlNode) );
625   node->key = k;
626   node->val = v;
627   oldV.b = False;
628   oldV.w = 0;
629   avl_insert_wrk( &fm->root, &oldV, node, fm->kCmp );
630   //if (oldV.b && fm->vFin)
631   //   fm->vFin( oldV.w );
632   if (oldV.b)
633      fm->dealloc(node);
634   return oldV.b;
635}
636
637// Delete key from fm, returning associated key and val if found
638Bool VG_(delFromFM) ( WordFM* fm,
639                      /*OUT*/UWord* oldK, /*OUT*/UWord* oldV, UWord key )
640{
641   AvlNode* node = avl_find_node( fm->root, key, fm->kCmp );
642   if (node) {
643      avl_remove_wrk( &fm->root, node, fm->kCmp );
644      if (oldK)
645         *oldK = node->key;
646      if (oldV)
647         *oldV = node->val;
648      fm->dealloc(node);
649      return True;
650   } else {
651      return False;
652   }
653}
654
655// Look up in fm, assigning found key & val at spec'd addresses
656Bool VG_(lookupFM) ( WordFM* fm,
657                     /*OUT*/UWord* keyP, /*OUT*/UWord* valP, UWord key )
658{
659   AvlNode* node = avl_find_node( fm->root, key, fm->kCmp );
660   if (node) {
661      if (keyP)
662         *keyP = node->key;
663      if (valP)
664         *valP = node->val;
665      return True;
666   } else {
667      return False;
668   }
669}
670
671// See comment in pub_tool_wordfm.h for explanation
672Bool VG_(findBoundsFM)( WordFM* fm,
673                        /*OUT*/UWord* kMinP, /*OUT*/UWord* vMinP,
674                        /*OUT*/UWord* kMaxP, /*OUT*/UWord* vMaxP,
675                        UWord minKey, UWord minVal,
676                        UWord maxKey, UWord maxVal,
677                        UWord key )
678{
679   /* really we should assert that minKey <= key <= maxKey,
680      where <= is as defined by fm->kCmp. */
681   return avl_find_bounds( fm->root, kMinP, vMinP,
682                                     kMaxP, vMaxP,
683                                     minKey, minVal,
684                                     maxKey, maxVal,
685                                     key, fm->kCmp );
686}
687
688// See comment in pub_tool_wordfm.h for performance warning
689UWord VG_(sizeFM) ( WordFM* fm )
690{
691   // Hmm, this is a bad way to do this
692   return fm->root ? size_avl_nonNull( fm->root ) : 0;
693}
694
695// NB UNTESTED!  TEST BEFORE USE!
696//Bool VG_(isEmptyFM)( WordFM* fm )
697//{
698//   return fm->root ? False : True;
699//}
700
701// set up FM for iteration
702void VG_(initIterFM) ( WordFM* fm )
703{
704   tl_assert(fm);
705   stackClear(fm);
706   if (fm->root)
707      stackPush(fm, fm->root, 1);
708}
709
710// set up FM for iteration so that the first key subsequently produced
711// by VG_(nextIterFM) is the smallest key in the map >= start_at.
712// Naturally ">=" is defined by the comparison function supplied to
713// VG_(newFM), as documented above.
714void VG_(initIterAtFM) ( WordFM* fm, UWord start_at )
715{
716   Int     i;
717   AvlNode *n, *t;
718   Word    cmpresS; /* signed */
719   UWord   cmpresU; /* unsigned */
720
721   tl_assert(fm);
722   stackClear(fm);
723
724   if (!fm->root)
725      return;
726
727   n = NULL;
728   // We need to do regular search and fill in the stack.
729   t = fm->root;
730
731   while (True) {
732      if (t == NULL) return;
733
734      cmpresS
735         = fm->kCmp ? /*boxed*/   fm->kCmp( t->key, start_at )
736                    : /*unboxed*/ cmp_unsigned_Words( t->key, start_at );
737
738      if (cmpresS == 0) {
739         // We found the exact key -- we are done.
740         // The iteration should start with this node.
741         stackPush(fm, t, 2);
742         // The stack now looks like {2, 2, ... ,2, 2}
743         return;
744      }
745      cmpresU = (UWord)cmpresS;
746      cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
747      if (!cmpresU) {
748         // Push this node only if we go to the left child.
749         stackPush(fm, t, 2);
750      }
751      t = t->child[cmpresU];
752   }
753   if (stackPop(fm, &n, &i)) {
754      // If we've pushed something to stack and did not find the exact key,
755      // we must fix the top element of stack.
756      tl_assert(i == 2);
757      stackPush(fm, n, 3);
758      // the stack looks like {2, 2, ..., 2, 3}
759   }
760}
761
762// get next key/val pair.  Will tl_assert if fm has been modified
763// or looked up in since initIter{,At}FM was called.
764Bool VG_(nextIterFM) ( WordFM* fm, /*OUT*/UWord* pKey, /*OUT*/UWord* pVal )
765{
766   Int i = 0;
767   AvlNode* n = NULL;
768
769   tl_assert(fm);
770
771   // This in-order traversal requires each node to be pushed and popped
772   // three times.  These could be avoided by updating nodes in-situ on the
773   // top of the stack, but the push/pop cost is so small that it's worth
774   // keeping this loop in this simpler form.
775   while (stackPop(fm, &n, &i)) {
776      switch (i) {
777      case 1: case_1:
778         stackPush(fm, n, 2);
779         /* if (n->child[0])  stackPush(fm, n->child[0], 1); */
780         if (n->child[0]) { n = n->child[0]; goto case_1; }
781         break;
782      case 2:
783         stackPush(fm, n, 3);
784         if (pKey) *pKey = n->key;
785         if (pVal) *pVal = n->val;
786         return True;
787      case 3:
788         /* if (n->child[1]) stackPush(fm, n->child[1], 1); */
789         if (n->child[1]) { n = n->child[1]; goto case_1; }
790         break;
791      default:
792         tl_assert(0);
793      }
794   }
795
796   // Stack empty, iterator is exhausted, return NULL
797   return False;
798}
799
800// clear the I'm iterating flag
801void VG_(doneIterFM) ( WordFM* fm )
802{
803}
804
805WordFM* VG_(dopyFM) ( WordFM* fm, UWord(*dopyK)(UWord), UWord(*dopyV)(UWord) )
806{
807   WordFM* nyu;
808
809   /* can't clone the fm whilst iterating on it */
810   tl_assert(fm->stackTop == 0);
811
812   nyu = fm->alloc_nofail( fm->cc, sizeof(WordFM) );
813   tl_assert(nyu);
814
815   *nyu = *fm;
816
817   fm->stackTop = 0;
818   VG_(memset)(fm->nodeStack, 0, sizeof(fm->nodeStack));
819   VG_(memset)(fm->numStack, 0,  sizeof(fm->numStack));
820
821   if (nyu->root) {
822      nyu->root = avl_dopy( nyu->root, dopyK, dopyV,
823                            fm->alloc_nofail, fm->cc );
824      if (! nyu->root)
825         return NULL;
826   }
827
828   return nyu;
829}
830
831// admin: what's the 'common' allocation size (for tree nodes?)
832SizeT VG_(getNodeSizeFM)( void )
833{
834   return sizeof(AvlNode);
835}
836
837//------------------------------------------------------------------//
838//---                         end WordFM                         ---//
839//---                       Implementation                       ---//
840//------------------------------------------------------------------//
841
842//------------------------------------------------------------------//
843//---                WordBag (unboxed words only)                ---//
844//---                       Implementation                       ---//
845//------------------------------------------------------------------//
846
847/* A trivial container, to make it opaque. */
848struct _WordBag {
849   WordFM* fm;
850};
851
852WordBag* VG_(newBag) ( void* (*alloc_nofail)( const HChar*, SizeT ),
853                       const HChar* cc,
854                       void  (*dealloc)(void*) )
855{
856   WordBag* bag = alloc_nofail(cc, sizeof(WordBag));
857   bag->fm = VG_(newFM)( alloc_nofail, cc, dealloc, NULL );
858   return bag;
859}
860
861void VG_(deleteBag) ( WordBag* bag )
862{
863   void (*dealloc)(void*) = bag->fm->dealloc;
864   VG_(deleteFM)( bag->fm, NULL, NULL );
865   VG_(memset)(bag, 0, sizeof(WordBag));
866   dealloc(bag);
867}
868
869void VG_(addToBag)( WordBag* bag, UWord w )
870{
871   UWord key, count;
872   if (VG_(lookupFM)(bag->fm, &key, &count, w)) {
873      tl_assert(key == w);
874      tl_assert(count >= 1);
875      VG_(addToFM)(bag->fm, w, count+1);
876   } else {
877      VG_(addToFM)(bag->fm, w, 1);
878   }
879}
880
881UWord VG_(elemBag) ( WordBag* bag, UWord w )
882{
883   UWord key, count;
884   if (VG_(lookupFM)( bag->fm, &key, &count, w)) {
885      tl_assert(key == w);
886      tl_assert(count >= 1);
887      return count;
888   } else {
889      return 0;
890   }
891}
892
893UWord VG_(sizeUniqueBag) ( WordBag* bag )
894{
895   return VG_(sizeFM)( bag->fm );
896}
897
898static UWord sizeTotalBag_wrk ( AvlNode* nd )
899{
900   /* unchecked pre: nd is non-NULL */
901   UWord w = nd->val;
902   tl_assert(w >= 1);
903   if (nd->child[0])
904      w += sizeTotalBag_wrk(nd->child[0]);
905   if (nd->child[1])
906      w += sizeTotalBag_wrk(nd->child[1]);
907   return w;
908}
909UWord VG_(sizeTotalBag)( WordBag* bag )
910{
911   if (bag->fm->root)
912      return sizeTotalBag_wrk(bag->fm->root);
913   else
914      return 0;
915}
916
917Bool VG_(delFromBag)( WordBag* bag, UWord w )
918{
919   UWord key, count;
920   if (VG_(lookupFM)(bag->fm, &key, &count, w)) {
921      tl_assert(key == w);
922      tl_assert(count >= 1);
923      if (count > 1) {
924         VG_(addToFM)(bag->fm, w, count-1);
925      } else {
926         tl_assert(count == 1);
927         VG_(delFromFM)( bag->fm, NULL, NULL, w );
928      }
929      return True;
930   } else {
931      return False;
932   }
933}
934
935Bool VG_(isEmptyBag)( WordBag* bag )
936{
937   return VG_(sizeFM)(bag->fm) == 0;
938}
939
940Bool VG_(isSingletonTotalBag)( WordBag* bag )
941{
942   AvlNode* nd;
943   if (VG_(sizeFM)(bag->fm) != 1)
944      return False;
945   nd = bag->fm->root;
946   tl_assert(nd);
947   tl_assert(!nd->child[0]);
948   tl_assert(!nd->child[1]);
949   return nd->val == 1;
950}
951
952UWord VG_(anyElementOfBag)( WordBag* bag )
953{
954   /* Return an arbitrarily chosen element in the bag.  We might as
955      well return the one at the root of the underlying AVL tree. */
956   AvlNode* nd = bag->fm->root;
957   tl_assert(nd); /* if this fails, 'bag' is empty - caller is in error. */
958   tl_assert(nd->val >= 1);
959   return nd->key;
960}
961
962void VG_(initIterBag)( WordBag* bag )
963{
964   VG_(initIterFM)(bag->fm);
965}
966
967Bool VG_(nextIterBag)( WordBag* bag, /*OUT*/UWord* pVal, /*OUT*/UWord* pCount )
968{
969   return VG_(nextIterFM)( bag->fm, pVal, pCount );
970}
971
972void VG_(doneIterBag)( WordBag* bag )
973{
974   VG_(doneIterFM)( bag->fm );
975}
976
977//------------------------------------------------------------------//
978//---             end WordBag (unboxed words only)               ---//
979//---                       Implementation                       ---//
980//------------------------------------------------------------------//
981
982/*--------------------------------------------------------------------*/
983/*--- end                                               m_wordfm.c ---*/
984/*--------------------------------------------------------------------*/
985