Home | History | Annotate | Download | only in coregrind
      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 (at) 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 (at) 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 */
     64 typedef
     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 
     73 typedef
     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 
     82 struct _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 */
     94 static Bool avl_removeroot_wrk(AvlNode** t, Word(*kCmp)(UWord,UWord));
     95 
     96 /* Swing to the left.  Warning: no balance maintainance. */
     97 static 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. */
    107 static 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. */
    117 static 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. */
    139 static 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. */
    148 static 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 */
    160 static
    161 Bool 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 */
    262 static
    263 Bool 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  */
    355 static
    356 Bool 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 
    391 static
    392 AvlNode* 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 
    419 static
    420 Bool 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.
    462 static 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.
    474 static 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.
    484 static 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 
    501 static
    502 AvlNode* 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. */
    548 static 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. */
    571 WordFM* 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 
    581 static 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. */
    602 void 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. */
    611 Bool 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
    629 Bool 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
    647 Bool 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
    663 Bool 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
    680 UWord 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
    693 void 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.
    705 void 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.
    755 Bool 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
    792 void VG_(doneIterFM) ( WordFM* fm )
    793 {
    794 }
    795 
    796 WordFM* 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. */
    833 struct _WordBag {
    834    WordFM* fm;
    835 };
    836 
    837 WordBag* 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 
    846 void 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 
    854 void 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 
    866 UWord 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 
    878 UWord VG_(sizeUniqueBag) ( const WordBag* bag )
    879 {
    880    return VG_(sizeFM)( bag->fm );
    881 }
    882 
    883 static 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 }
    894 UWord 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 
    902 Bool 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 
    920 Bool VG_(isEmptyBag)( const WordBag* bag )
    921 {
    922    return VG_(sizeFM)(bag->fm) == 0;
    923 }
    924 
    925 Bool 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 
    937 UWord 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 
    947 void VG_(initIterBag)( WordBag* bag )
    948 {
    949    VG_(initIterFM)(bag->fm);
    950 }
    951 
    952 Bool VG_(nextIterBag)( WordBag* bag, /*OUT*/UWord* pVal, /*OUT*/UWord* pCount )
    953 {
    954    return VG_(nextIterFM)( bag->fm, pVal, pCount );
    955 }
    956 
    957 void 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