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