coregrind/m_wordfm.c

   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 */
  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 /*--------------------------------------------------------------------*/