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-2013 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; // allocator
116 const HChar* cc; // cc for allocator
117 OSetFree_t free; // deallocator
118 PoolAlloc* node_pa; // (optional) pool allocator for nodes.
119 SizeT maxEltSize; // for node_pa, must be > 0. Otherwise unused.
120 Word 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(AvlTree * t,AvlNode * n)168 void* slow_key_of_node(AvlTree* t, AvlNode* n)
169 {
170 return (void*)((Addr)elem_of_node(n) + t->keyOff);
171 }
172
173 static inline
fast_key_of_node(AvlNode * n)174 void* fast_key_of_node(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, const HChar* _cc,
290 OSetFree_t _free)
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);
299 vg_assert(_free);
300 if (!_cmp) vg_assert(0 == _keyOff); // If no cmp, offset must be zero
301
302 t = _alloc(_cc, sizeof(AvlTree));
303 t->keyOff = _keyOff;
304 t->cmp = _cmp;
305 t->alloc = _alloc;
306 t->cc = _cc;
307 t->free = _free;
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, const HChar* _cc,
319 OSetFree_t _free,
320 SizeT _poolSize,
321 SizeT _maxEltSize)
322 {
323 AvlTree* t;
324
325 t = VG_(OSetGen_Create) (_keyOff, _cmp,
326 _alloc, _cc,
327 _free);
328
329 vg_assert (_poolSize > 0);
330 vg_assert (_maxEltSize > 0);
331 t->maxEltSize = _maxEltSize;
332 t->node_pa = VG_(newPA)(sizeof(AvlNode)
333 + VG_ROUNDUP(_maxEltSize, sizeof(void*)),
334 _poolSize,
335 t->alloc,
336 _cc,
337 t->free);
338 VG_(addRefPA) (t->node_pa);
339
340 return t;
341 }
342
VG_(OSetGen_EmptyClone)343 AvlTree* VG_(OSetGen_EmptyClone) (AvlTree* os)
344 {
345 AvlTree* t;
346
347 vg_assert(os);
348
349 t = os->alloc(os->cc, sizeof(AvlTree));
350 t->keyOff = os->keyOff;
351 t->cmp = os->cmp;
352 t->alloc = os->alloc;
353 t->cc = os->cc;
354 t->free = os->free;
355 t->node_pa = os->node_pa;
356 if (t->node_pa)
357 VG_(addRefPA) (t->node_pa);
358 t->maxEltSize = os->maxEltSize;
359 t->nElems = 0;
360 t->root = NULL;
361 stackClear(t);
362
363 return t;
364 }
365
VG_(OSetWord_Create)366 AvlTree* VG_(OSetWord_Create)(OSetAlloc_t _alloc, const HChar* _cc,
367 OSetFree_t _free)
368 {
369 return VG_(OSetGen_Create)(/*keyOff*/0, /*cmp*/NULL, _alloc, _cc, _free);
370 }
371
372 // Destructor, frees up all memory held by remaining nodes.
VG_(OSetGen_Destroy)373 void VG_(OSetGen_Destroy)(AvlTree* t)
374 {
375 Bool has_node_pa;
376 vg_assert(t);
377
378 has_node_pa = t->node_pa != NULL;
379
380 /*
381 * If we are the only remaining user of this pool allocator, release all
382 * the elements by deleting the pool allocator. That's more efficient than
383 * deleting tree nodes one by one.
384 */
385 if (!has_node_pa || VG_(releasePA)(t->node_pa) > 0) {
386 AvlNode* n = NULL;
387 Int i = 0;
388 Word sz = 0;
389
390 stackClear(t);
391 if (t->root)
392 stackPush(t, t->root, 1);
393
394 /* Free all the AvlNodes. This is a post-order traversal, because we */
395 /* must free all children of a node before the node itself. */
396 while (stackPop(t, &n, &i)) {
397 switch (i) {
398 case 1:
399 stackPush(t, n, 2);
400 if (n->left) stackPush(t, n->left, 1);
401 break;
402 case 2:
403 stackPush(t, n, 3);
404 if (n->right) stackPush(t, n->right, 1);
405 break;
406 case 3:
407 if (has_node_pa)
408 VG_(freeEltPA) (t->node_pa, n);
409 else
410 t->free(n);
411 sz++;
412 break;
413 }
414 }
415 vg_assert(sz == t->nElems);
416 }
417
418 /* Free the AvlTree itself. */
419 t->free(t);
420 }
421
VG_(OSetWord_Destroy)422 void VG_(OSetWord_Destroy)(AvlTree* t)
423 {
424 VG_(OSetGen_Destroy)(t);
425 }
426
427 // Allocate and initialise a new node.
VG_(OSetGen_AllocNode)428 void* VG_(OSetGen_AllocNode)(AvlTree* t, SizeT elemSize)
429 {
430 AvlNode* n;
431 Int nodeSize = sizeof(AvlNode) + elemSize;
432 vg_assert(elemSize > 0);
433 if (t->node_pa) {
434 vg_assert(elemSize <= t->maxEltSize);
435 n = VG_(allocEltPA) (t->node_pa);
436 } else {
437 n = t->alloc( t->cc, nodeSize );
438 }
439 VG_(memset)(n, 0, nodeSize);
440 n->magic = OSET_MAGIC;
441 return elem_of_node(n);
442 }
443
VG_(OSetGen_FreeNode)444 void VG_(OSetGen_FreeNode)(AvlTree* t, void* e)
445 {
446 if (t->node_pa)
447 VG_(freeEltPA) (t->node_pa, node_of_elem (e));
448 else
449 t->free( node_of_elem(e) );
450 }
451
452 /*--------------------------------------------------------------------*/
453 /*--- Insertion ---*/
454 /*--------------------------------------------------------------------*/
455
cmp_key_root(AvlTree * t,AvlNode * n)456 static inline Word cmp_key_root(AvlTree* t, AvlNode* n)
457 {
458 return t->cmp
459 ? slow_cmp(t, slow_key_of_node(t, n), t->root)
460 : fast_cmp( fast_key_of_node( n), t->root);
461 }
462
463 // Insert element e into the non-empty AVL tree t.
464 // Returns True if the depth of the tree has grown.
avl_insert(AvlTree * t,AvlNode * n)465 static Bool avl_insert(AvlTree* t, AvlNode* n)
466 {
467 Word cmpres = cmp_key_root(t, n);
468
469 if (cmpres < 0) {
470 // Insert into the left subtree.
471 if (t->root->left) {
472 // Only need to set the used fields in the subtree.
473 AvlTree left_subtree;
474 left_subtree.root = t->root->left;
475 left_subtree.cmp = t->cmp;
476 left_subtree.keyOff = t->keyOff;
477 if (avl_insert(&left_subtree, n)) {
478 switch (t->root->balance--) {
479 case 1: return False;
480 case 0: return True;
481 }
482 if (t->root->left->balance < 0) {
483 avl_swr(&(t->root));
484 t->root->balance = 0;
485 t->root->right->balance = 0;
486 } else {
487 avl_swl(&(t->root->left));
488 avl_swr(&(t->root));
489 avl_nasty(t->root);
490 }
491 } else {
492 t->root->left=left_subtree.root;
493 }
494 return False;
495 } else {
496 t->root->left = n;
497 if (t->root->balance--) return False;
498 return True;
499 }
500
501 } else if (cmpres > 0) {
502 // Insert into the right subtree
503 if (t->root->right) {
504 // Only need to set the used fields in the subtree.
505 AvlTree right_subtree;
506 right_subtree.root = t->root->right;
507 right_subtree.cmp = t->cmp;
508 right_subtree.keyOff = t->keyOff;
509 if (avl_insert(&right_subtree, n)) {
510 switch (t->root->balance++) {
511 case -1: return False;
512 case 0: return True;
513 }
514 if (t->root->right->balance > 0) {
515 avl_swl(&(t->root));
516 t->root->balance = 0;
517 t->root->left->balance = 0;
518 } else {
519 avl_swr(&(t->root->right));
520 avl_swl(&(t->root));
521 avl_nasty(t->root);
522 }
523 } else {
524 t->root->right=right_subtree.root;
525 }
526 return False;
527 } else {
528 t->root->right = n;
529 if (t->root->balance++) return False;
530 return True;
531 }
532
533 } else {
534 vg_assert2(0, "OSet{Word,Gen}_Insert: duplicate element added");
535 }
536 }
537
538 // Insert element e into the AVL tree t. This is just a wrapper for
539 // avl_insert() which doesn't return a Bool.
VG_(OSetGen_Insert)540 void VG_(OSetGen_Insert)(AvlTree* t, void* e)
541 {
542 AvlNode* n;
543
544 vg_assert(t);
545
546 // Initialise. Even though OSetGen_AllocNode zeroes these fields,
547 // we should do it again in case a node is removed and then
548 // re-added to the tree.
549 n = node_of_elem(e);
550 n->left = 0;
551 n->right = 0;
552 n->balance = 0;
553
554 // Insert into an empty tree
555 if (!t->root) {
556 t->root = n;
557 } else {
558 avl_insert(t, n);
559 }
560
561 t->nElems++;
562 t->stackTop = 0; // So the iterator can't get out of sync
563 }
564
VG_(OSetWord_Insert)565 void VG_(OSetWord_Insert)(AvlTree* t, UWord val)
566 {
567 Word* node = VG_(OSetGen_AllocNode)(t, sizeof(UWord));
568 *node = val;
569 VG_(OSetGen_Insert)(t, node);
570 }
571
572 /*--------------------------------------------------------------------*/
573 /*--- Lookup ---*/
574 /*--------------------------------------------------------------------*/
575
576 // Find the *node* in t matching k, or NULL if not found.
avl_lookup(const AvlTree * t,const void * k)577 static AvlNode* avl_lookup(const AvlTree* t, const void* k)
578 {
579 Word cmpres;
580 AvlNode* curr = t->root;
581
582 if (t->cmp) {
583 // General case
584 while (True) {
585 if (curr == NULL) return NULL;
586 cmpres = slow_cmp(t, k, curr);
587 if (cmpres < 0) curr = curr->left;
588 else if (cmpres > 0) curr = curr->right;
589 else return curr;
590 }
591 } else {
592 // Fast-track special case. We use the no-check version of
593 // elem_of_node because it saves about 10% on lookup time. This
594 // shouldn't be very dangerous because each node will have been
595 // checked on insertion.
596 UWord w1 = *(const UWord*)k;
597 UWord w2;
598 while (True) {
599 if (curr == NULL) return NULL;
600 w2 = *(UWord*)elem_of_node_no_check(curr);
601 if (w1 < w2) curr = curr->left;
602 else if (w1 > w2) curr = curr->right;
603 else return curr;
604 }
605 }
606 }
607
608 // Find the *element* in t matching k, or NULL if not found.
VG_(OSetGen_Lookup)609 void* VG_(OSetGen_Lookup)(const AvlTree* t, const void* k)
610 {
611 AvlNode* n;
612 vg_assert(t);
613 n = avl_lookup(t, k);
614 return ( n ? elem_of_node(n) : NULL );
615 }
616
617 // Find the *element* in t matching k, or NULL if not found; use the given
618 // comparison function rather than the standard one.
VG_(OSetGen_LookupWithCmp)619 void* VG_(OSetGen_LookupWithCmp)(AvlTree* t, const void* k, OSetCmp_t cmp)
620 {
621 // Save the normal one to the side, then restore once we're done.
622 void* e;
623 OSetCmp_t tmpcmp;
624 vg_assert(t);
625 tmpcmp = t->cmp;
626 t->cmp = cmp;
627 e = VG_(OSetGen_Lookup)(t, k);
628 t->cmp = tmpcmp;
629 return e;
630 }
631
632 // Is there an element matching k?
VG_(OSetGen_Contains)633 Bool VG_(OSetGen_Contains)(const AvlTree* t, const void* k)
634 {
635 return (NULL != VG_(OSetGen_Lookup)(t, k));
636 }
637
VG_(OSetWord_Contains)638 Bool VG_(OSetWord_Contains)(AvlTree* t, UWord val)
639 {
640 return (NULL != VG_(OSetGen_Lookup)(t, &val));
641 }
642
643 /*--------------------------------------------------------------------*/
644 /*--- Deletion ---*/
645 /*--------------------------------------------------------------------*/
646
647 static Bool avl_removeroot(AvlTree* t);
648
649 // Remove an already-selected node n from the AVL tree t.
650 // Returns True if the depth of the tree has shrunk.
avl_remove(AvlTree * t,AvlNode * n)651 static Bool avl_remove(AvlTree* t, AvlNode* n)
652 {
653 Bool ch;
654 Word cmpres = cmp_key_root(t, n);
655
656 if (cmpres < 0) {
657 AvlTree left_subtree;
658 // Remove from the left subtree
659 vg_assert(t->root->left);
660 // Only need to set the used fields in the subtree.
661 left_subtree.root = t->root->left;
662 left_subtree.cmp = t->cmp;
663 left_subtree.keyOff = t->keyOff;
664 ch = avl_remove(&left_subtree, n);
665 t->root->left = left_subtree.root;
666 if (ch) {
667 switch (t->root->balance++) {
668 case -1: return True;
669 case 0: return False;
670 }
671 switch (t->root->right->balance) {
672 case 0:
673 avl_swl(&(t->root));
674 t->root->balance = -1;
675 t->root->left->balance = 1;
676 return False;
677 case 1:
678 avl_swl(&(t->root));
679 t->root->balance = 0;
680 t->root->left->balance = 0;
681 return True;
682 }
683 avl_swr(&(t->root->right));
684 avl_swl(&(t->root));
685 avl_nasty(t->root);
686 return True;
687 } else {
688 return False;
689 }
690
691 } else if (cmpres > 0) {
692 // Remove from the right subtree
693 AvlTree right_subtree;
694 vg_assert(t->root->right);
695 // Only need to set the used fields in the subtree.
696 right_subtree.root = t->root->right;
697 right_subtree.cmp = t->cmp;
698 right_subtree.keyOff = t->keyOff;
699 ch = avl_remove(&right_subtree, n);
700 t->root->right = right_subtree.root;
701 if (ch) {
702 switch (t->root->balance--) {
703 case 1: return True;
704 case 0: return False;
705 }
706 switch (t->root->left->balance) {
707 case 0:
708 avl_swr(&(t->root));
709 t->root->balance = 1;
710 t->root->right->balance = -1;
711 return False;
712 case -1:
713 avl_swr(&(t->root));
714 t->root->balance = 0;
715 t->root->right->balance = 0;
716 return True;
717 }
718 avl_swl(&(t->root->left));
719 avl_swr(&(t->root));
720 avl_nasty(t->root);
721 return True;
722 } else {
723 return False;
724 }
725
726 } else {
727 // Found the node to be removed.
728 vg_assert(t->root == n);
729 return avl_removeroot(t);
730 }
731 }
732
733 // Remove the root of the AVL tree t.
734 // Returns True if the depth of the tree has shrunk.
avl_removeroot(AvlTree * t)735 static Bool avl_removeroot(AvlTree* t)
736 {
737 Bool ch;
738 AvlNode* n;
739
740 if (!t->root->left) {
741 if (!t->root->right) {
742 t->root = NULL;
743 return True;
744 }
745 t->root = t->root->right;
746 return True;
747 }
748 if (!t->root->right) {
749 t->root = t->root->left;
750 return True;
751 }
752 if (t->root->balance < 0) {
753 // Remove from the left subtree
754 n = t->root->left;
755 while (n->right) n = n->right;
756 } else {
757 // Remove from the right subtree
758 n = t->root->right;
759 while (n->left) n = n->left;
760 }
761 ch = avl_remove(t, n);
762 n->left = t->root->left;
763 n->right = t->root->right;
764 n->balance = t->root->balance;
765 t->root = n;
766 if (n->balance == 0) return ch;
767 return False;
768 }
769
770 // Remove and return the element matching the key 'k', or NULL
771 // if not present.
VG_(OSetGen_Remove)772 void* VG_(OSetGen_Remove)(AvlTree* t, const void* k)
773 {
774 // Have to find the node first, then remove it.
775 AvlNode* n = avl_lookup(t, k);
776 if (n) {
777 avl_remove(t, n);
778 t->nElems--;
779 t->stackTop = 0; // So the iterator can't get out of sync
780 return elem_of_node(n);
781 } else {
782 return NULL;
783 }
784 }
785
VG_(OSetWord_Remove)786 Bool VG_(OSetWord_Remove)(AvlTree* t, UWord val)
787 {
788 void* n = VG_(OSetGen_Remove)(t, &val);
789 if (n) {
790 VG_(OSetGen_FreeNode)(t, n);
791 return True;
792 } else {
793 return False;
794 }
795 }
796
797 /*--------------------------------------------------------------------*/
798 /*--- Iterator ---*/
799 /*--------------------------------------------------------------------*/
800
801 // The iterator is implemented using in-order traversal with an explicit
802 // stack, which lets us do the traversal one step at a time and remember
803 // where we are between each call to OSetGen_Next().
804
VG_(OSetGen_ResetIter)805 void VG_(OSetGen_ResetIter)(AvlTree* t)
806 {
807 vg_assert(t);
808 stackClear(t);
809 if (t->root)
810 stackPush(t, t->root, 1);
811 }
812
VG_(OSetWord_ResetIter)813 void VG_(OSetWord_ResetIter)(AvlTree* t)
814 {
815 VG_(OSetGen_ResetIter)(t);
816 }
817
VG_(OSetGen_Next)818 void* VG_(OSetGen_Next)(AvlTree* t)
819 {
820 Int i = 0;
821 OSetNode* n = NULL;
822
823 vg_assert(t);
824
825 // This in-order traversal requires each node to be pushed and popped
826 // three times. These could be avoided by updating nodes in-situ on the
827 // top of the stack, but the push/pop cost is so small that it's worth
828 // keeping this loop in this simpler form.
829 while (stackPop(t, &n, &i)) {
830 switch (i) {
831 case 1: case_1:
832 stackPush(t, n, 2);
833 /* if (n->left) stackPush(t, n->left, 1); */
834 if (n->left) { n = n->left; goto case_1; }
835 break;
836 case 2:
837 stackPush(t, n, 3);
838 return elem_of_node(n);
839 case 3:
840 /* if (n->right) stackPush(t, n->right, 1); */
841 if (n->right) { n = n->right; goto case_1; }
842 break;
843 }
844 }
845
846 // Stack empty, iterator is exhausted, return NULL
847 return NULL;
848 }
849
VG_(OSetWord_Next)850 Bool VG_(OSetWord_Next)(AvlTree* t, UWord* val)
851 {
852 UWord* n = VG_(OSetGen_Next)(t);
853 if (n) {
854 *val = *n;
855 return True;
856 } else {
857 return False;
858 }
859 }
860
861 // set up 'oset' for iteration so that the first key subsequently
862 // produced VG_(OSetGen_Next) is the smallest key in the map
863 // >= start_at. Naturally ">=" is defined by the comparison
864 // function supplied to VG_(OSetGen_Create).
VG_(OSetGen_ResetIterAt)865 void VG_(OSetGen_ResetIterAt)(AvlTree* oset, const void* k)
866 {
867 AvlNode *t;
868 Word cmpresS; /* signed */
869 UWord cmpresU; /* unsigned */
870
871 vg_assert(oset);
872 stackClear(oset);
873
874 if (!oset->root)
875 return;
876
877 // We need to do regular search and fill in the stack.
878 t = oset->root;
879
880 while (True) {
881 if (t == NULL) return;
882
883 if (oset->cmp) {
884 cmpresS = (Word)slow_cmp(oset, k, t);
885 } else {
886 cmpresS = fast_cmp(k, t);
887 }
888
889 /* Switch the sense of the comparison, since the comparison
890 order of args (k vs t) above is opposite to that of the
891 corresponding code in hg_wordfm.c. */
892 if (cmpresS < 0) { cmpresS = 1; }
893 else if (cmpresS > 0) { cmpresS = -1; }
894
895 if (cmpresS == 0) {
896 // We found the exact key -- we are done.
897 // The iteration should start with this node.
898 stackPush(oset, t, 2);
899 // The stack now looks like {2, 2, ... ,2, 2}
900 return;
901 }
902 cmpresU = (UWord)cmpresS;
903 cmpresU >>=/*unsigned*/ (8 * sizeof(cmpresU) - 1);
904 vg_assert(cmpresU == 0 || cmpresU == 1);
905 if (!cmpresU) {
906 // Push this node only if we go to the left child.
907 stackPush(oset, t, 2);
908 }
909 t = cmpresU==0 ? t->left : t->right;
910 }
911 }
912
913 /*--------------------------------------------------------------------*/
914 /*--- Miscellaneous operations ---*/
915 /*--------------------------------------------------------------------*/
916
VG_(OSetGen_Size)917 Word VG_(OSetGen_Size)(const AvlTree* t)
918 {
919 vg_assert(t);
920 return t->nElems;
921 }
922
VG_(OSetWord_Size)923 Word VG_(OSetWord_Size)(AvlTree* t)
924 {
925 return VG_(OSetGen_Size)(t);
926 }
927
OSet_Print2(AvlTree * t,AvlNode * n,HChar * (* strElem)(void *),Int p)928 static void OSet_Print2( AvlTree* t, AvlNode* n,
929 HChar*(*strElem)(void *), Int p )
930 {
931 // This is a recursive in-order traversal.
932 Int q = p;
933 if (NULL == n) return;
934 if (n->right) OSet_Print2(t, n->right, strElem, p+1);
935 while (q--) VG_(printf)(".. ");
936 VG_(printf)("%s\n", strElem(elem_of_node(n)));
937 if (n->left) OSet_Print2(t, n->left, strElem, p+1);
938 }
939
940 __attribute__((unused))
OSet_Print(AvlTree * t,const HChar * where,HChar * (* strElem)(void *))941 static void OSet_Print( AvlTree* t, const HChar *where,
942 HChar*(*strElem)(void *) )
943 {
944 VG_(printf)("-- start %s ----------------\n", where);
945 OSet_Print2(t, t->root, strElem, 0);
946 VG_(printf)("-- end %s ----------------\n", where);
947 }
948
949 /*--------------------------------------------------------------------*/
950 /*--- end ---*/
951 /*--------------------------------------------------------------------*/
952