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-2015 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-2015 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         vg_assert(0);
134   }
135   root->balance=0;
136}
137
138/* Find size of a non-NULL tree. */
139static UWord size_avl_nonNull ( const 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: vg_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      vg_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: vg_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      vg_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      vg_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: vg_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               vg_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      vg_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: vg_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               vg_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      vg_assert(cmpres == 0);
346      vg_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 ( const 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   vg_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   vg_assert(fm->stackTop < WFM_STKMAX);
477   vg_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   vg_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      vg_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 ( const 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
510   vg_assert(nd != NULL);
511
512   nyu = alloc_nofail(cc, sizeof(AvlNode));
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   } else {
522      /* copying assumedly unboxed keys */
523      nyu->key = nd->key;
524   }
525
526   /* Copy val */
527   if (dopyV) {
528      nyu->val = dopyV( nd->val );
529   } else {
530      /* copying assumedly unboxed vals */
531      nyu->val = nd->val;
532   }
533
534   /* Copy subtrees */
535   if (nyu->child[0]) {
536      nyu->child[0] = avl_dopy( nyu->child[0], dopyK, dopyV,
537                                alloc_nofail, cc );
538   }
539   if (nyu->child[1]) {
540      nyu->child[1] = avl_dopy( nyu->child[1], dopyK, dopyV,
541                                alloc_nofail, cc );
542   }
543
544   return nyu;
545}
546
547/* Initialise a WordFM. */
548static void initFM ( WordFM* fm,
549                     void*   (*alloc_nofail)( const HChar*, SizeT ),
550                     const HChar* cc,
551                     void    (*dealloc)(void*),
552                     Word    (*kCmp)(UWord,UWord) )
553{
554   fm->root         = 0;
555   fm->kCmp         = kCmp;
556   fm->alloc_nofail = alloc_nofail;
557   fm->cc           = cc;
558   fm->dealloc      = dealloc;
559   fm->stackTop     = 0;
560}
561
562/* --- Public interface functions --- */
563
564/* Allocate and initialise a WordFM.  If kCmp is non-NULL, elements in
565   the set are ordered according to the ordering specified by kCmp,
566   which becomes obvious if you use VG_(initIterFM),
567   VG_(initIterAtFM), VG_(nextIterFM), VG_(doneIterFM) to iterate over
568   sections of the map, or the whole thing.  If kCmp is NULL then the
569   ordering used is unsigned word ordering (UWord) on the key
570   values. */
571WordFM* VG_(newFM) ( void* (*alloc_nofail)( const HChar*, SizeT ),
572                     const HChar* cc,
573                     void  (*dealloc)(void*),
574                     Word  (*kCmp)(UWord,UWord) )
575{
576   WordFM* fm = alloc_nofail(cc, sizeof(WordFM));
577   initFM(fm, alloc_nofail, cc, dealloc, kCmp);
578   return fm;
579}
580
581static void avl_free ( AvlNode* nd,
582                       void(*kFin)(UWord),
583                       void(*vFin)(UWord),
584                       void(*dealloc)(void*) )
585{
586   if (!nd)
587      return;
588   if (nd->child[0])
589      avl_free(nd->child[0], kFin, vFin, dealloc);
590   if (nd->child[1])
591      avl_free(nd->child[1], kFin, vFin, dealloc);
592   if (kFin)
593      kFin( nd->key );
594   if (vFin)
595      vFin( nd->val );
596   VG_(memset)(nd, 0, sizeof(AvlNode));
597   dealloc(nd);
598}
599
600/* Free up the FM.  If kFin is non-NULL, it is applied to keys
601   before the FM is deleted; ditto with vFin for vals. */
602void VG_(deleteFM) ( WordFM* fm, void(*kFin)(UWord), void(*vFin)(UWord) )
603{
604   void(*dealloc)(void*) = fm->dealloc;
605   avl_free( fm->root, kFin, vFin, dealloc );
606   VG_(memset)(fm, 0, sizeof(WordFM) );
607   dealloc(fm);
608}
609
610/* Add (k,v) to fm. */
611Bool VG_(addToFM) ( WordFM* fm, UWord k, UWord v )
612{
613   MaybeWord oldV;
614   AvlNode* node;
615   node = fm->alloc_nofail( fm->cc, sizeof(AvlNode) );
616   node->key = k;
617   node->val = v;
618   oldV.b = False;
619   oldV.w = 0;
620   avl_insert_wrk( &fm->root, &oldV, node, fm->kCmp );
621   //if (oldV.b && fm->vFin)
622   //   fm->vFin( oldV.w );
623   if (oldV.b)
624      fm->dealloc(node);
625   return oldV.b;
626}
627
628// Delete key from fm, returning associated key and val if found
629Bool VG_(delFromFM) ( WordFM* fm,
630                      /*OUT*/UWord* oldK, /*OUT*/UWord* oldV, UWord key )
631{
632   AvlNode* node = avl_find_node( fm->root, key, fm->kCmp );
633   if (node) {
634      avl_remove_wrk( &fm->root, node, fm->kCmp );
635      if (oldK)
636         *oldK = node->key;
637      if (oldV)
638         *oldV = node->val;
639      fm->dealloc(node);
640      return True;
641   } else {
642      return False;
643   }
644}
645
646// Look up in fm, assigning found key & val at spec'd addresses
647Bool VG_(lookupFM) ( const WordFM* fm,
648                     /*OUT*/UWord* keyP, /*OUT*/UWord* valP, UWord key )
649{
650   AvlNode* node = avl_find_node( fm->root, key, fm->kCmp );
651   if (node) {
652      if (keyP)
653         *keyP = node->key;
654      if (valP)
655         *valP = node->val;
656      return True;
657   } else {
658      return False;
659   }
660}
661
662// See comment in pub_tool_wordfm.h for explanation
663Bool VG_(findBoundsFM)( const WordFM* fm,
664                        /*OUT*/UWord* kMinP, /*OUT*/UWord* vMinP,
665                        /*OUT*/UWord* kMaxP, /*OUT*/UWord* vMaxP,
666                        UWord minKey, UWord minVal,
667                        UWord maxKey, UWord maxVal,
668                        UWord key )
669{
670   /* really we should assert that minKey <= key <= maxKey,
671      where <= is as defined by fm->kCmp. */
672   return avl_find_bounds( fm->root, kMinP, vMinP,
673                                     kMaxP, vMaxP,
674                                     minKey, minVal,
675                                     maxKey, maxVal,
676                                     key, fm->kCmp );
677}
678
679// See comment in pub_tool_wordfm.h for performance warning
680UWord VG_(sizeFM) ( const WordFM* fm )
681{
682   // Hmm, this is a bad way to do this
683   return fm->root ? size_avl_nonNull( fm->root ) : 0;
684}
685
686// NB UNTESTED!  TEST BEFORE USE!
687//Bool VG_(isEmptyFM)( const WordFM* fm )
688//{
689//   return fm->root ? False : True;
690//}
691
692// set up FM for iteration
693void VG_(initIterFM) ( WordFM* fm )
694{
695   vg_assert(fm);
696   stackClear(fm);
697   if (fm->root)
698      stackPush(fm, fm->root, 1);
699}
700
701// set up FM for iteration so that the first key subsequently produced
702// by VG_(nextIterFM) is the smallest key in the map >= start_at.
703// Naturally ">=" is defined by the comparison function supplied to
704// VG_(newFM), as documented above.
705void VG_(initIterAtFM) ( WordFM* fm, UWord start_at )
706{
707   Int     i;
708   AvlNode *n, *t;
709   Word    cmpresS; /* signed */
710   UWord   cmpresU; /* unsigned */
711
712   vg_assert(fm);
713   stackClear(fm);
714
715   if (!fm->root)
716      return;
717
718   n = NULL;
719   // We need to do regular search and fill in the stack.
720   t = fm->root;
721
722   while (True) {
723      if (t == NULL) return;
724
725      cmpresS
726         = fm->kCmp ? /*boxed*/   fm->kCmp( t->key, start_at )
727                    : /*unboxed*/ cmp_unsigned_Words( t->key, start_at );
728
729      if (cmpresS == 0) {
730         // We found the exact key -- we are done.
731         // The iteration should start with this node.
732         stackPush(fm, t, 2);
733         // The stack now looks like {2, 2, ... ,2, 2}
734         return;
735      }
736      cmpresU = (UWord)cmpresS;
737      cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
738      if (!cmpresU) {
739         // Push this node only if we go to the left child.
740         stackPush(fm, t, 2);
741      }
742      t = t->child[cmpresU];
743   }
744   if (stackPop(fm, &n, &i)) {
745      // If we've pushed something to stack and did not find the exact key,
746      // we must fix the top element of stack.
747      vg_assert(i == 2);
748      stackPush(fm, n, 3);
749      // the stack looks like {2, 2, ..., 2, 3}
750   }
751}
752
753// get next key/val pair.  Will vg_assert if fm has been modified
754// or looked up in since initIter{,At}FM was called.
755Bool VG_(nextIterFM) ( WordFM* fm, /*OUT*/UWord* pKey, /*OUT*/UWord* pVal )
756{
757   Int i = 0;
758   AvlNode* n = NULL;
759
760   vg_assert(fm);
761
762   // This in-order traversal requires each node to be pushed and popped
763   // three times.  These could be avoided by updating nodes in-situ on the
764   // top of the stack, but the push/pop cost is so small that it's worth
765   // keeping this loop in this simpler form.
766   while (stackPop(fm, &n, &i)) {
767      switch (i) {
768      case 1: case_1:
769         stackPush(fm, n, 2);
770         /* if (n->child[0])  stackPush(fm, n->child[0], 1); */
771         if (n->child[0]) { n = n->child[0]; goto case_1; }
772         break;
773      case 2:
774         stackPush(fm, n, 3);
775         if (pKey) *pKey = n->key;
776         if (pVal) *pVal = n->val;
777         return True;
778      case 3:
779         /* if (n->child[1]) stackPush(fm, n->child[1], 1); */
780         if (n->child[1]) { n = n->child[1]; goto case_1; }
781         break;
782      default:
783         vg_assert(0);
784      }
785   }
786
787   // Stack empty, iterator is exhausted, return NULL
788   return False;
789}
790
791// Finish an FM iteration
792void VG_(doneIterFM) ( WordFM* fm )
793{
794}
795
796WordFM* VG_(dopyFM) ( WordFM* fm, UWord(*dopyK)(UWord),
797                      UWord(*dopyV)(UWord) )
798{
799   WordFM* nyu;
800
801   /* can't clone the fm whilst iterating on it */
802   vg_assert(fm->stackTop == 0);
803
804   nyu = fm->alloc_nofail( fm->cc, sizeof(WordFM) );
805
806   *nyu = *fm;
807
808   fm->stackTop = 0;
809   VG_(memset)(fm->nodeStack, 0, sizeof(fm->nodeStack));
810   VG_(memset)(fm->numStack, 0,  sizeof(fm->numStack));
811
812   if (nyu->root) {
813      nyu->root = avl_dopy( nyu->root, dopyK, dopyV,
814                            fm->alloc_nofail, fm->cc );
815      if (! nyu->root)
816         return NULL;
817   }
818
819   return nyu;
820}
821
822//------------------------------------------------------------------//
823//---                         end WordFM                         ---//
824//---                       Implementation                       ---//
825//------------------------------------------------------------------//
826
827//------------------------------------------------------------------//
828//---                WordBag (unboxed words only)                ---//
829//---                       Implementation                       ---//
830//------------------------------------------------------------------//
831
832/* A trivial container, to make it opaque. */
833struct _WordBag {
834   WordFM* fm;
835};
836
837WordBag* VG_(newBag) ( void* (*alloc_nofail)( const HChar*, SizeT ),
838                       const HChar* cc,
839                       void  (*dealloc)(void*) )
840{
841   WordBag* bag = alloc_nofail(cc, sizeof(WordBag));
842   bag->fm = VG_(newFM)( alloc_nofail, cc, dealloc, NULL );
843   return bag;
844}
845
846void VG_(deleteBag) ( WordBag* bag )
847{
848   void (*dealloc)(void*) = bag->fm->dealloc;
849   VG_(deleteFM)( bag->fm, NULL, NULL );
850   VG_(memset)(bag, 0, sizeof(WordBag));
851   dealloc(bag);
852}
853
854void VG_(addToBag)( WordBag* bag, UWord w )
855{
856   UWord key, count;
857   if (VG_(lookupFM)(bag->fm, &key, &count, w)) {
858      vg_assert(key == w);
859      vg_assert(count >= 1);
860      VG_(addToFM)(bag->fm, w, count+1);
861   } else {
862      VG_(addToFM)(bag->fm, w, 1);
863   }
864}
865
866UWord VG_(elemBag) ( const WordBag* bag, UWord w )
867{
868   UWord key, count;
869   if (VG_(lookupFM)( bag->fm, &key, &count, w)) {
870      vg_assert(key == w);
871      vg_assert(count >= 1);
872      return count;
873   } else {
874      return 0;
875   }
876}
877
878UWord VG_(sizeUniqueBag) ( const WordBag* bag )
879{
880   return VG_(sizeFM)( bag->fm );
881}
882
883static UWord sizeTotalBag_wrk ( const AvlNode* nd )
884{
885   /* unchecked pre: nd is non-NULL */
886   UWord w = nd->val;
887   vg_assert(w >= 1);
888   if (nd->child[0])
889      w += sizeTotalBag_wrk(nd->child[0]);
890   if (nd->child[1])
891      w += sizeTotalBag_wrk(nd->child[1]);
892   return w;
893}
894UWord VG_(sizeTotalBag)( const WordBag* bag )
895{
896   if (bag->fm->root)
897      return sizeTotalBag_wrk(bag->fm->root);
898   else
899      return 0;
900}
901
902Bool VG_(delFromBag)( WordBag* bag, UWord w )
903{
904   UWord key, count;
905   if (VG_(lookupFM)(bag->fm, &key, &count, w)) {
906      vg_assert(key == w);
907      vg_assert(count >= 1);
908      if (count > 1) {
909         VG_(addToFM)(bag->fm, w, count-1);
910      } else {
911         vg_assert(count == 1);
912         VG_(delFromFM)( bag->fm, NULL, NULL, w );
913      }
914      return True;
915   } else {
916      return False;
917   }
918}
919
920Bool VG_(isEmptyBag)( const WordBag* bag )
921{
922   return VG_(sizeFM)(bag->fm) == 0;
923}
924
925Bool VG_(isSingletonTotalBag)( const WordBag* bag )
926{
927   AvlNode* nd;
928   if (VG_(sizeFM)(bag->fm) != 1)
929      return False;
930   nd = bag->fm->root;
931   vg_assert(nd);
932   vg_assert(!nd->child[0]);
933   vg_assert(!nd->child[1]);
934   return nd->val == 1;
935}
936
937UWord VG_(anyElementOfBag)( const WordBag* bag )
938{
939   /* Return an arbitrarily chosen element in the bag.  We might as
940      well return the one at the root of the underlying AVL tree. */
941   AvlNode* nd = bag->fm->root;
942   vg_assert(nd); /* if this fails, 'bag' is empty - caller is in error. */
943   vg_assert(nd->val >= 1);
944   return nd->key;
945}
946
947void VG_(initIterBag)( WordBag* bag )
948{
949   VG_(initIterFM)(bag->fm);
950}
951
952Bool VG_(nextIterBag)( WordBag* bag, /*OUT*/UWord* pVal, /*OUT*/UWord* pCount )
953{
954   return VG_(nextIterFM)( bag->fm, pVal, pCount );
955}
956
957void VG_(doneIterBag)( WordBag* bag )
958{
959   VG_(doneIterFM)( bag->fm );
960}
961
962//------------------------------------------------------------------//
963//---             end WordBag (unboxed words only)               ---//
964//---                       Implementation                       ---//
965//------------------------------------------------------------------//
966
967/*--------------------------------------------------------------------*/
968/*--- end                                               m_wordfm.c ---*/
969/*--------------------------------------------------------------------*/
970