1
2 /*--------------------------------------------------------------------*/
3 /*--- An ordered set implemented using an AVL tree. m_oset.c ---*/
4 /*--------------------------------------------------------------------*/
5
6 /*
7 This file is part of Valgrind, a dynamic binary instrumentation
8 framework.
9
10 Copyright (C) 2005-2015 Nicholas Nethercote
11 njn@valgrind.org
12
13 This program is free software; you can redistribute it and/or
14 modify it under the terms of the GNU General Public License as
15 published by the Free Software Foundation; either version 2 of the
16 License, or (at your option) any later version.
17
18 This program is distributed in the hope that it will be useful, but
19 WITHOUT ANY WARRANTY; without even the implied warranty of
20 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
21 General Public License for more details.
22
23 You should have received a copy of the GNU General Public License
24 along with this program; if not, write to the Free Software
25 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
26 02111-1307, USA.
27
28 The GNU General Public License is contained in the file COPYING.
29 */
30
31 //----------------------------------------------------------------------
32 // This file is based on:
33 //
34 // ANSI C Library for maintainance of AVL Balanced Trees
35 // (C) 2000 Daniel Nagy, Budapest University of Technology and Economics
36 // Released under GNU General Public License (GPL) version 2
37 //----------------------------------------------------------------------
38
39 // This file implements a generic ordered set using an AVL tree.
40 //
41 // Each node in the tree has two parts.
42 // - First is the AVL metadata, which is three words: a left pointer, a
43 // right pointer, and a word containing balancing information and a
44 // "magic" value which provides some checking that the user has not
45 // corrupted the metadata. So the overhead is 12 bytes on 32-bit
46 // platforms and 24 bytes on 64-bit platforms.
47 // - Second is the user's data. This can be anything. Note that because it
48 // comes after the metadata, it will only be word-aligned, even if the
49 // user data is a struct that would normally be doubleword-aligned.
50 //
51 // AvlNode* node -> +---------------+ V
52 // | struct |
53 // | AvlNode |
54 // void* element -> +---------------+ ^
55 // | element | |
56 // keyOff -> | key | elemSize
57 // +---------------+ v
58 //
59 // Users have to allocate AvlNodes with OSetGen_AllocNode(), which allocates
60 // space for the metadata.
61 //
62 // The terminology used throughout this file:
63 // - a "node", usually called "n", is a pointer to the metadata.
64 // - an "element", usually called "e", is a pointer to the user data.
65 // - a "key", usually called "k", is a pointer to a key.
66 //
67 // The helper functions elem_of_node and node_of_elem do the pointer
68 // arithmetic to switch between the node and the element. The node magic is
69 // checked after each operation to make sure that we're really operating on
70 // an AvlNode.
71 //
72 // Each tree also has an iterator. Note that we cannot use the iterator
73 // internally within this file (eg. we could implement OSetGen_Size() by
74 // stepping through with the iterator and counting nodes) because it's
75 // non-reentrant -- the user might be using it themselves, and the
76 // concurrent uses would screw things up.
77
78 #include "pub_core_basics.h"
79 #include "pub_core_libcbase.h"
80 #include "pub_core_libcassert.h"
81 #include "pub_core_libcprint.h"
82 #include "pub_core_oset.h"
83 #include "pub_core_poolalloc.h"
84
85 /*--------------------------------------------------------------------*/
86 /*--- Types and constants ---*/
87 /*--------------------------------------------------------------------*/
88
89 typedef struct _OSetNode OSetNode;
90
91 // Internal names for the OSet types.
92 typedef OSet AvlTree;
93 typedef OSetNode AvlNode;
94
95 // The padding ensures that magic is right at the end of the node,
96 // regardless of the machine's word size, so that any overwrites will be
97 // detected earlier.
98 struct _OSetNode {
99 AvlNode* left;
100 AvlNode* right;
101 Char balance;
102 Char padding[sizeof(void*)-sizeof(Char)-sizeof(Short)];
103 Short magic;
104 };
105
106 #define STACK_MAX 32 // At most 2**32 entries can be iterated over
107 #define OSET_MAGIC 0x5b1f
108
109 // An OSet (AVL tree). If cmp is NULL, the key must be a UWord, and must
110 // be the first word in the element. If cmp is set, arbitrary keys in
111 // arbitrary positions can be used.
112 struct _OSet {
113 SizeT keyOff; // key offset
114 OSetCmp_t cmp; // compare a key and an element, or NULL
115 OSetAlloc_t alloc_fn; // allocator
116 const HChar* cc; // cost centre for allocator
117 OSetFree_t free_fn; // deallocator
118 PoolAlloc* node_pa; // (optional) pool allocator for nodes.
119 SizeT maxEltSize; // for node_pa, must be > 0. Otherwise unused.
120 UInt nElems; // number of elements in the tree
121 AvlNode* root; // root node
122
123 AvlNode* nodeStack[STACK_MAX]; // Iterator node stack
124 Int numStack[STACK_MAX]; // Iterator num stack
125 Int stackTop; // Iterator stack pointer, one past end
126 };
127
128 /*--------------------------------------------------------------------*/
129 /*--- Helper operations ---*/
130 /*--------------------------------------------------------------------*/
131
132 // Given a pointer to the node's element, return the pointer to the AvlNode
133 // structure. If the node has a bad magic number, it will die with an
134 // assertion failure.
135 static inline
node_of_elem(const void * elem)136 AvlNode* node_of_elem(const void *elem)
137 {
138 AvlNode* n = (AvlNode*)((Addr)elem - sizeof(AvlNode));
139 vg_assert2(n->magic == OSET_MAGIC,
140 "bad magic on node %p = %x (expected %x)\n"
141 "possible causes:\n"
142 " - node not allocated with VG_(OSetGen_AllocNode)()?\n"
143 " - node metadata corrupted by underwriting start of element?\n",
144 n, n->magic, OSET_MAGIC);
145 return n;
146 }
147
148 // Given an AvlNode, return the pointer to the element.
149 static inline
elem_of_node(const AvlNode * n)150 void* elem_of_node(const AvlNode *n)
151 {
152 vg_assert2(n->magic == OSET_MAGIC,
153 "bad magic on node %p = %x (expected %x)\n"
154 "possible causes:\n"
155 " - node metadata corrupted by overwriting end of element?\n",
156 n, n->magic, OSET_MAGIC);
157 return (void*)((Addr)n + sizeof(AvlNode));
158 }
159
160 // Like elem_of_node, but no magic checking.
161 static inline
elem_of_node_no_check(const AvlNode * n)162 void* elem_of_node_no_check(const AvlNode *n)
163 {
164 return (void*)((Addr)n + sizeof(AvlNode));
165 }
166
167 static inline
slow_key_of_node(const AvlTree * t,const AvlNode * n)168 void* slow_key_of_node(const AvlTree* t, const AvlNode* n)
169 {
170 return (void*)((Addr)elem_of_node(n) + t->keyOff);
171 }
172
173 static inline
fast_key_of_node(const AvlNode * n)174 void* fast_key_of_node(const AvlNode* n)
175 {
176 return elem_of_node(n);
177 }
178
179 // Compare the first word of each element. Inlining is *crucial*.
fast_cmp(const void * k,const AvlNode * n)180 static inline Word fast_cmp(const void* k, const AvlNode* n)
181 {
182 UWord w1 = *(const UWord*)k;
183 UWord w2 = *(const UWord*)elem_of_node(n);
184 // In previous versions, we tried to do this faster by doing
185 // "return w1 - w2". But it didn't work reliably, because the
186 // complete result of subtracting two N-bit numbers is an N+1-bit
187 // number, and what the caller is interested in is the sign of
188 // the complete N+1-bit result. The branching version is slightly
189 // slower, but safer and easier to understand.
190 if (w1 > w2) return 1;
191 if (w1 < w2) return -1;
192 return 0;
193 }
194
195 // Compare a key and an element. Inlining is *crucial*.
196 static
slow_cmp(const AvlTree * t,const void * k,const AvlNode * n)197 inline Word slow_cmp(const AvlTree* t, const void* k, const AvlNode* n)
198 {
199 return t->cmp(k, elem_of_node(n));
200 }
201
202
203 // Swing to the left. Warning: no balance maintainance.
avl_swl(AvlNode ** root)204 static void avl_swl ( AvlNode** root )
205 {
206 AvlNode* a = *root;
207 AvlNode* b = a->right;
208 *root = b;
209 a->right = b->left;
210 b->left = a;
211 }
212
213 // Swing to the right. Warning: no balance maintainance.
avl_swr(AvlNode ** root)214 static void avl_swr ( AvlNode** root )
215 {
216 AvlNode* a = *root;
217 AvlNode* b = a->left;
218 *root = b;
219 a->left = b->right;
220 b->right = a;
221 }
222
223 // Balance maintainance after especially nasty swings.
avl_nasty(AvlNode * root)224 static void avl_nasty ( AvlNode* root )
225 {
226 switch (root->balance) {
227 case -1:
228 root->left->balance = 0;
229 root->right->balance = 1;
230 break;
231 case 1:
232 root->left->balance =-1;
233 root->right->balance = 0;
234 break;
235 case 0:
236 root->left->balance = 0;
237 root->right->balance = 0;
238 }
239 root->balance = 0;
240 }
241
242
243 // Clear the iterator stack.
stackClear(AvlTree * t)244 static void stackClear(AvlTree* t)
245 {
246 Int i;
247 vg_assert(t);
248 for (i = 0; i < STACK_MAX; i++) {
249 t->nodeStack[i] = NULL;
250 t->numStack[i] = 0;
251 }
252 t->stackTop = 0;
253 }
254
255 // Push onto the iterator stack.
stackPush(AvlTree * t,AvlNode * n,Int i)256 static inline void stackPush(AvlTree* t, AvlNode* n, Int i)
257 {
258 vg_assert(t->stackTop < STACK_MAX);
259 vg_assert(1 <= i && i <= 3);
260 t->nodeStack[t->stackTop] = n;
261 t-> numStack[t->stackTop] = i;
262 t->stackTop++;
263 }
264
265 // Pop from the iterator stack.
stackPop(AvlTree * t,AvlNode ** n,Int * i)266 static inline Bool stackPop(AvlTree* t, AvlNode** n, Int* i)
267 {
268 vg_assert(t->stackTop <= STACK_MAX);
269
270 if (t->stackTop > 0) {
271 t->stackTop--;
272 *n = t->nodeStack[t->stackTop];
273 *i = t-> numStack[t->stackTop];
274 vg_assert(1 <= *i && *i <= 3);
275 t->nodeStack[t->stackTop] = NULL;
276 t-> numStack[t->stackTop] = 0;
277 return True;
278 } else {
279 return False;
280 }
281 }
282
283 /*--------------------------------------------------------------------*/
284 /*--- Creating and destroying AvlTrees and AvlNodes ---*/
285 /*--------------------------------------------------------------------*/
286
287 // The underscores avoid GCC complaints about overshadowing global names.
VG_(OSetGen_Create)288 AvlTree* VG_(OSetGen_Create)(PtrdiffT keyOff, OSetCmp_t cmp,
289 OSetAlloc_t alloc_fn, const HChar* cc,
290 OSetFree_t free_fn)
291 {
292 AvlTree* t;
293
294 // Check the padding is right and the AvlNode is the expected size.
295 vg_assert(sizeof(AvlNode) == 3*sizeof(void*));
296
297 // Sanity check args
298 vg_assert(alloc_fn);
299 vg_assert(free_fn);
300 if (!cmp) vg_assert(0 == keyOff); // If no cmp, offset must be zero
301
302 t = alloc_fn(cc, sizeof(AvlTree));
303 t->keyOff = keyOff;
304 t->cmp = cmp;
305 t->alloc_fn = alloc_fn;
306 t->cc = cc;
307 t->free_fn = free_fn;
308 t->node_pa = NULL;
309 t->maxEltSize = 0; // Just in case it would be wrongly used.
310 t->nElems = 0;
311 t->root = NULL;
312 stackClear(t);
313
314 return t;
315 }
316
VG_(OSetGen_Create_With_Pool)317 AvlTree* VG_(OSetGen_Create_With_Pool)(PtrdiffT keyOff, OSetCmp_t cmp,
318 OSetAlloc_t alloc_fn, const HChar* cc,
319 OSetFree_t free_fn,
320 SizeT poolSize,
321 SizeT maxEltSize)
322 {
323 AvlTree* t;
324
325 t = VG_(OSetGen_Create) (keyOff, cmp, alloc_fn, cc, free_fn);
326
327 vg_assert (poolSize > 0);
328 vg_assert (maxEltSize > 0);
329 t->maxEltSize = maxEltSize;
330 t->node_pa = VG_(newPA)(sizeof(AvlNode)
331 + VG_ROUNDUP(maxEltSize, sizeof(void*)),
332 poolSize,
333 t->alloc_fn,
334 cc,
335 t->free_fn);
336 VG_(addRefPA) (t->node_pa);
337
338 return t;
339 }
340
VG_(OSetGen_EmptyClone)341 AvlTree* VG_(OSetGen_EmptyClone) (const AvlTree* os)
342 {
343 AvlTree* t;
344
345 vg_assert(os);
346
347 t = os->alloc_fn(os->cc, sizeof(AvlTree));
348 t->keyOff = os->keyOff;
349 t->cmp = os->cmp;
350 t->alloc_fn = os->alloc_fn;
351 t->cc = os->cc;
352 t->free_fn = os->free_fn;
353 t->node_pa = os->node_pa;
354 if (t->node_pa)
355 VG_(addRefPA) (t->node_pa);
356 t->maxEltSize = os->maxEltSize;
357 t->nElems = 0;
358 t->root = NULL;
359 stackClear(t);
360
361 return t;
362 }
363
VG_(OSetWord_Create)364 AvlTree* VG_(OSetWord_Create)(OSetAlloc_t alloc_fn, const HChar* cc,
365 OSetFree_t free_fn)
366 {
367 return VG_(OSetGen_Create)(/*keyOff*/0, /*cmp*/NULL, alloc_fn, cc, free_fn);
368 }
369
370 // Destructor, frees up all memory held by remaining nodes.
VG_(OSetGen_Destroy)371 void VG_(OSetGen_Destroy)(AvlTree* t)
372 {
373 Bool has_node_pa;
374 vg_assert(t);
375
376 has_node_pa = t->node_pa != NULL;
377
378 /*
379 * If we are the only remaining user of this pool allocator, release all
380 * the elements by deleting the pool allocator. That's more efficient than
381 * deleting tree nodes one by one.
382 */
383 if (!has_node_pa || VG_(releasePA)(t->node_pa) > 0) {
384 AvlNode* n = NULL;
385 Int i = 0;
386 UWord sz = 0;
387
388 stackClear(t);
389 if (t->root)
390 stackPush(t, t->root, 1);
391
392 /* Free all the AvlNodes. This is a post-order traversal, because we */
393 /* must free all children of a node before the node itself. */
394 while (stackPop(t, &n, &i)) {
395 switch (i) {
396 case 1:
397 stackPush(t, n, 2);
398 if (n->left) stackPush(t, n->left, 1);
399 break;
400 case 2:
401 stackPush(t, n, 3);
402 if (n->right) stackPush(t, n->right, 1);
403 break;
404 case 3:
405 if (has_node_pa)
406 VG_(freeEltPA) (t->node_pa, n);
407 else
408 t->free_fn(n);
409 sz++;
410 break;
411 }
412 }
413 vg_assert(sz == t->nElems);
414 }
415
416 /* Free the AvlTree itself. */
417 t->free_fn(t);
418 }
419
VG_(OSetWord_Destroy)420 void VG_(OSetWord_Destroy)(AvlTree* t)
421 {
422 VG_(OSetGen_Destroy)(t);
423 }
424
425 // Allocate and initialise a new node.
VG_(OSetGen_AllocNode)426 void* VG_(OSetGen_AllocNode)(const AvlTree* t, SizeT elemSize)
427 {
428 AvlNode* n;
429 Int nodeSize = sizeof(AvlNode) + elemSize;
430 vg_assert(elemSize > 0);
431 if (t->node_pa) {
432 vg_assert(elemSize <= t->maxEltSize);
433 n = VG_(allocEltPA) (t->node_pa);
434 } else {
435 n = t->alloc_fn( t->cc, nodeSize );
436 }
437 VG_(memset)(n, 0, nodeSize);
438 n->magic = OSET_MAGIC;
439 return elem_of_node(n);
440 }
441
VG_(OSetGen_FreeNode)442 void VG_(OSetGen_FreeNode)(const AvlTree* t, void* e)
443 {
444 if (t->node_pa)
445 VG_(freeEltPA) (t->node_pa, node_of_elem (e));
446 else
447 t->free_fn( node_of_elem(e) );
448 }
449
450 /*--------------------------------------------------------------------*/
451 /*--- Insertion ---*/
452 /*--------------------------------------------------------------------*/
453
cmp_key_root(const AvlTree * t,const AvlNode * n)454 static inline Word cmp_key_root(const AvlTree* t, const AvlNode* n)
455 {
456 return t->cmp
457 ? slow_cmp(t, slow_key_of_node(t, n), t->root)
458 : fast_cmp( fast_key_of_node( n), t->root);
459 }
460
461 // Insert element e into the non-empty AVL tree t.
462 // Returns True if the depth of the tree has grown.
avl_insert(AvlTree * t,AvlNode * n)463 static Bool avl_insert(AvlTree* t, AvlNode* n)
464 {
465 Word cmpres = cmp_key_root(t, n);
466
467 if (cmpres < 0) {
468 // Insert into the left subtree.
469 if (t->root->left) {
470 // Only need to set the used fields in the subtree.
471 AvlTree left_subtree;
472 left_subtree.root = t->root->left;
473 left_subtree.cmp = t->cmp;
474 left_subtree.keyOff = t->keyOff;
475 if (avl_insert(&left_subtree, n)) {
476 switch (t->root->balance--) {
477 case 1: return False;
478 case 0: return True;
479 }
480 if (t->root->left->balance < 0) {
481 avl_swr(&(t->root));
482 t->root->balance = 0;
483 t->root->right->balance = 0;
484 } else {
485 avl_swl(&(t->root->left));
486 avl_swr(&(t->root));
487 avl_nasty(t->root);
488 }
489 } else {
490 t->root->left=left_subtree.root;
491 }
492 return False;
493 } else {
494 t->root->left = n;
495 if (t->root->balance--) return False;
496 return True;
497 }
498
499 } else if (cmpres > 0) {
500 // Insert into the right subtree
501 if (t->root->right) {
502 // Only need to set the used fields in the subtree.
503 AvlTree right_subtree;
504 right_subtree.root = t->root->right;
505 right_subtree.cmp = t->cmp;
506 right_subtree.keyOff = t->keyOff;
507 if (avl_insert(&right_subtree, n)) {
508 switch (t->root->balance++) {
509 case -1: return False;
510 case 0: return True;
511 }
512 if (t->root->right->balance > 0) {
513 avl_swl(&(t->root));
514 t->root->balance = 0;
515 t->root->left->balance = 0;
516 } else {
517 avl_swr(&(t->root->right));
518 avl_swl(&(t->root));
519 avl_nasty(t->root);
520 }
521 } else {
522 t->root->right=right_subtree.root;
523 }
524 return False;
525 } else {
526 t->root->right = n;
527 if (t->root->balance++) return False;
528 return True;
529 }
530
531 } else {
532 vg_assert2(0, "OSet{Word,Gen}_Insert: duplicate element added");
533 }
534 }
535
536 // Insert element e into the AVL tree t. This is just a wrapper for
537 // avl_insert() which doesn't return a Bool.
VG_(OSetGen_Insert)538 void VG_(OSetGen_Insert)(AvlTree* t, void* e)
539 {
540 AvlNode* n;
541
542 vg_assert(t);
543
544 // Initialise. Even though OSetGen_AllocNode zeroes these fields,
545 // we should do it again in case a node is removed and then
546 // re-added to the tree.
547 n = node_of_elem(e);
548 n->left = 0;
549 n->right = 0;
550 n->balance = 0;
551
552 // Insert into an empty tree
553 if (!t->root) {
554 t->root = n;
555 } else {
556 avl_insert(t, n);
557 }
558
559 t->nElems++;
560 t->stackTop = 0; // So the iterator can't get out of sync
561 }
562
VG_(OSetWord_Insert)563 void VG_(OSetWord_Insert)(AvlTree* t, UWord val)
564 {
565 Word* node = VG_(OSetGen_AllocNode)(t, sizeof(UWord));
566 *node = val;
567 VG_(OSetGen_Insert)(t, node);
568 }
569
570 /*--------------------------------------------------------------------*/
571 /*--- Lookup ---*/
572 /*--------------------------------------------------------------------*/
573
574 // Find the *node* in t matching k, or NULL if not found.
avl_lookup(const AvlTree * t,const void * k)575 static AvlNode* avl_lookup(const AvlTree* t, const void* k)
576 {
577 Word cmpres;
578 AvlNode* curr = t->root;
579
580 if (t->cmp) {
581 // General case
582 while (True) {
583 if (curr == NULL) return NULL;
584 cmpres = slow_cmp(t, k, curr);
585 if (cmpres < 0) curr = curr->left;
586 else if (cmpres > 0) curr = curr->right;
587 else return curr;
588 }
589 } else {
590 // Fast-track special case. We use the no-check version of
591 // elem_of_node because it saves about 10% on lookup time. This
592 // shouldn't be very dangerous because each node will have been
593 // checked on insertion.
594 UWord w1 = *(const UWord*)k;
595 UWord w2;
596 while (True) {
597 if (curr == NULL) return NULL;
598 w2 = *(UWord*)elem_of_node_no_check(curr);
599 if (w1 < w2) curr = curr->left;
600 else if (w1 > w2) curr = curr->right;
601 else return curr;
602 }
603 }
604 }
605
606 // Find the *element* in t matching k, or NULL if not found.
VG_(OSetGen_Lookup)607 void* VG_(OSetGen_Lookup)(const AvlTree* t, const void* k)
608 {
609 AvlNode* n;
610 vg_assert(t);
611 n = avl_lookup(t, k);
612 return ( n ? elem_of_node(n) : NULL );
613 }
614
615 // Find the *element* in t matching k, or NULL if not found; use the given
616 // comparison function rather than the standard one.
VG_(OSetGen_LookupWithCmp)617 void* VG_(OSetGen_LookupWithCmp)(AvlTree* t, const void* k, OSetCmp_t cmp)
618 {
619 // Save the normal one to the side, then restore once we're done.
620 void* e;
621 OSetCmp_t tmpcmp;
622 vg_assert(t);
623 tmpcmp = t->cmp;
624 t->cmp = cmp;
625 e = VG_(OSetGen_Lookup)(t, k);
626 t->cmp = tmpcmp;
627 return e;
628 }
629
630 // Is there an element matching k?
VG_(OSetGen_Contains)631 Bool VG_(OSetGen_Contains)(const AvlTree* t, const void* k)
632 {
633 return (NULL != VG_(OSetGen_Lookup)(t, k));
634 }
635
VG_(OSetWord_Contains)636 Bool VG_(OSetWord_Contains)(const AvlTree* t, UWord val)
637 {
638 return (NULL != VG_(OSetGen_Lookup)(t, &val));
639 }
640
641 /*--------------------------------------------------------------------*/
642 /*--- Deletion ---*/
643 /*--------------------------------------------------------------------*/
644
645 static Bool avl_removeroot(AvlTree* t);
646
647 // Remove an already-selected node n from the AVL tree t.
648 // Returns True if the depth of the tree has shrunk.
avl_remove(AvlTree * t,const AvlNode * n)649 static Bool avl_remove(AvlTree* t, const AvlNode* n)
650 {
651 Bool ch;
652 Word cmpres = cmp_key_root(t, n);
653
654 if (cmpres < 0) {
655 AvlTree left_subtree;
656 // Remove from the left subtree
657 vg_assert(t->root->left);
658 // Only need to set the used fields in the subtree.
659 left_subtree.root = t->root->left;
660 left_subtree.cmp = t->cmp;
661 left_subtree.keyOff = t->keyOff;
662 ch = avl_remove(&left_subtree, n);
663 t->root->left = left_subtree.root;
664 if (ch) {
665 switch (t->root->balance++) {
666 case -1: return True;
667 case 0: return False;
668 }
669 switch (t->root->right->balance) {
670 case 0:
671 avl_swl(&(t->root));
672 t->root->balance = -1;
673 t->root->left->balance = 1;
674 return False;
675 case 1:
676 avl_swl(&(t->root));
677 t->root->balance = 0;
678 t->root->left->balance = 0;
679 return True;
680 }
681 avl_swr(&(t->root->right));
682 avl_swl(&(t->root));
683 avl_nasty(t->root);
684 return True;
685 } else {
686 return False;
687 }
688
689 } else if (cmpres > 0) {
690 // Remove from the right subtree
691 AvlTree right_subtree;
692 vg_assert(t->root->right);
693 // Only need to set the used fields in the subtree.
694 right_subtree.root = t->root->right;
695 right_subtree.cmp = t->cmp;
696 right_subtree.keyOff = t->keyOff;
697 ch = avl_remove(&right_subtree, n);
698 t->root->right = right_subtree.root;
699 if (ch) {
700 switch (t->root->balance--) {
701 case 1: return True;
702 case 0: return False;
703 }
704 switch (t->root->left->balance) {
705 case 0:
706 avl_swr(&(t->root));
707 t->root->balance = 1;
708 t->root->right->balance = -1;
709 return False;
710 case -1:
711 avl_swr(&(t->root));
712 t->root->balance = 0;
713 t->root->right->balance = 0;
714 return True;
715 }
716 avl_swl(&(t->root->left));
717 avl_swr(&(t->root));
718 avl_nasty(t->root);
719 return True;
720 } else {
721 return False;
722 }
723
724 } else {
725 // Found the node to be removed.
726 vg_assert(t->root == n);
727 return avl_removeroot(t);
728 }
729 }
730
731 // Remove the root of the AVL tree t.
732 // Returns True if the depth of the tree has shrunk.
avl_removeroot(AvlTree * t)733 static Bool avl_removeroot(AvlTree* t)
734 {
735 Bool ch;
736 AvlNode* n;
737
738 if (!t->root->left) {
739 if (!t->root->right) {
740 t->root = NULL;
741 return True;
742 }
743 t->root = t->root->right;
744 return True;
745 }
746 if (!t->root->right) {
747 t->root = t->root->left;
748 return True;
749 }
750 if (t->root->balance < 0) {
751 // Remove from the left subtree
752 n = t->root->left;
753 while (n->right) n = n->right;
754 } else {
755 // Remove from the right subtree
756 n = t->root->right;
757 while (n->left) n = n->left;
758 }
759 ch = avl_remove(t, n);
760 n->left = t->root->left;
761 n->right = t->root->right;
762 n->balance = t->root->balance;
763 t->root = n;
764 if (n->balance == 0) return ch;
765 return False;
766 }
767
768 // Remove and return the element matching the key 'k', or NULL
769 // if not present.
VG_(OSetGen_Remove)770 void* VG_(OSetGen_Remove)(AvlTree* t, const void* k)
771 {
772 // Have to find the node first, then remove it.
773 AvlNode* n = avl_lookup(t, k);
774 if (n) {
775 avl_remove(t, n);
776 t->nElems--;
777 t->stackTop = 0; // So the iterator can't get out of sync
778 return elem_of_node(n);
779 } else {
780 return NULL;
781 }
782 }
783
VG_(OSetWord_Remove)784 Bool VG_(OSetWord_Remove)(AvlTree* t, UWord val)
785 {
786 void* n = VG_(OSetGen_Remove)(t, &val);
787 if (n) {
788 VG_(OSetGen_FreeNode)(t, n);
789 return True;
790 } else {
791 return False;
792 }
793 }
794
795 /*--------------------------------------------------------------------*/
796 /*--- Iterator ---*/
797 /*--------------------------------------------------------------------*/
798
799 // The iterator is implemented using in-order traversal with an explicit
800 // stack, which lets us do the traversal one step at a time and remember
801 // where we are between each call to OSetGen_Next().
802
VG_(OSetGen_ResetIter)803 void VG_(OSetGen_ResetIter)(AvlTree* t)
804 {
805 vg_assert(t);
806 stackClear(t);
807 if (t->root)
808 stackPush(t, t->root, 1);
809 }
810
VG_(OSetWord_ResetIter)811 void VG_(OSetWord_ResetIter)(AvlTree* t)
812 {
813 VG_(OSetGen_ResetIter)(t);
814 }
815
VG_(OSetGen_Next)816 void* VG_(OSetGen_Next)(AvlTree* t)
817 {
818 Int i = 0;
819 OSetNode* n = NULL;
820
821 vg_assert(t);
822
823 // This in-order traversal requires each node to be pushed and popped
824 // three times. These could be avoided by updating nodes in-situ on the
825 // top of the stack, but the push/pop cost is so small that it's worth
826 // keeping this loop in this simpler form.
827 while (stackPop(t, &n, &i)) {
828 switch (i) {
829 case 1: case_1:
830 stackPush(t, n, 2);
831 /* if (n->left) stackPush(t, n->left, 1); */
832 if (n->left) { n = n->left; goto case_1; }
833 break;
834 case 2:
835 stackPush(t, n, 3);
836 return elem_of_node(n);
837 case 3:
838 /* if (n->right) stackPush(t, n->right, 1); */
839 if (n->right) { n = n->right; goto case_1; }
840 break;
841 }
842 }
843
844 // Stack empty, iterator is exhausted, return NULL
845 return NULL;
846 }
847
VG_(OSetWord_Next)848 Bool VG_(OSetWord_Next)(AvlTree* t, UWord* val)
849 {
850 UWord* n = VG_(OSetGen_Next)(t);
851 if (n) {
852 *val = *n;
853 return True;
854 } else {
855 return False;
856 }
857 }
858
859 // set up 'oset' for iteration so that the first key subsequently
860 // produced VG_(OSetGen_Next) is the smallest key in the map
861 // >= start_at. Naturally ">=" is defined by the comparison
862 // function supplied to VG_(OSetGen_Create).
VG_(OSetGen_ResetIterAt)863 void VG_(OSetGen_ResetIterAt)(AvlTree* oset, const void* k)
864 {
865 AvlNode *t;
866 Word cmpresS; /* signed */
867 UWord cmpresU; /* unsigned */
868
869 vg_assert(oset);
870 stackClear(oset);
871
872 if (!oset->root)
873 return;
874
875 // We need to do regular search and fill in the stack.
876 t = oset->root;
877
878 while (True) {
879 if (t == NULL) return;
880
881 if (oset->cmp) {
882 cmpresS = (Word)slow_cmp(oset, k, t);
883 } else {
884 cmpresS = fast_cmp(k, t);
885 }
886
887 /* Switch the sense of the comparison, since the comparison
888 order of args (k vs t) above is opposite to that of the
889 corresponding code in hg_wordfm.c. */
890 if (cmpresS < 0) { cmpresS = 1; }
891 else if (cmpresS > 0) { cmpresS = -1; }
892
893 if (cmpresS == 0) {
894 // We found the exact key -- we are done.
895 // The iteration should start with this node.
896 stackPush(oset, t, 2);
897 // The stack now looks like {2, 2, ... ,2, 2}
898 return;
899 }
900 cmpresU = (UWord)cmpresS;
901 cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
902 vg_assert(cmpresU == 0 || cmpresU == 1);
903 if (!cmpresU) {
904 // Push this node only if we go to the left child.
905 stackPush(oset, t, 2);
906 }
907 t = cmpresU==0 ? t->left : t->right;
908 }
909 }
910
911 /*--------------------------------------------------------------------*/
912 /*--- Miscellaneous operations ---*/
913 /*--------------------------------------------------------------------*/
914
VG_(OSetGen_Size)915 UInt VG_(OSetGen_Size)(const AvlTree* t)
916 {
917 vg_assert(t);
918 return t->nElems;
919 }
920
VG_(OSetWord_Size)921 Word VG_(OSetWord_Size)(const AvlTree* t)
922 {
923 return VG_(OSetGen_Size)(t);
924 }
925
OSet_Print2(const AvlTree * t,const AvlNode * n,const HChar * (* strElem)(const void *),Int p)926 static void OSet_Print2( const AvlTree* t, const AvlNode* n,
927 const HChar*(*strElem)(const void *), Int p )
928 {
929 // This is a recursive in-order traversal.
930 Int q = p;
931 if (NULL == n) return;
932 if (n->right) OSet_Print2(t, n->right, strElem, p+1);
933 while (q--) VG_(printf)(".. ");
934 VG_(printf)("%s\n", strElem(elem_of_node(n)));
935 if (n->left) OSet_Print2(t, n->left, strElem, p+1);
936 }
937
938 __attribute__((unused))
OSet_Print(const AvlTree * t,const HChar * where,const HChar * (* strElem)(const void *))939 static void OSet_Print( const AvlTree* t, const HChar *where,
940 const HChar*(*strElem)(const void *) )
941 {
942 VG_(printf)("-- start %s ----------------\n", where);
943 OSet_Print2(t, t->root, strElem, 0);
944 VG_(printf)("-- end %s ----------------\n", where);
945 }
946
947 /*--------------------------------------------------------------------*/
948 /*--- end ---*/
949 /*--------------------------------------------------------------------*/
950