1 /*
2 * This file is part of the Yices SMT Solver.
3 * Copyright (C) 2017 SRI International.
4 *
5 * Yices is free software: you can redistribute it and/or modify
6 * it under the terms of the GNU General Public License as published by
7 * the Free Software Foundation, either version 3 of the License, or
8 * (at your option) any later version.
9 *
10 * Yices is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 * GNU General Public License for more details.
14 *
15 * You should have received a copy of the GNU General Public License
16 * along with Yices. If not, see <http://www.gnu.org/licenses/>.
17 */
18
19 /*
20 * RED-BLACK TREES TO STORE SETS OF UNSIGNED 32-BIT INTEGERS
21 */
22
23 #include "utils/memalloc.h"
24 #include "utils/uint_rbtrees.h"
25
26
27 /*
28 * Initialize tree
29 * - n = initial size
30 * - if n = 0, the default size is used
31 */
init_rbtree(rbtree_t * tree,uint32_t n)32 void init_rbtree(rbtree_t *tree, uint32_t n) {
33 if (n == 0) {
34 n = DEF_RBTREE_SIZE;
35 }
36 if (n >= MAX_RBTREE_SIZE) {
37 out_of_memory();
38 }
39
40 tree->data = (uint32_t *) safe_malloc(n * sizeof(uint32_t));
41 tree->node = (rbnode_t *) safe_malloc(n * sizeof(rbnode_t));
42 tree->isred = allocate_bitvector(n);
43 tree->size = n;
44 init_ivector(&tree->stack, 20);
45
46 assert(n > 0 && null_rbnode == 0);
47
48 // initialize the null_node
49 tree->data[0] = 0;
50 tree->node[0][0] = 0;
51 tree->node[0][1] = 0;
52 clr_bit(tree->isred, 0); // black node
53
54 tree->nbnodes = 1;
55 tree->root = 0;
56 }
57
58
59 /*
60 * Resize the tree: make it 50% larger
61 */
extend_rbtree(rbtree_t * tree)62 static void extend_rbtree(rbtree_t *tree) {
63 uint32_t n;
64
65 n = tree->size + 1;
66 n += n>>1;
67 if (n >= MAX_RBTREE_SIZE) {
68 out_of_memory();
69 }
70
71 tree->data = (uint32_t *) safe_realloc(tree->data, n * sizeof(uint32_t));
72 tree->node = (rbnode_t *) safe_realloc(tree->node, n * sizeof(rbnode_t));
73 tree->isred = extend_bitvector(tree->isred, n);
74 tree->size = n;
75 }
76
77
78 /*
79 * Allocate a new node and return its id
80 * - the node is not initialized
81 */
rbtree_new_node(rbtree_t * tree)82 static uint32_t rbtree_new_node(rbtree_t *tree) {
83 uint32_t i;
84
85 i = tree->nbnodes;
86 if (i == tree->size) {
87 extend_rbtree(tree);
88 }
89 assert(i < tree->size);
90 tree->nbnodes = i+1;
91
92 return i;
93 }
94
95
96 /*
97 * Delete tree: free memory
98 */
delete_rbtree(rbtree_t * tree)99 void delete_rbtree(rbtree_t *tree) {
100 safe_free(tree->data);
101 safe_free(tree->node);
102 delete_bitvector(tree->isred);
103 delete_ivector(&tree->stack);
104 tree->data = NULL;
105 tree->node = NULL;
106 tree->isred = NULL;
107 }
108
109
110 /*
111 * Reset: empty the tree
112 */
reset_rbtree(rbtree_t * tree)113 void reset_rbtree(rbtree_t *tree) {
114 tree->nbnodes = 1;
115 tree->root = null_rbnode;
116 ivector_reset(&tree->stack);
117 }
118
119
120 /*
121 * Search for a node of value x
122 * - return its id or null_rbnode if there's no such node
123 */
rbtree_find(rbtree_t * tree,uint32_t x)124 uint32_t rbtree_find(rbtree_t *tree, uint32_t x) {
125 uint32_t i, k;
126
127 // to force termination: store x in the null_node
128 tree->data[0] = x;
129
130 i = tree->root;
131 assert(i < tree->nbnodes);
132 while (tree->data[i] != x) {
133 k = (tree->data[i] < x);
134 assert((k == 0 && x < tree->data[i]) ||
135 (k == 1 && x > tree->data[i]));
136 i = tree->node[i][k];
137 assert(i < tree->nbnodes);
138 }
139
140 return i;
141 }
142
143
144
145 /*
146 * Auxiliary functions for balancing the tree
147 */
148
149 /*
150 * Check whether p is parent of q
151 * - both must be valid node indices
152 */
153 #ifndef NDEBUG
is_parent_node(rbtree_t * tree,uint32_t p,uint32_t q)154 static inline bool is_parent_node(rbtree_t *tree, uint32_t p, uint32_t q) {
155 assert(p < tree->nbnodes && q < tree->nbnodes);
156 return tree->node[p][0] == q || tree->node[p][1] == q;
157 }
158 #endif
159
160 /*
161 * Child-index(p, q):
162 * - q must be a child of node p
163 * - returns 0 if q is the left child
164 * returns 1 if q is the right child
165 * So i = child_index(treee, p, q) implies tree->node[p][i] = q
166 */
child_index(rbtree_t * tree,uint32_t p,uint32_t q)167 static inline uint32_t child_index(rbtree_t *tree, uint32_t p, uint32_t q) {
168 assert(is_parent_node(tree, p, q));
169 return tree->node[p][1] == q;
170 }
171
172
173 /*
174 * Get sibling of q in p
175 * - both p and q must be valid node indices
176 * - q must be a child of p
177 */
sibling(rbtree_t * tree,uint32_t p,uint32_t q)178 static inline uint32_t sibling(rbtree_t *tree, uint32_t p, uint32_t q) {
179 assert(is_parent_node(tree, p, q));
180 return (tree->node[p][0] ^ tree->node[p][1]) ^ q;
181 }
182
183
184 /*
185 * Check color of node p
186 */
is_red(rbtree_t * tree,uint32_t p)187 static inline bool is_red(rbtree_t *tree, uint32_t p) {
188 assert(p < tree->nbnodes);
189 return tst_bit(tree->isred, p);
190 }
191
192 #ifndef NDEBUG
is_black(rbtree_t * tree,uint32_t p)193 static inline bool is_black(rbtree_t *tree, uint32_t p) {
194 return ! is_red(tree, p);
195 }
196 #endif
197
198 /*
199 * Set the color of node p
200 */
mark_red(rbtree_t * tree,uint32_t p)201 static inline void mark_red(rbtree_t *tree, uint32_t p) {
202 assert(p < tree->nbnodes);
203 set_bit(tree->isred, p);
204 }
205
mark_black(rbtree_t * tree,uint32_t p)206 static inline void mark_black(rbtree_t *tree, uint32_t p) {
207 assert(p < tree->nbnodes);
208 clr_bit(tree->isred, p);
209 }
210
211
212 /*
213 * Balance the tree:
214 * - p = new node just added (must be red)
215 * - q = parent of p
216 * - tree->stack contains [null_rbnode, root, ..., r],
217 * which describes a path form the root to r where r = parent of q.
218 * - the root must be black
219 */
rbtree_balance(rbtree_t * tree,uint32_t p,uint32_t q)220 static void rbtree_balance(rbtree_t *tree, uint32_t p, uint32_t q) {
221 uint32_t r, s;
222 uint32_t i, j;
223
224 assert(is_parent_node(tree, q, p) && is_red(tree, p) && is_black(tree, tree->root));
225
226 while (is_red(tree, q)) {
227 r = ivector_pop2(&tree->stack); // r = parent of q
228 assert(is_black(tree, r));
229
230 s = sibling(tree, r, q); // s = sibling of q
231 if (is_red(tree, s)) {
232 // flip colors of q and s
233 mark_black(tree, s);
234 mark_black(tree, q);
235 // if r is the root, we're done
236 if (r == tree->root) break;
237 // otherwise, we color r red
238 // and move up to p := r, q := parent of r
239 mark_red(tree, r);
240 p = r;
241 q = ivector_pop2(&tree->stack); // q = parent of r
242 assert(is_parent_node(tree, q, p));
243
244 } else {
245 // Balance the tree with one or two rotations
246 i = child_index(tree, r, q);
247 j = child_index(tree, q, p);
248 if (i != j) {
249 /*
250 * Rotate p and q
251 * q becomes a child of p
252 * p becomes a child of r
253 */
254 assert(q != 0 && p != 0 && r != 0 &&
255 tree->node[r][i] == q &&
256 tree->node[q][j] == p);
257 tree->node[r][i] = p;
258 tree->node[q][j] = tree->node[p][i];
259 tree->node[p][i] = q;
260
261 // prepare for second rotation:
262 q = p;
263 }
264
265 /*
266 * rotate r and q
267 * and fix the colors: r becomes red, q becomes black
268 */
269 assert(tree->node[r][i] == q);
270 p = ivector_pop2(&tree->stack);
271 if (p == null_rbnode) {
272 assert(r == tree->root);
273 tree->root = q;
274 } else {
275 // p is r's parent
276 j = child_index(tree, p, r);
277 assert(tree->node[p][j] == r);
278 tree->node[p][j] = q;
279 }
280 tree->node[r][i] = tree->node[q][1-i];
281 tree->node[q][1-i] = r;
282 mark_red(tree, r);
283 mark_black(tree, q);
284
285 break;
286 }
287 }
288 }
289
290
291
292 /*
293 * Search or add node of value x
294 * - return the node id
295 * - set new_node to true if that's a new node (x was not present)
296 */
rbtree_get(rbtree_t * tree,uint32_t x,bool * new_node)297 uint32_t rbtree_get(rbtree_t *tree, uint32_t x, bool *new_node) {
298 uint32_t k, i, p;
299
300 assert(tree->stack.size == 0);
301
302 // to force termination: store x in the null_node
303 tree->data[0] = x;
304
305 k = 0; // stop GCC bogus warning
306 p = null_rbnode; // parent
307 i = tree->root;
308 assert(i < tree->nbnodes);
309 while (tree->data[i] != x) {
310 k = (tree->data[i] < x);
311 assert((k == 0 && x < tree->data[i]) ||
312 (k == 1 && x > tree->data[i]));
313
314 // save p on the stack for balancing
315 ivector_push(&tree->stack, p);
316 p = i;
317 i = tree->node[i][k];
318 assert(i < tree->nbnodes);
319 }
320
321 *new_node = false;
322 if (i == 0) {
323 // x is not in the current tree: add it
324 *new_node = true;
325
326 i = rbtree_new_node(tree);
327 tree->data[i] = x;
328 tree->node[i][0] = null_rbnode;
329 tree->node[i][1] = null_rbnode;
330 if (p == null_rbnode) {
331 // make sure the root is always black
332 tree->root = i;
333 mark_black(tree, i);
334 } else {
335 // add i as child of p then balance the tree
336 assert(p < tree->nbnodes && tree->node[p][k] == null_rbnode);
337 tree->node[p][k] = i;
338 mark_red(tree, i);
339 rbtree_balance(tree, i, p);
340 }
341 }
342
343 ivector_reset(&tree->stack);
344
345 assert(i > 0 && tree->data[i] == x);
346
347 return i;
348 }
349
350
351
352 /*
353 * To scan the tree:
354 * - return the id of the node whose value is the smallest
355 * element in the tree that's >= x.
356 * - return null_rbnode (i.e., 0) if all elements are smaller than x.
357 */
rbtree_find_sup(rbtree_t * tree,uint32_t x)358 uint32_t rbtree_find_sup(rbtree_t *tree, uint32_t x) {
359 uint32_t i, j;
360
361 // to force termination: store x in the null_node
362 tree->data[0] = x;
363
364 // j = either null_node or
365 // such that tree->data[j] > x
366 j = null_rbnode;
367 i = tree->root;
368 while (tree->data[i] != x) {
369 if (tree->data[i] < x) {
370 i = tree->node[i][1];
371 } else {
372 j = i;
373 i = tree->node[i][0];
374 }
375 }
376
377 if (i > 0) {
378 assert(tree->data[i] == x);
379 return i;
380 } else {
381 return j;
382 }
383 }
384
385
386
387 #if 0
388
389 /*
390 * Search or add node of value x
391 * - return the node id
392 * - set new_node to true if that's a new node (x was not present)
393 * Don't balance the tree
394 */
395 uint32_t rbtree_get_var(rbtree_t *tree, uint32_t x, bool *new_node) {
396 uint32_t k, i, p;
397
398 // to force termination: store x in the null_node
399 tree->data[0] = x;
400
401 k = 0; // stop GCC bogus warning
402 p = null_rbnode; // parent
403 i = tree->root;
404 assert(i < tree->nbnodes);
405 while (tree->data[i] != x) {
406 k = (tree->data[i] < x);
407 assert((k == 0 && x < tree->data[i]) ||
408 (k == 1 && x > tree->data[i]));
409
410 p = i;
411 i = tree->node[i][k];
412 assert(i < tree->nbnodes);
413 }
414
415 *new_node = false;
416 if (i == 0) {
417 // x is not in the current tree: add it
418 *new_node = true;
419
420 i = rbtree_new_node(tree);
421 tree->data[i] = x;
422 tree->node[i][0] = null_rbnode;
423 tree->node[i][1] = null_rbnode;
424 mark_black(tree, i);
425
426 if (p == null_rbnode) {
427 tree->root = i;
428 } else {
429 assert(p < tree->nbnodes && tree->node[p][k] == null_rbnode);
430 tree->node[p][k] = i;
431 }
432 }
433
434 assert(i > 0 && tree->data[i] == x);
435
436 return i;
437 }
438
439 #endif
440