xref: /openbsd-src/sys/kern/subr_tree.c (revision 7ea23c36e10aeb2dd2178bbcaa02755a2489b93a)
1 /*	$OpenBSD */
2 
3 /*
4  * Copyright 2002 Niels Provos <provos@citi.umich.edu>
5  * All rights reserved.
6  *
7  * Redistribution and use in source and binary forms, with or without
8  * modification, are permitted provided that the following conditions
9  * are met:
10  * 1. Redistributions of source code must retain the above copyright
11  *    notice, this list of conditions and the following disclaimer.
12  * 2. Redistributions in binary form must reproduce the above copyright
13  *    notice, this list of conditions and the following disclaimer in the
14  *    documentation and/or other materials provided with the distribution.
15  *
16  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
17  * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
18  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
19  * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
20  * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
21  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
22  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
23  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
25  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26  */
27 
28 /*
29  * Copyright (c) 2016 David Gwynne <dlg@openbsd.org>
30  *
31  * Permission to use, copy, modify, and distribute this software for any
32  * purpose with or without fee is hereby granted, provided that the above
33  * copyright notice and this permission notice appear in all copies.
34  *
35  * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
36  * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
37  * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
38  * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
39  * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
40  * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
41  * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
42  */
43 
44 #include <sys/param.h>
45 #include <sys/tree.h>
46 
47 static inline void *
48 rb_n2e(const struct rb_type *t, void *node)
49 {
50 	caddr_t addr = (caddr_t)node;
51 
52 	return ((void *)(addr + t->t_offset));
53 }
54 
55 static inline void *
56 rb_e2n(const struct rb_type *t, struct rb_entry *rbe)
57 {
58 	caddr_t addr = (caddr_t)rbe;
59 
60 	return ((void *)(addr - t->t_offset));
61 }
62 
63 #define RBE_LEFT(_rbe)		(_rbe)->rbe_left
64 #define RBE_RIGHT(_rbe)		(_rbe)->rbe_right
65 #define RBE_PARENT(_rbe)	(_rbe)->rbe_parent
66 #define RBE_COLOR(_rbe)		(_rbe)->rbe_color
67 
68 #define RBH_ROOT(_rbt)		(_rbt)->rbt_root
69 
70 static inline void
71 rbe_set(struct rb_entry *rbe, struct rb_entry *parent)
72 {
73 	RBE_PARENT(rbe) = parent;
74 	RBE_LEFT(rbe) = RBE_RIGHT(rbe) = NULL;
75 	RBE_COLOR(rbe) = RB_RED;
76 }
77 
78 static inline void
79 rbe_set_blackred(struct rb_entry *black, struct rb_entry *red)
80 {
81 	RBE_COLOR(black) = RB_BLACK;
82 	RBE_COLOR(red) = RB_RED;
83 }
84 
85 static inline void
86 rbe_augment(const struct rb_type *t, struct rb_entry *rbe)
87 {
88 	(*t->t_augment)(rb_e2n(t, rbe));
89 }
90 
91 static inline void
92 rbe_if_augment(const struct rb_type *t, struct rb_entry *rbe)
93 {
94 	if (t->t_augment != NULL)
95 		rbe_augment(t, rbe);
96 }
97 
98 static inline void
99 rbe_rotate_left(const struct rb_type *t, struct rb_tree *rbt,
100     struct rb_entry *rbe)
101 {
102 	struct rb_entry *parent;
103 	struct rb_entry *tmp;
104 
105 	tmp = RBE_RIGHT(rbe);
106 	RBE_RIGHT(rbe) = RBE_LEFT(tmp);
107 	if (RBE_RIGHT(rbe) != NULL)
108 		RBE_PARENT(RBE_LEFT(tmp)) = rbe;
109 
110 	parent = RBE_PARENT(rbe);
111 	RBE_PARENT(tmp) = parent;
112 	if (parent != NULL) {
113 		if (rbe == RBE_LEFT(parent))
114 			RBE_LEFT(parent) = tmp;
115                 else
116 			RBE_RIGHT(parent) = tmp;
117 	} else
118 		RBH_ROOT(rbt) = tmp;
119 
120 	RBE_LEFT(tmp) = rbe;
121 	RBE_PARENT(rbe) = tmp;
122 
123 	if (t->t_augment != NULL) {
124 		rbe_augment(t, rbe);
125 		rbe_augment(t, tmp);
126 		parent = RBE_PARENT(tmp);
127 		if (parent != NULL)
128 			rbe_augment(t, parent);
129 	}
130 }
131 
132 static inline void
133 rbe_rotate_right(const struct rb_type *t, struct rb_tree *rbt,
134     struct rb_entry *rbe)
135 {
136 	struct rb_entry *parent;
137 	struct rb_entry *tmp;
138 
139 	tmp = RBE_LEFT(rbe);
140 	RBE_LEFT(rbe) = RBE_RIGHT(tmp);
141 	if (RBE_LEFT(rbe) != NULL)
142 		RBE_PARENT(RBE_RIGHT(tmp)) = rbe;
143 
144 	parent = RBE_PARENT(rbe);
145 	RBE_PARENT(tmp) = parent;
146 	if (parent != NULL) {
147 		if (rbe == RBE_LEFT(parent))
148 			RBE_LEFT(parent) = tmp;
149                 else
150 			RBE_RIGHT(parent) = tmp;
151 	} else
152 		RBH_ROOT(rbt) = tmp;
153 
154 	RBE_RIGHT(tmp) = rbe;
155 	RBE_PARENT(rbe) = tmp;
156 
157 	if (t->t_augment != NULL) {
158 		rbe_augment(t, rbe);
159 		rbe_augment(t, tmp);
160 		parent = RBE_PARENT(tmp);
161 		if (parent != NULL)
162 			rbe_augment(t, parent);
163 	}
164 }
165 
166 static inline void
167 rbe_insert_color(const struct rb_type *t, struct rb_tree *rbt,
168     struct rb_entry *rbe)
169 {
170 	struct rb_entry *parent, *gparent, *tmp;
171 
172 	while ((parent = RBE_PARENT(rbe)) != NULL &&
173 	    RBE_COLOR(parent) == RB_RED) {
174 		gparent = RBE_PARENT(parent);
175 
176 		if (parent == RBE_LEFT(gparent)) {
177 			tmp = RBE_RIGHT(gparent);
178 			if (tmp != NULL && RBE_COLOR(tmp) == RB_RED) {
179 				RBE_COLOR(tmp) = RB_BLACK;
180 				rbe_set_blackred(parent, gparent);
181 				rbe = gparent;
182 				continue;
183 			}
184 
185 			if (RBE_RIGHT(parent) == rbe) {
186 				rbe_rotate_left(t, rbt, parent);
187 				tmp = parent;
188 				parent = rbe;
189 				rbe = tmp;
190 			}
191 
192 			rbe_set_blackred(parent, gparent);
193 			rbe_rotate_right(t, rbt, gparent);
194 		} else {
195 			tmp = RBE_LEFT(gparent);
196 			if (tmp != NULL && RBE_COLOR(tmp) == RB_RED) {
197 				RBE_COLOR(tmp) = RB_BLACK;
198 				rbe_set_blackred(parent, gparent);
199 				rbe = gparent;
200 				continue;
201 			}
202 
203 			if (RBE_LEFT(parent) == rbe) {
204 				rbe_rotate_right(t, rbt, parent);
205 				tmp = parent;
206 				parent = rbe;
207 				rbe = tmp;
208 			}
209 
210 			rbe_set_blackred(parent, gparent);
211 			rbe_rotate_left(t, rbt, gparent);
212 		}
213 	}
214 
215 	RBE_COLOR(RBH_ROOT(rbt)) = RB_BLACK;
216 }
217 
218 static inline void
219 rbe_remove_color(const struct rb_type *t, struct rb_tree *rbt,
220     struct rb_entry *parent, struct rb_entry *rbe)
221 {
222 	struct rb_entry *tmp;
223 
224 	while ((rbe == NULL || RBE_COLOR(rbe) == RB_BLACK) &&
225 	    rbe != RBH_ROOT(rbt)) {
226 		if (RBE_LEFT(parent) == rbe) {
227 			tmp = RBE_RIGHT(parent);
228 			if (RBE_COLOR(tmp) == RB_RED) {
229 				rbe_set_blackred(tmp, parent);
230 				rbe_rotate_left(t, rbt, parent);
231 				tmp = RBE_RIGHT(parent);
232 			}
233 			if ((RBE_LEFT(tmp) == NULL ||
234 			     RBE_COLOR(RBE_LEFT(tmp)) == RB_BLACK) &&
235 			    (RBE_RIGHT(tmp) == NULL ||
236 			     RBE_COLOR(RBE_RIGHT(tmp)) == RB_BLACK)) {
237 				RBE_COLOR(tmp) = RB_RED;
238 				rbe = parent;
239 				parent = RBE_PARENT(rbe);
240 			} else {
241 				if (RBE_RIGHT(tmp) == NULL ||
242 				    RBE_COLOR(RBE_RIGHT(tmp)) == RB_BLACK) {
243 					struct rb_entry *oleft;
244 
245 					oleft = RBE_LEFT(tmp);
246 					if (oleft != NULL)
247 						RBE_COLOR(oleft) = RB_BLACK;
248 
249 					RBE_COLOR(tmp) = RB_RED;
250 					rbe_rotate_right(t, rbt, tmp);
251 					tmp = RBE_RIGHT(parent);
252 				}
253 
254 				RBE_COLOR(tmp) = RBE_COLOR(parent);
255 				RBE_COLOR(parent) = RB_BLACK;
256 				if (RBE_RIGHT(tmp))
257 					RBE_COLOR(RBE_RIGHT(tmp)) = RB_BLACK;
258 
259 				rbe_rotate_left(t, rbt, parent);
260 				rbe = RBH_ROOT(rbt);
261 				break;
262 			}
263 		} else {
264 			tmp = RBE_LEFT(parent);
265 			if (RBE_COLOR(tmp) == RB_RED) {
266 				rbe_set_blackred(tmp, parent);
267 				rbe_rotate_right(t, rbt, parent);
268 				tmp = RBE_LEFT(parent);
269 			}
270 
271 			if ((RBE_LEFT(tmp) == NULL ||
272 			     RBE_COLOR(RBE_LEFT(tmp)) == RB_BLACK) &&
273 			    (RBE_RIGHT(tmp) == NULL ||
274 			     RBE_COLOR(RBE_RIGHT(tmp)) == RB_BLACK)) {
275 				RBE_COLOR(tmp) = RB_RED;
276 				rbe = parent;
277 				parent = RBE_PARENT(rbe);
278 			} else {
279 				if (RBE_LEFT(tmp) == NULL ||
280 				    RBE_COLOR(RBE_LEFT(tmp)) == RB_BLACK) {
281 					struct rb_entry *oright;
282 
283 					oright = RBE_RIGHT(tmp);
284 					if (oright != NULL)
285 						RBE_COLOR(oright) = RB_BLACK;
286 
287 					RBE_COLOR(tmp) = RB_RED;
288 					rbe_rotate_left(t, rbt, tmp);
289 					tmp = RBE_LEFT(parent);
290 				}
291 
292 				RBE_COLOR(tmp) = RBE_COLOR(parent);
293 				RBE_COLOR(parent) = RB_BLACK;
294 				if (RBE_LEFT(tmp) != NULL)
295 					RBE_COLOR(RBE_LEFT(tmp)) = RB_BLACK;
296 
297 				rbe_rotate_right(t, rbt, parent);
298 				rbe = RBH_ROOT(rbt);
299 				break;
300 			}
301 		}
302 	}
303 
304 	if (rbe != NULL)
305 		RBE_COLOR(rbe) = RB_BLACK;
306 }
307 
308 static inline struct rb_entry *
309 rbe_remove(const struct rb_type *t, struct rb_tree *rbt, struct rb_entry *rbe)
310 {
311 	struct rb_entry *child, *parent, *old = rbe;
312 	unsigned int color;
313 
314 	if (RBE_LEFT(rbe) == NULL)
315 		child = RBE_RIGHT(rbe);
316 	else if (RBE_RIGHT(rbe) == NULL)
317 		child = RBE_LEFT(rbe);
318 	else {
319 		struct rb_entry *tmp;
320 
321 		rbe = RBE_RIGHT(rbe);
322 		while ((tmp = RBE_LEFT(rbe)) != NULL)
323 			rbe = tmp;
324 
325 		child = RBE_RIGHT(rbe);
326 		parent = RBE_PARENT(rbe);
327 		color = RBE_COLOR(rbe);
328 		if (child != NULL)
329 			RBE_PARENT(child) = parent;
330 		if (parent != NULL) {
331 			if (RBE_LEFT(parent) == rbe)
332 				RBE_LEFT(parent) = child;
333 			else
334 				RBE_RIGHT(parent) = child;
335 
336 			rbe_if_augment(t, parent);
337 		} else
338 			RBH_ROOT(rbt) = child;
339 		if (RBE_PARENT(rbe) == old)
340 			parent = rbe;
341 		*rbe = *old;
342 
343 		tmp = RBE_PARENT(old);
344 		if (tmp != NULL) {
345 			if (RBE_LEFT(tmp) == old)
346 				RBE_LEFT(tmp) = rbe;
347 			else
348 				RBE_RIGHT(tmp) = rbe;
349 
350 			rbe_if_augment(t, parent);
351 		} else
352 			RBH_ROOT(rbt) = rbe;
353 
354 		RBE_PARENT(RBE_LEFT(old)) = rbe;
355 		if (RBE_RIGHT(old))
356 			RBE_PARENT(RBE_RIGHT(old)) = rbe;
357 
358 		if (t->t_augment != NULL && parent != NULL) {
359 			tmp = parent;
360 			do {
361 				rbe_augment(t, tmp);
362 				tmp = RBE_PARENT(tmp);
363 			} while (tmp != NULL);
364 		}
365 
366 		goto color;
367 	}
368 
369 	parent = RBE_PARENT(rbe);
370 	color = RBE_COLOR(rbe);
371 
372 	if (child != NULL)
373 		RBE_PARENT(child) = parent;
374 	if (parent != NULL) {
375 		if (RBE_LEFT(parent) == rbe)
376 			RBE_LEFT(parent) = child;
377 		else
378 			RBE_RIGHT(parent) = child;
379 
380 		rbe_if_augment(t, parent);
381 	} else
382 		RBH_ROOT(rbt) = child;
383 color:
384 	if (color == RB_BLACK)
385 		rbe_remove_color(t, rbt, parent, child);
386 
387 	return (old);
388 }
389 
390 void *
391 _rb_remove(const struct rb_type *t, struct rb_tree *rbt, void *elm)
392 {
393 	struct rb_entry *rbe = rb_n2e(t, elm);
394 	struct rb_entry *old;
395 
396 	old = rbe_remove(t, rbt, rbe);
397 
398 	return (old == NULL ? NULL : rb_e2n(t, old));
399 }
400 
401 void *
402 _rb_insert(const struct rb_type *t, struct rb_tree *rbt, void *elm)
403 {
404 	struct rb_entry *rbe = rb_n2e(t, elm);
405 	struct rb_entry *tmp;
406 	struct rb_entry *parent = NULL;
407 	void *node;
408 	int comp = 0;
409 
410 	tmp = RBH_ROOT(rbt);
411 	while (tmp != NULL) {
412 		parent = tmp;
413 
414 		node = rb_e2n(t, tmp);
415 		comp = (*t->t_compare)(elm, node);
416 		if (comp < 0)
417 			tmp = RBE_LEFT(tmp);
418 		else if (comp > 0)
419 			tmp = RBE_RIGHT(tmp);
420 		else
421 			return (node);
422 	}
423 
424 	rbe_set(rbe, parent);
425 
426 	if (parent != NULL) {
427 		if (comp < 0)
428 			RBE_LEFT(parent) = rbe;
429 		else
430 			RBE_RIGHT(parent) = rbe;
431 
432 		rbe_if_augment(t, parent);
433 	} else
434 		RBH_ROOT(rbt) = rbe;
435 
436 	rbe_insert_color(t, rbt, rbe);
437 
438 	return (NULL);
439 }
440 
441 /* Finds the node with the same key as elm */
442 void *
443 _rb_find(const struct rb_type *t, struct rb_tree *rbt, const void *key)
444 {
445 	struct rb_entry *tmp = RBH_ROOT(rbt);
446 	void *node;
447         int comp;
448 
449 	while (tmp != NULL) {
450 		node = rb_e2n(t, tmp);
451 		comp = (*t->t_compare)(key, node);
452 		if (comp < 0)
453 			tmp = RBE_LEFT(tmp);
454 		else if (comp > 0)
455 			tmp = RBE_RIGHT(tmp);
456 		else
457 			return (node);
458 	}
459 
460 	return (NULL);
461 }
462 
463 /* Finds the first node greater than or equal to the search key */      \
464 void *
465 _rb_nfind(const struct rb_type *t, struct rb_tree *rbt, const void *key)
466 {
467         struct rb_entry *tmp = RBH_ROOT(rbt);
468 	void *node;
469 	void *res = NULL;
470 	int comp;
471 
472         while (tmp != NULL) {
473 		node = rb_e2n(t, tmp);
474                 comp = (*t->t_compare)(key, node);
475 		if (comp < 0) {
476 			res = node;
477 			tmp = RBE_LEFT(tmp);
478 		} else if (comp > 0)
479 			tmp = RBE_RIGHT(tmp);
480 		else
481 			return (node);
482 	}
483 
484 	return (res);
485 }
486 
487 void *
488 _rb_next(const struct rb_type *t, void *elm)
489 {
490 	struct rb_entry *rbe = rb_n2e(t, elm);
491 
492 	if (RBE_RIGHT(rbe) != NULL) {
493 		rbe = RBE_RIGHT(rbe);
494 		while (RBE_LEFT(rbe) != NULL)
495 			rbe = RBE_LEFT(rbe);
496 	} else {
497 		if (RBE_PARENT(rbe) &&
498 		    (rbe == RBE_LEFT(RBE_PARENT(rbe))))
499 			rbe = RBE_PARENT(rbe);
500 		else {
501 			while (RBE_PARENT(rbe) &&
502 			    (rbe == RBE_RIGHT(RBE_PARENT(rbe))))
503 				rbe = RBE_PARENT(rbe);
504 			rbe = RBE_PARENT(rbe);
505 		}
506 	}
507 
508 	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
509 }
510 
511 void *
512 _rb_prev(const struct rb_type *t, void *elm)
513 {
514 	struct rb_entry *rbe = rb_n2e(t, elm);
515 
516 	if (RBE_LEFT(rbe)) {
517 		rbe = RBE_LEFT(rbe);
518 		while (RBE_RIGHT(rbe))
519 			rbe = RBE_RIGHT(rbe);
520 	} else {
521 		if (RBE_PARENT(rbe) &&
522 		    (rbe == RBE_RIGHT(RBE_PARENT(rbe))))
523 			rbe = RBE_PARENT(rbe);
524 		else {
525 			while (RBE_PARENT(rbe) &&
526 			    (rbe == RBE_LEFT(RBE_PARENT(rbe))))
527 				rbe = RBE_PARENT(rbe);
528 			rbe = RBE_PARENT(rbe);
529 		}
530 	}
531 
532 	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
533 }
534 
535 void *
536 _rb_root(const struct rb_type *t, struct rb_tree *rbt)
537 {
538 	struct rb_entry *rbe = RBH_ROOT(rbt);
539 
540 	return (rbe == NULL ? rbe : rb_e2n(t, rbe));
541 }
542 
543 void *
544 _rb_min(const struct rb_type *t, struct rb_tree *rbt)
545 {
546 	struct rb_entry *rbe = RBH_ROOT(rbt);
547 	struct rb_entry *parent = NULL;
548 
549 	while (rbe != NULL) {
550 		parent = rbe;
551 		rbe = RBE_LEFT(rbe);
552 	}
553 
554 	return (parent == NULL ? NULL : rb_e2n(t, parent));
555 }
556 
557 void *
558 _rb_max(const struct rb_type *t, struct rb_tree *rbt)
559 {
560 	struct rb_entry *rbe = RBH_ROOT(rbt);
561 	struct rb_entry *parent = NULL;
562 
563 	while (rbe != NULL) {
564 		parent = rbe;
565 		rbe = RBE_RIGHT(rbe);
566 	}
567 
568 	return (parent == NULL ? NULL : rb_e2n(t, parent));
569 }
570 
571 void *
572 _rb_left(const struct rb_type *t, void *node)
573 {
574 	struct rb_entry *rbe = rb_n2e(t, node);
575 	rbe = RBE_LEFT(rbe);
576 	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
577 }
578 
579 void *
580 _rb_right(const struct rb_type *t, void *node)
581 {
582 	struct rb_entry *rbe = rb_n2e(t, node);
583 	rbe = RBE_RIGHT(rbe);
584 	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
585 }
586 
587 void *
588 _rb_parent(const struct rb_type *t, void *node)
589 {
590 	struct rb_entry *rbe = rb_n2e(t, node);
591 	rbe = RBE_PARENT(rbe);
592 	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
593 }
594 
595