xref: /netbsd-src/common/lib/libc/gen/rb.c (revision 72e44f84cb5b48f30b82cc15884a628bba6d3928)
1 /*	$NetBSD: rb.c,v 1.16 2021/09/16 21:29:41 andvar Exp $	*/
2 
3 /*-
4  * Copyright (c) 2001 The NetBSD Foundation, Inc.
5  * All rights reserved.
6  *
7  * This code is derived from software contributed to The NetBSD Foundation
8  * by Matt Thomas <matt@3am-software.com>.
9  *
10  * Redistribution and use in source and binary forms, with or without
11  * modification, are permitted provided that the following conditions
12  * are met:
13  * 1. Redistributions of source code must retain the above copyright
14  *    notice, this list of conditions and the following disclaimer.
15  * 2. Redistributions in binary form must reproduce the above copyright
16  *    notice, this list of conditions and the following disclaimer in the
17  *    documentation and/or other materials provided with the distribution.
18  *
19  * THIS SOFTWARE IS PROVIDED BY THE NETBSD FOUNDATION, INC. AND CONTRIBUTORS
20  * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
21  * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
22  * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR CONTRIBUTORS
23  * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
24  * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
25  * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
26  * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
27  * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29  * POSSIBILITY OF SUCH DAMAGE.
30  */
31 
32 #if HAVE_NBTOOL_CONFIG_H
33 #include "nbtool_config.h"
34 #endif
35 
36 #if !defined(_KERNEL) && !defined(_STANDALONE)
37 #include <sys/types.h>
38 #include <stddef.h>
39 #include <assert.h>
40 #include <stdbool.h>
41 #ifdef RBDEBUG
42 #define	KASSERT(s)	assert(s)
43 #define	__rbt_unused
44 #else
45 #define KASSERT(s)	do { } while (/*CONSTCOND*/ 0)
46 #define	__rbt_unused	__unused
47 #endif
48 __RCSID("$NetBSD: rb.c,v 1.16 2021/09/16 21:29:41 andvar Exp $");
49 #else
50 #include <lib/libkern/libkern.h>
51 __KERNEL_RCSID(0, "$NetBSD: rb.c,v 1.16 2021/09/16 21:29:41 andvar Exp $");
52 #ifndef DIAGNOSTIC
53 #define	__rbt_unused	__unused
54 #else
55 #define	__rbt_unused
56 #endif
57 #endif
58 
59 #ifdef _LIBC
60 __weak_alias(rb_tree_init, _rb_tree_init)
61 __weak_alias(rb_tree_find_node, _rb_tree_find_node)
62 __weak_alias(rb_tree_find_node_geq, _rb_tree_find_node_geq)
63 __weak_alias(rb_tree_find_node_leq, _rb_tree_find_node_leq)
64 __weak_alias(rb_tree_insert_node, _rb_tree_insert_node)
65 __weak_alias(rb_tree_remove_node, _rb_tree_remove_node)
66 __weak_alias(rb_tree_iterate, _rb_tree_iterate)
67 #ifdef RBDEBUG
68 __weak_alias(rb_tree_check, _rb_tree_check)
69 __weak_alias(rb_tree_depths, _rb_tree_depths)
70 #endif
71 
72 #include "namespace.h"
73 #endif
74 
75 #ifdef RBTEST
76 #include "rbtree.h"
77 #else
78 #include <sys/rbtree.h>
79 #endif
80 
81 static void rb_tree_insert_rebalance(struct rb_tree *, struct rb_node *);
82 static void rb_tree_removal_rebalance(struct rb_tree *, struct rb_node *,
83 	unsigned int);
84 #ifdef RBDEBUG
85 static const struct rb_node *rb_tree_iterate_const(const struct rb_tree *,
86 	const struct rb_node *, const unsigned int);
87 static bool rb_tree_check_node(const struct rb_tree *, const struct rb_node *,
88 	const struct rb_node *, bool);
89 #else
90 #define	rb_tree_check_node(a, b, c, d)	true
91 #endif
92 
93 #define	RB_NODETOITEM(rbto, rbn)	\
94     ((void *)((uintptr_t)(rbn) - (rbto)->rbto_node_offset))
95 #define	RB_ITEMTONODE(rbto, rbn)	\
96     ((rb_node_t *)((uintptr_t)(rbn) + (rbto)->rbto_node_offset))
97 
98 #define	RB_SENTINEL_NODE	NULL
99 
100 void
rb_tree_init(struct rb_tree * rbt,const rb_tree_ops_t * ops)101 rb_tree_init(struct rb_tree *rbt, const rb_tree_ops_t *ops)
102 {
103 
104 	rbt->rbt_ops = ops;
105 	rbt->rbt_root = RB_SENTINEL_NODE;
106 	RB_TAILQ_INIT(&rbt->rbt_nodes);
107 #ifndef RBSMALL
108 	rbt->rbt_minmax[RB_DIR_LEFT] = rbt->rbt_root;	/* minimum node */
109 	rbt->rbt_minmax[RB_DIR_RIGHT] = rbt->rbt_root;	/* maximum node */
110 #endif
111 #ifdef RBSTATS
112 	rbt->rbt_count = 0;
113 	rbt->rbt_insertions = 0;
114 	rbt->rbt_removals = 0;
115 	rbt->rbt_insertion_rebalance_calls = 0;
116 	rbt->rbt_insertion_rebalance_passes = 0;
117 	rbt->rbt_removal_rebalance_calls = 0;
118 	rbt->rbt_removal_rebalance_passes = 0;
119 #endif
120 }
121 
122 void *
rb_tree_find_node(struct rb_tree * rbt,const void * key)123 rb_tree_find_node(struct rb_tree *rbt, const void *key)
124 {
125 	const rb_tree_ops_t *rbto = rbt->rbt_ops;
126 	rbto_compare_key_fn compare_key = rbto->rbto_compare_key;
127 	struct rb_node *parent = rbt->rbt_root;
128 
129 	while (!RB_SENTINEL_P(parent)) {
130 		void *pobj = RB_NODETOITEM(rbto, parent);
131 		const signed int diff = (*compare_key)(rbto->rbto_context,
132 		    pobj, key);
133 		if (diff == 0)
134 			return pobj;
135 		parent = parent->rb_nodes[diff < 0];
136 	}
137 
138 	return NULL;
139 }
140 
141 void *
rb_tree_find_node_geq(struct rb_tree * rbt,const void * key)142 rb_tree_find_node_geq(struct rb_tree *rbt, const void *key)
143 {
144 	const rb_tree_ops_t *rbto = rbt->rbt_ops;
145 	rbto_compare_key_fn compare_key = rbto->rbto_compare_key;
146 	struct rb_node *parent = rbt->rbt_root, *last = NULL;
147 
148 	while (!RB_SENTINEL_P(parent)) {
149 		void *pobj = RB_NODETOITEM(rbto, parent);
150 		const signed int diff = (*compare_key)(rbto->rbto_context,
151 		    pobj, key);
152 		if (diff == 0)
153 			return pobj;
154 		if (diff > 0)
155 			last = parent;
156 		parent = parent->rb_nodes[diff < 0];
157 	}
158 
159 	return last == NULL ? NULL : RB_NODETOITEM(rbto, last);
160 }
161 
162 void *
rb_tree_find_node_leq(struct rb_tree * rbt,const void * key)163 rb_tree_find_node_leq(struct rb_tree *rbt, const void *key)
164 {
165 	const rb_tree_ops_t *rbto = rbt->rbt_ops;
166 	rbto_compare_key_fn compare_key = rbto->rbto_compare_key;
167 	struct rb_node *parent = rbt->rbt_root, *last = NULL;
168 
169 	while (!RB_SENTINEL_P(parent)) {
170 		void *pobj = RB_NODETOITEM(rbto, parent);
171 		const signed int diff = (*compare_key)(rbto->rbto_context,
172 		    pobj, key);
173 		if (diff == 0)
174 			return pobj;
175 		if (diff < 0)
176 			last = parent;
177 		parent = parent->rb_nodes[diff < 0];
178 	}
179 
180 	return last == NULL ? NULL : RB_NODETOITEM(rbto, last);
181 }
182 
183 void *
rb_tree_insert_node(struct rb_tree * rbt,void * object)184 rb_tree_insert_node(struct rb_tree *rbt, void *object)
185 {
186 	const rb_tree_ops_t *rbto = rbt->rbt_ops;
187 	rbto_compare_nodes_fn compare_nodes = rbto->rbto_compare_nodes;
188 	struct rb_node *parent, *tmp, *self = RB_ITEMTONODE(rbto, object);
189 	unsigned int position;
190 	bool rebalance;
191 
192 	RBSTAT_INC(rbt->rbt_insertions);
193 
194 	tmp = rbt->rbt_root;
195 	/*
196 	 * This is a hack.  Because rbt->rbt_root is just a struct rb_node *,
197 	 * just like rb_node->rb_nodes[RB_DIR_LEFT], we can use this fact to
198 	 * avoid a lot of tests for root and know that even at root,
199 	 * updating RB_FATHER(rb_node)->rb_nodes[RB_POSITION(rb_node)] will
200 	 * update rbt->rbt_root.
201 	 */
202 	parent = (struct rb_node *)(void *)&rbt->rbt_root;
203 	position = RB_DIR_LEFT;
204 
205 	/*
206 	 * Find out where to place this new leaf.
207 	 */
208 	while (!RB_SENTINEL_P(tmp)) {
209 		void *tobj = RB_NODETOITEM(rbto, tmp);
210 		const signed int diff = (*compare_nodes)(rbto->rbto_context,
211 		    tobj, object);
212 		if (__predict_false(diff == 0)) {
213 			/*
214 			 * Node already exists; return it.
215 			 */
216 			return tobj;
217 		}
218 		parent = tmp;
219 		position = (diff < 0);
220 		tmp = parent->rb_nodes[position];
221 	}
222 
223 #ifdef RBDEBUG
224 	{
225 		struct rb_node *prev = NULL, *next = NULL;
226 
227 		if (position == RB_DIR_RIGHT)
228 			prev = parent;
229 		else if (tmp != rbt->rbt_root)
230 			next = parent;
231 
232 		/*
233 		 * Verify our sequential position
234 		 */
235 		KASSERT(prev == NULL || !RB_SENTINEL_P(prev));
236 		KASSERT(next == NULL || !RB_SENTINEL_P(next));
237 		if (prev != NULL && next == NULL)
238 			next = TAILQ_NEXT(prev, rb_link);
239 		if (prev == NULL && next != NULL)
240 			prev = TAILQ_PREV(next, rb_node_qh, rb_link);
241 		KASSERT(prev == NULL || !RB_SENTINEL_P(prev));
242 		KASSERT(next == NULL || !RB_SENTINEL_P(next));
243 		KASSERT(prev == NULL || (*compare_nodes)(rbto->rbto_context,
244 		    RB_NODETOITEM(rbto, prev), RB_NODETOITEM(rbto, self)) < 0);
245 		KASSERT(next == NULL || (*compare_nodes)(rbto->rbto_context,
246 		    RB_NODETOITEM(rbto, self), RB_NODETOITEM(rbto, next)) < 0);
247 	}
248 #endif
249 
250 	/*
251 	 * Initialize the node and insert as a leaf into the tree.
252 	 */
253 	RB_SET_FATHER(self, parent);
254 	RB_SET_POSITION(self, position);
255 	if (__predict_false(parent == (struct rb_node *)(void *)&rbt->rbt_root)) {
256 		RB_MARK_BLACK(self);		/* root is always black */
257 #ifndef RBSMALL
258 		rbt->rbt_minmax[RB_DIR_LEFT] = self;
259 		rbt->rbt_minmax[RB_DIR_RIGHT] = self;
260 #endif
261 		rebalance = false;
262 	} else {
263 		KASSERT(position == RB_DIR_LEFT || position == RB_DIR_RIGHT);
264 #ifndef RBSMALL
265 		/*
266 		 * Keep track of the minimum and maximum nodes.  If our
267 		 * parent is a minmax node and we on their min/max side,
268 		 * we must be the new min/max node.
269 		 */
270 		if (parent == rbt->rbt_minmax[position])
271 			rbt->rbt_minmax[position] = self;
272 #endif /* !RBSMALL */
273 		/*
274 		 * All new nodes are colored red.  We only need to rebalance
275 		 * if our parent is also red.
276 		 */
277 		RB_MARK_RED(self);
278 		rebalance = RB_RED_P(parent);
279 	}
280 	KASSERT(RB_SENTINEL_P(parent->rb_nodes[position]));
281 	self->rb_left = parent->rb_nodes[position];
282 	self->rb_right = parent->rb_nodes[position];
283 	parent->rb_nodes[position] = self;
284 	KASSERT(RB_CHILDLESS_P(self));
285 
286 	/*
287 	 * Insert the new node into a sorted list for easy sequential access
288 	 */
289 	RBSTAT_INC(rbt->rbt_count);
290 #ifdef RBDEBUG
291 	if (RB_ROOT_P(rbt, self)) {
292 		RB_TAILQ_INSERT_HEAD(&rbt->rbt_nodes, self, rb_link);
293 	} else if (position == RB_DIR_LEFT) {
294 		KASSERT((*compare_nodes)(rbto->rbto_context,
295 		    RB_NODETOITEM(rbto, self),
296 		    RB_NODETOITEM(rbto, RB_FATHER(self))) < 0);
297 		RB_TAILQ_INSERT_BEFORE(RB_FATHER(self), self, rb_link);
298 	} else {
299 		KASSERT((*compare_nodes)(rbto->rbto_context,
300 		    RB_NODETOITEM(rbto, RB_FATHER(self)),
301 		    RB_NODETOITEM(rbto, self)) < 0);
302 		RB_TAILQ_INSERT_AFTER(&rbt->rbt_nodes, RB_FATHER(self),
303 		    self, rb_link);
304 	}
305 #endif
306 	KASSERT(rb_tree_check_node(rbt, self, NULL, !rebalance));
307 
308 	/*
309 	 * Rebalance tree after insertion
310 	 */
311 	if (rebalance) {
312 		rb_tree_insert_rebalance(rbt, self);
313 		KASSERT(rb_tree_check_node(rbt, self, NULL, true));
314 	}
315 
316 	/* Successfully inserted, return our node pointer. */
317 	return object;
318 }
319 
320 /*
321  * Swap the location and colors of 'self' and its child @ which.  The child
322  * can not be a sentinel node.  This is our rotation function.  However,
323  * since it preserves coloring, it great simplifies both insertion and
324  * removal since rotation almost always involves the exchanging of colors
325  * as a separate step.
326  */
327 static void
rb_tree_reparent_nodes(__rbt_unused struct rb_tree * rbt,struct rb_node * old_father,const unsigned int which)328 rb_tree_reparent_nodes(__rbt_unused struct rb_tree *rbt,
329 	struct rb_node *old_father, const unsigned int which)
330 {
331 	const unsigned int other = which ^ RB_DIR_OTHER;
332 	struct rb_node * const grandpa = RB_FATHER(old_father);
333 	struct rb_node * const old_child = old_father->rb_nodes[which];
334 	struct rb_node * const new_father = old_child;
335 	struct rb_node * const new_child = old_father;
336 
337 	KASSERT(which == RB_DIR_LEFT || which == RB_DIR_RIGHT);
338 
339 	KASSERT(!RB_SENTINEL_P(old_child));
340 	KASSERT(RB_FATHER(old_child) == old_father);
341 
342 	KASSERT(rb_tree_check_node(rbt, old_father, NULL, false));
343 	KASSERT(rb_tree_check_node(rbt, old_child, NULL, false));
344 	KASSERT(RB_ROOT_P(rbt, old_father) ||
345 	    rb_tree_check_node(rbt, grandpa, NULL, false));
346 
347 	/*
348 	 * Exchange descendant linkages.
349 	 */
350 	grandpa->rb_nodes[RB_POSITION(old_father)] = new_father;
351 	new_child->rb_nodes[which] = old_child->rb_nodes[other];
352 	new_father->rb_nodes[other] = new_child;
353 
354 	/*
355 	 * Update ancestor linkages
356 	 */
357 	RB_SET_FATHER(new_father, grandpa);
358 	RB_SET_FATHER(new_child, new_father);
359 
360 	/*
361 	 * Exchange properties between new_father and new_child.  The only
362 	 * change is that new_child's position is now on the other side.
363 	 */
364 #if 0
365 	{
366 		struct rb_node tmp;
367 		tmp.rb_info = 0;
368 		RB_COPY_PROPERTIES(&tmp, old_child);
369 		RB_COPY_PROPERTIES(new_father, old_father);
370 		RB_COPY_PROPERTIES(new_child, &tmp);
371 	}
372 #else
373 	RB_SWAP_PROPERTIES(new_father, new_child);
374 #endif
375 	RB_SET_POSITION(new_child, other);
376 
377 	/*
378 	 * Make sure to reparent the new child to ourself.
379 	 */
380 	if (!RB_SENTINEL_P(new_child->rb_nodes[which])) {
381 		RB_SET_FATHER(new_child->rb_nodes[which], new_child);
382 		RB_SET_POSITION(new_child->rb_nodes[which], which);
383 	}
384 
385 	KASSERT(rb_tree_check_node(rbt, new_father, NULL, false));
386 	KASSERT(rb_tree_check_node(rbt, new_child, NULL, false));
387 	KASSERT(RB_ROOT_P(rbt, new_father) ||
388 	    rb_tree_check_node(rbt, grandpa, NULL, false));
389 }
390 
391 static void
rb_tree_insert_rebalance(struct rb_tree * rbt,struct rb_node * self)392 rb_tree_insert_rebalance(struct rb_tree *rbt, struct rb_node *self)
393 {
394 	struct rb_node * father = RB_FATHER(self);
395 	struct rb_node * grandpa = RB_FATHER(father);
396 	struct rb_node * uncle;
397 	unsigned int which;
398 	unsigned int other;
399 
400 	KASSERT(!RB_ROOT_P(rbt, self));
401 	KASSERT(RB_RED_P(self));
402 	KASSERT(RB_RED_P(father));
403 	RBSTAT_INC(rbt->rbt_insertion_rebalance_calls);
404 
405 	for (;;) {
406 		KASSERT(!RB_SENTINEL_P(self));
407 
408 		KASSERT(RB_RED_P(self));
409 		KASSERT(RB_RED_P(father));
410 		/*
411 		 * We are red and our parent is red, therefore we must have a
412 		 * grandfather and he must be black.
413 		 */
414 		grandpa = RB_FATHER(father);
415 		KASSERT(RB_BLACK_P(grandpa));
416 		KASSERT(RB_DIR_RIGHT == 1 && RB_DIR_LEFT == 0);
417 		which = (father == grandpa->rb_right);
418 		other = which ^ RB_DIR_OTHER;
419 		uncle = grandpa->rb_nodes[other];
420 
421 		if (RB_BLACK_P(uncle))
422 			break;
423 
424 		RBSTAT_INC(rbt->rbt_insertion_rebalance_passes);
425 		/*
426 		 * Case 1: our uncle is red
427 		 *   Simply invert the colors of our parent and
428 		 *   uncle and make our grandparent red.  And
429 		 *   then solve the problem up at his level.
430 		 */
431 		RB_MARK_BLACK(uncle);
432 		RB_MARK_BLACK(father);
433 		if (__predict_false(RB_ROOT_P(rbt, grandpa))) {
434 			/*
435 			 * If our grandpa is root, don't bother
436 			 * setting him to red, just return.
437 			 */
438 			KASSERT(RB_BLACK_P(grandpa));
439 			return;
440 		}
441 		RB_MARK_RED(grandpa);
442 		self = grandpa;
443 		father = RB_FATHER(self);
444 		KASSERT(RB_RED_P(self));
445 		if (RB_BLACK_P(father)) {
446 			/*
447 			 * If our greatgrandpa is black, we're done.
448 			 */
449 			KASSERT(RB_BLACK_P(rbt->rbt_root));
450 			return;
451 		}
452 	}
453 
454 	KASSERT(!RB_ROOT_P(rbt, self));
455 	KASSERT(RB_RED_P(self));
456 	KASSERT(RB_RED_P(father));
457 	KASSERT(RB_BLACK_P(uncle));
458 	KASSERT(RB_BLACK_P(grandpa));
459 	/*
460 	 * Case 2&3: our uncle is black.
461 	 */
462 	if (self == father->rb_nodes[other]) {
463 		/*
464 		 * Case 2: we are on the same side as our uncle
465 		 *   Swap ourselves with our parent so this case
466 		 *   becomes case 3.  Basically our parent becomes our
467 		 *   child.
468 		 */
469 		rb_tree_reparent_nodes(rbt, father, other);
470 		KASSERT(RB_FATHER(father) == self);
471 		KASSERT(self->rb_nodes[which] == father);
472 		KASSERT(RB_FATHER(self) == grandpa);
473 		self = father;
474 		father = RB_FATHER(self);
475 	}
476 	KASSERT(RB_RED_P(self) && RB_RED_P(father));
477 	KASSERT(grandpa->rb_nodes[which] == father);
478 	/*
479 	 * Case 3: we are opposite a child of a black uncle.
480 	 *   Swap our parent and grandparent.  Since our grandfather
481 	 *   is black, our father will become black and our new sibling
482 	 *   (former grandparent) will become red.
483 	 */
484 	rb_tree_reparent_nodes(rbt, grandpa, which);
485 	KASSERT(RB_FATHER(self) == father);
486 	KASSERT(RB_FATHER(self)->rb_nodes[RB_POSITION(self) ^ RB_DIR_OTHER] == grandpa);
487 	KASSERT(RB_RED_P(self));
488 	KASSERT(RB_BLACK_P(father));
489 	KASSERT(RB_RED_P(grandpa));
490 
491 	/*
492 	 * Final step: Set the root to black.
493 	 */
494 	RB_MARK_BLACK(rbt->rbt_root);
495 }
496 
497 static void
rb_tree_prune_node(struct rb_tree * rbt,struct rb_node * self,bool rebalance)498 rb_tree_prune_node(struct rb_tree *rbt, struct rb_node *self, bool rebalance)
499 {
500 	const unsigned int which = RB_POSITION(self);
501 	struct rb_node *father = RB_FATHER(self);
502 #ifndef RBSMALL
503 	const bool was_root = RB_ROOT_P(rbt, self);
504 #endif
505 
506 	KASSERT(rebalance || (RB_ROOT_P(rbt, self) || RB_RED_P(self)));
507 	KASSERT(!rebalance || RB_BLACK_P(self));
508 	KASSERT(RB_CHILDLESS_P(self));
509 	KASSERT(rb_tree_check_node(rbt, self, NULL, false));
510 
511 	/*
512 	 * Since we are childless, we know that self->rb_left is pointing
513 	 * to the sentinel node.
514 	 */
515 	father->rb_nodes[which] = self->rb_left;
516 
517 	/*
518 	 * Remove ourselves from the node list, decrement the count,
519 	 * and update min/max.
520 	 */
521 	RB_TAILQ_REMOVE(&rbt->rbt_nodes, self, rb_link);
522 	RBSTAT_DEC(rbt->rbt_count);
523 #ifndef RBSMALL
524 	if (__predict_false(rbt->rbt_minmax[RB_POSITION(self)] == self)) {
525 		rbt->rbt_minmax[RB_POSITION(self)] = father;
526 		/*
527 		 * When removing the root, rbt->rbt_minmax[RB_DIR_LEFT] is
528 		 * updated automatically, but we also need to update
529 		 * rbt->rbt_minmax[RB_DIR_RIGHT];
530 		 */
531 		if (__predict_false(was_root)) {
532 			rbt->rbt_minmax[RB_DIR_RIGHT] = father;
533 		}
534 	}
535 	RB_SET_FATHER(self, NULL);
536 #endif
537 
538 	/*
539 	 * Rebalance if requested.
540 	 */
541 	if (rebalance)
542 		rb_tree_removal_rebalance(rbt, father, which);
543 	KASSERT(was_root || rb_tree_check_node(rbt, father, NULL, true));
544 }
545 
546 /*
547  * When deleting an interior node
548  */
549 static void
rb_tree_swap_prune_and_rebalance(struct rb_tree * rbt,struct rb_node * self,struct rb_node * standin)550 rb_tree_swap_prune_and_rebalance(struct rb_tree *rbt, struct rb_node *self,
551 	struct rb_node *standin)
552 {
553 	const unsigned int standin_which = RB_POSITION(standin);
554 	unsigned int standin_other = standin_which ^ RB_DIR_OTHER;
555 	struct rb_node *standin_son;
556 	struct rb_node *standin_father = RB_FATHER(standin);
557 	bool rebalance = RB_BLACK_P(standin);
558 
559 	if (standin_father == self) {
560 		/*
561 		 * As a child of self, any childen would be opposite of
562 		 * our parent.
563 		 */
564 		KASSERT(RB_SENTINEL_P(standin->rb_nodes[standin_other]));
565 		standin_son = standin->rb_nodes[standin_which];
566 	} else {
567 		/*
568 		 * Since we aren't a child of self, any childen would be
569 		 * on the same side as our parent.
570 		 */
571 		KASSERT(RB_SENTINEL_P(standin->rb_nodes[standin_which]));
572 		standin_son = standin->rb_nodes[standin_other];
573 	}
574 
575 	/*
576 	 * the node we are removing must have two children.
577 	 */
578 	KASSERT(RB_TWOCHILDREN_P(self));
579 	/*
580 	 * If standin has a child, it must be red.
581 	 */
582 	KASSERT(RB_SENTINEL_P(standin_son) || RB_RED_P(standin_son));
583 
584 	/*
585 	 * Verify things are sane.
586 	 */
587 	KASSERT(rb_tree_check_node(rbt, self, NULL, false));
588 	KASSERT(rb_tree_check_node(rbt, standin, NULL, false));
589 
590 	if (__predict_false(RB_RED_P(standin_son))) {
591 		/*
592 		 * We know we have a red child so if we flip it to black
593 		 * we don't have to rebalance.
594 		 */
595 		KASSERT(rb_tree_check_node(rbt, standin_son, NULL, true));
596 		RB_MARK_BLACK(standin_son);
597 		rebalance = false;
598 
599 		if (standin_father == self) {
600 			KASSERT(RB_POSITION(standin_son) == standin_which);
601 		} else {
602 			KASSERT(RB_POSITION(standin_son) == standin_other);
603 			/*
604 			 * Change the son's parentage to point to his grandpa.
605 			 */
606 			RB_SET_FATHER(standin_son, standin_father);
607 			RB_SET_POSITION(standin_son, standin_which);
608 		}
609 	}
610 
611 	if (standin_father == self) {
612 		/*
613 		 * If we are about to delete the standin's father, then when
614 		 * we call rebalance, we need to use ourselves as our father.
615 		 * Otherwise remember our original father.  Also, sincef we are
616 		 * our standin's father we only need to reparent the standin's
617 		 * brother.
618 		 *
619 		 * |    R      -->     S    |
620 		 * |  Q   S    -->   Q   T  |
621 		 * |        t  -->          |
622 		 */
623 		KASSERT(RB_SENTINEL_P(standin->rb_nodes[standin_other]));
624 		KASSERT(!RB_SENTINEL_P(self->rb_nodes[standin_other]));
625 		KASSERT(self->rb_nodes[standin_which] == standin);
626 		/*
627 		 * Have our son/standin adopt his brother as his new son.
628 		 */
629 		standin_father = standin;
630 	} else {
631 		/*
632 		 * |    R          -->    S       .  |
633 		 * |   / \  |   T  -->   / \  |  /   |
634 		 * |  ..... | S    -->  ..... | T    |
635 		 *
636 		 * Sever standin's connection to his father.
637 		 */
638 		standin_father->rb_nodes[standin_which] = standin_son;
639 		/*
640 		 * Adopt the far son.
641 		 */
642 		standin->rb_nodes[standin_other] = self->rb_nodes[standin_other];
643 		RB_SET_FATHER(standin->rb_nodes[standin_other], standin);
644 		KASSERT(RB_POSITION(self->rb_nodes[standin_other]) == standin_other);
645 		/*
646 		 * Use standin_other because we need to preserve standin_which
647 		 * for the removal_rebalance.
648 		 */
649 		standin_other = standin_which;
650 	}
651 
652 	/*
653 	 * Move the only remaining son to our standin.  If our standin is our
654 	 * son, this will be the only son needed to be moved.
655 	 */
656 	KASSERT(standin->rb_nodes[standin_other] != self->rb_nodes[standin_other]);
657 	standin->rb_nodes[standin_other] = self->rb_nodes[standin_other];
658 	RB_SET_FATHER(standin->rb_nodes[standin_other], standin);
659 
660 	/*
661 	 * Now copy the result of self to standin and then replace
662 	 * self with standin in the tree.
663 	 */
664 	RB_COPY_PROPERTIES(standin, self);
665 	RB_SET_FATHER(standin, RB_FATHER(self));
666 	RB_FATHER(standin)->rb_nodes[RB_POSITION(standin)] = standin;
667 
668 	/*
669 	 * Remove ourselves from the node list, decrement the count,
670 	 * and update min/max.
671 	 */
672 	RB_TAILQ_REMOVE(&rbt->rbt_nodes, self, rb_link);
673 	RBSTAT_DEC(rbt->rbt_count);
674 #ifndef RBSMALL
675 	if (__predict_false(rbt->rbt_minmax[RB_POSITION(self)] == self))
676 		rbt->rbt_minmax[RB_POSITION(self)] = RB_FATHER(self);
677 	RB_SET_FATHER(self, NULL);
678 #endif
679 
680 	KASSERT(rb_tree_check_node(rbt, standin, NULL, false));
681 	KASSERT(RB_FATHER_SENTINEL_P(standin)
682 		|| rb_tree_check_node(rbt, standin_father, NULL, false));
683 	KASSERT(RB_LEFT_SENTINEL_P(standin)
684 		|| rb_tree_check_node(rbt, standin->rb_left, NULL, false));
685 	KASSERT(RB_RIGHT_SENTINEL_P(standin)
686 		|| rb_tree_check_node(rbt, standin->rb_right, NULL, false));
687 
688 	if (!rebalance)
689 		return;
690 
691 	rb_tree_removal_rebalance(rbt, standin_father, standin_which);
692 	KASSERT(rb_tree_check_node(rbt, standin, NULL, true));
693 }
694 
695 /*
696  * We could do this by doing
697  *	rb_tree_node_swap(rbt, self, which);
698  *	rb_tree_prune_node(rbt, self, false);
699  *
700  * But it's more efficient to just evalate and recolor the child.
701  */
702 static void
rb_tree_prune_blackred_branch(struct rb_tree * rbt,struct rb_node * self,unsigned int which)703 rb_tree_prune_blackred_branch(struct rb_tree *rbt, struct rb_node *self,
704 	unsigned int which)
705 {
706 	struct rb_node *father = RB_FATHER(self);
707 	struct rb_node *son = self->rb_nodes[which];
708 #ifndef RBSMALL
709 	const bool was_root = RB_ROOT_P(rbt, self);
710 #endif
711 
712 	KASSERT(which == RB_DIR_LEFT || which == RB_DIR_RIGHT);
713 	KASSERT(RB_BLACK_P(self) && RB_RED_P(son));
714 	KASSERT(!RB_TWOCHILDREN_P(son));
715 	KASSERT(RB_CHILDLESS_P(son));
716 	KASSERT(rb_tree_check_node(rbt, self, NULL, false));
717 	KASSERT(rb_tree_check_node(rbt, son, NULL, false));
718 
719 	/*
720 	 * Remove ourselves from the tree and give our former child our
721 	 * properties (position, color, root).
722 	 */
723 	RB_COPY_PROPERTIES(son, self);
724 	father->rb_nodes[RB_POSITION(son)] = son;
725 	RB_SET_FATHER(son, father);
726 
727 	/*
728 	 * Remove ourselves from the node list, decrement the count,
729 	 * and update minmax.
730 	 */
731 	RB_TAILQ_REMOVE(&rbt->rbt_nodes, self, rb_link);
732 	RBSTAT_DEC(rbt->rbt_count);
733 #ifndef RBSMALL
734 	if (__predict_false(was_root)) {
735 		KASSERT(rbt->rbt_minmax[which] == son);
736 		rbt->rbt_minmax[which ^ RB_DIR_OTHER] = son;
737 	} else if (rbt->rbt_minmax[RB_POSITION(self)] == self) {
738 		rbt->rbt_minmax[RB_POSITION(self)] = son;
739 	}
740 	RB_SET_FATHER(self, NULL);
741 #endif
742 
743 	KASSERT(was_root || rb_tree_check_node(rbt, father, NULL, true));
744 	KASSERT(rb_tree_check_node(rbt, son, NULL, true));
745 }
746 
747 void
rb_tree_remove_node(struct rb_tree * rbt,void * object)748 rb_tree_remove_node(struct rb_tree *rbt, void *object)
749 {
750 	const rb_tree_ops_t *rbto = rbt->rbt_ops;
751 	struct rb_node *standin, *self = RB_ITEMTONODE(rbto, object);
752 	unsigned int which;
753 
754 	KASSERT(!RB_SENTINEL_P(self));
755 	RBSTAT_INC(rbt->rbt_removals);
756 
757 	/*
758 	 * In the following diagrams, we (the node to be removed) are S.  Red
759 	 * nodes are lowercase.  T could be either red or black.
760 	 *
761 	 * Remember the major axiom of the red-black tree: the number of
762 	 * black nodes from the root to each leaf is constant across all
763 	 * leaves, only the number of red nodes varies.
764 	 *
765 	 * Thus removing a red leaf doesn't require any other changes to a
766 	 * red-black tree.  So if we must remove a node, attempt to rearrange
767 	 * the tree so we can remove a red node.
768 	 *
769 	 * The simpliest case is a childless red node or a childless root node:
770 	 *
771 	 * |    T  -->    T  |    or    |  R  -->  *  |
772 	 * |  s    -->  *    |
773 	 */
774 	if (RB_CHILDLESS_P(self)) {
775 		const bool rebalance = RB_BLACK_P(self) && !RB_ROOT_P(rbt, self);
776 		rb_tree_prune_node(rbt, self, rebalance);
777 		return;
778 	}
779 	KASSERT(!RB_CHILDLESS_P(self));
780 	if (!RB_TWOCHILDREN_P(self)) {
781 		/*
782 		 * The next simpliest case is the node we are deleting is
783 		 * black and has one red child.
784 		 *
785 		 * |      T  -->      T  -->      T  |
786 		 * |    S    -->  R      -->  R      |
787 		 * |  r      -->    s    -->    *    |
788 		 */
789 		which = RB_LEFT_SENTINEL_P(self) ? RB_DIR_RIGHT : RB_DIR_LEFT;
790 		KASSERT(RB_BLACK_P(self));
791 		KASSERT(RB_RED_P(self->rb_nodes[which]));
792 		KASSERT(RB_CHILDLESS_P(self->rb_nodes[which]));
793 		rb_tree_prune_blackred_branch(rbt, self, which);
794 		return;
795 	}
796 	KASSERT(RB_TWOCHILDREN_P(self));
797 
798 	/*
799 	 * We invert these because we prefer to remove from the inside of
800 	 * the tree.
801 	 */
802 	which = RB_POSITION(self) ^ RB_DIR_OTHER;
803 
804 	/*
805 	 * Let's find the node closes to us opposite of our parent
806 	 * Now swap it with ourself, "prune" it, and rebalance, if needed.
807 	 */
808 	standin = RB_ITEMTONODE(rbto, rb_tree_iterate(rbt, object, which));
809 	rb_tree_swap_prune_and_rebalance(rbt, self, standin);
810 }
811 
812 static void
rb_tree_removal_rebalance(struct rb_tree * rbt,struct rb_node * parent,unsigned int which)813 rb_tree_removal_rebalance(struct rb_tree *rbt, struct rb_node *parent,
814 	unsigned int which)
815 {
816 	KASSERT(!RB_SENTINEL_P(parent));
817 	KASSERT(RB_SENTINEL_P(parent->rb_nodes[which]));
818 	KASSERT(which == RB_DIR_LEFT || which == RB_DIR_RIGHT);
819 	RBSTAT_INC(rbt->rbt_removal_rebalance_calls);
820 
821 	while (RB_BLACK_P(parent->rb_nodes[which])) {
822 		unsigned int other = which ^ RB_DIR_OTHER;
823 		struct rb_node *brother = parent->rb_nodes[other];
824 
825 		RBSTAT_INC(rbt->rbt_removal_rebalance_passes);
826 
827 		KASSERT(!RB_SENTINEL_P(brother));
828 		/*
829 		 * For cases 1, 2a, and 2b, our brother's children must
830 		 * be black and our father must be black
831 		 */
832 		if (RB_BLACK_P(parent)
833 		    && RB_BLACK_P(brother->rb_left)
834 		    && RB_BLACK_P(brother->rb_right)) {
835 			if (RB_RED_P(brother)) {
836 				/*
837 				 * Case 1: Our brother is red, swap its
838 				 * position (and colors) with our parent.
839 				 * This should now be case 2b (unless C or E
840 				 * has a red child which is case 3; thus no
841 				 * explicit branch to case 2b).
842 				 *
843 				 *    B         ->        D
844 				 *  A     d     ->    b     E
845 				 *      C   E   ->  A   C
846 				 */
847 				KASSERT(RB_BLACK_P(parent));
848 				rb_tree_reparent_nodes(rbt, parent, other);
849 				brother = parent->rb_nodes[other];
850 				KASSERT(!RB_SENTINEL_P(brother));
851 				KASSERT(RB_RED_P(parent));
852 				KASSERT(RB_BLACK_P(brother));
853 				KASSERT(rb_tree_check_node(rbt, brother, NULL, false));
854 				KASSERT(rb_tree_check_node(rbt, parent, NULL, false));
855 			} else {
856 				/*
857 				 * Both our parent and brother are black.
858 				 * Change our brother to red, advance up rank
859 				 * and go through the loop again.
860 				 *
861 				 *    B         ->   *B
862 				 * *A     D     ->  A     d
863 				 *      C   E   ->      C   E
864 				 */
865 				RB_MARK_RED(brother);
866 				KASSERT(RB_BLACK_P(brother->rb_left));
867 				KASSERT(RB_BLACK_P(brother->rb_right));
868 				if (RB_ROOT_P(rbt, parent))
869 					return;	/* root == parent == black */
870 				KASSERT(rb_tree_check_node(rbt, brother, NULL, false));
871 				KASSERT(rb_tree_check_node(rbt, parent, NULL, false));
872 				which = RB_POSITION(parent);
873 				parent = RB_FATHER(parent);
874 				continue;
875 			}
876 		}
877 		/*
878 		 * Avoid an else here so that case 2a above can hit either
879 		 * case 2b, 3, or 4.
880 		 */
881 		if (RB_RED_P(parent)
882 		    && RB_BLACK_P(brother)
883 		    && RB_BLACK_P(brother->rb_left)
884 		    && RB_BLACK_P(brother->rb_right)) {
885 			KASSERT(RB_RED_P(parent));
886 			KASSERT(RB_BLACK_P(brother));
887 			KASSERT(RB_BLACK_P(brother->rb_left));
888 			KASSERT(RB_BLACK_P(brother->rb_right));
889 			/*
890 			 * We are black, our father is red, our brother and
891 			 * both nephews are black.  Simply invert/exchange the
892 			 * colors of our father and brother (to black and red
893 			 * respectively).
894 			 *
895 			 *	|    f        -->    F        |
896 			 *	|  *     B    -->  *     b    |
897 			 *	|      N   N  -->      N   N  |
898 			 */
899 			RB_MARK_BLACK(parent);
900 			RB_MARK_RED(brother);
901 			KASSERT(rb_tree_check_node(rbt, brother, NULL, true));
902 			break;		/* We're done! */
903 		} else {
904 			/*
905 			 * Our brother must be black and have at least one
906 			 * red child (it may have two).
907 			 */
908 			KASSERT(RB_BLACK_P(brother));
909 			KASSERT(RB_RED_P(brother->rb_nodes[which]) ||
910 				RB_RED_P(brother->rb_nodes[other]));
911 			if (RB_BLACK_P(brother->rb_nodes[other])) {
912 				/*
913 				 * Case 3: our brother is black, our near
914 				 * nephew is red, and our far nephew is black.
915 				 * Swap our brother with our near nephew.
916 				 * This result in a tree that matches case 4.
917 				 * (Our father could be red or black).
918 				 *
919 				 *	|    F      -->    F      |
920 				 *	|  x     B  -->  x   B    |
921 				 *	|      n    -->        n  |
922 				 */
923 				KASSERT(RB_RED_P(brother->rb_nodes[which]));
924 				rb_tree_reparent_nodes(rbt, brother, which);
925 				KASSERT(RB_FATHER(brother) == parent->rb_nodes[other]);
926 				brother = parent->rb_nodes[other];
927 				KASSERT(RB_RED_P(brother->rb_nodes[other]));
928 			}
929 			/*
930 			 * Case 4: our brother is black and our far nephew
931 			 * is red.  Swap our father and brother locations and
932 			 * change our far nephew to black.  (these can be
933 			 * done in either order so we change the color first).
934 			 * The result is a valid red-black tree and is a
935 			 * terminal case.  (again we don't care about the
936 			 * father's color)
937 			 *
938 			 * If the father is red, we will get a red-black-black
939 			 * tree:
940 			 *	|  f      ->  f      -->    b    |
941 			 *	|    B    ->    B    -->  F   N  |
942 			 *	|      n  ->      N  -->         |
943 			 *
944 			 * If the father is black, we will get an all black
945 			 * tree:
946 			 *	|  F      ->  F      -->    B    |
947 			 *	|    B    ->    B    -->  F   N  |
948 			 *	|      n  ->      N  -->         |
949 			 *
950 			 * If we had two red nephews, then after the swap,
951 			 * our former father would have a red grandson.
952 			 */
953 			KASSERT(RB_BLACK_P(brother));
954 			KASSERT(RB_RED_P(brother->rb_nodes[other]));
955 			RB_MARK_BLACK(brother->rb_nodes[other]);
956 			rb_tree_reparent_nodes(rbt, parent, other);
957 			break;		/* We're done! */
958 		}
959 	}
960 	KASSERT(rb_tree_check_node(rbt, parent, NULL, true));
961 }
962 
963 void *
rb_tree_iterate(struct rb_tree * rbt,void * object,const unsigned int direction)964 rb_tree_iterate(struct rb_tree *rbt, void *object, const unsigned int direction)
965 {
966 	const rb_tree_ops_t *rbto = rbt->rbt_ops;
967 	const unsigned int other = direction ^ RB_DIR_OTHER;
968 	struct rb_node *self;
969 
970 	KASSERT(direction == RB_DIR_LEFT || direction == RB_DIR_RIGHT);
971 
972 	if (object == NULL) {
973 #ifndef RBSMALL
974 		if (RB_SENTINEL_P(rbt->rbt_root))
975 			return NULL;
976 		return RB_NODETOITEM(rbto, rbt->rbt_minmax[direction]);
977 #else
978 		self = rbt->rbt_root;
979 		if (RB_SENTINEL_P(self))
980 			return NULL;
981 		while (!RB_SENTINEL_P(self->rb_nodes[direction]))
982 			self = self->rb_nodes[direction];
983 		return RB_NODETOITEM(rbto, self);
984 #endif /* !RBSMALL */
985 	}
986 	self = RB_ITEMTONODE(rbto, object);
987 	KASSERT(!RB_SENTINEL_P(self));
988 	/*
989 	 * We can't go any further in this direction.  We proceed up in the
990 	 * opposite direction until our parent is in direction we want to go.
991 	 */
992 	if (RB_SENTINEL_P(self->rb_nodes[direction])) {
993 		while (!RB_ROOT_P(rbt, self)) {
994 			if (other == RB_POSITION(self))
995 				return RB_NODETOITEM(rbto, RB_FATHER(self));
996 			self = RB_FATHER(self);
997 		}
998 		return NULL;
999 	}
1000 
1001 	/*
1002 	 * Advance down one in current direction and go down as far as possible
1003 	 * in the opposite direction.
1004 	 */
1005 	self = self->rb_nodes[direction];
1006 	KASSERT(!RB_SENTINEL_P(self));
1007 	while (!RB_SENTINEL_P(self->rb_nodes[other]))
1008 		self = self->rb_nodes[other];
1009 	return RB_NODETOITEM(rbto, self);
1010 }
1011 
1012 #ifdef RBDEBUG
1013 static const struct rb_node *
rb_tree_iterate_const(const struct rb_tree * rbt,const struct rb_node * self,const unsigned int direction)1014 rb_tree_iterate_const(const struct rb_tree *rbt, const struct rb_node *self,
1015 	const unsigned int direction)
1016 {
1017 	const unsigned int other = direction ^ RB_DIR_OTHER;
1018 	KASSERT(direction == RB_DIR_LEFT || direction == RB_DIR_RIGHT);
1019 
1020 	if (self == NULL) {
1021 #ifndef RBSMALL
1022 		if (RB_SENTINEL_P(rbt->rbt_root))
1023 			return NULL;
1024 		return rbt->rbt_minmax[direction];
1025 #else
1026 		self = rbt->rbt_root;
1027 		if (RB_SENTINEL_P(self))
1028 			return NULL;
1029 		while (!RB_SENTINEL_P(self->rb_nodes[direction]))
1030 			self = self->rb_nodes[direction];
1031 		return self;
1032 #endif /* !RBSMALL */
1033 	}
1034 	KASSERT(!RB_SENTINEL_P(self));
1035 	/*
1036 	 * We can't go any further in this direction.  We proceed up in the
1037 	 * opposite direction until our parent is in direction we want to go.
1038 	 */
1039 	if (RB_SENTINEL_P(self->rb_nodes[direction])) {
1040 		while (!RB_ROOT_P(rbt, self)) {
1041 			if (other == RB_POSITION(self))
1042 				return RB_FATHER(self);
1043 			self = RB_FATHER(self);
1044 		}
1045 		return NULL;
1046 	}
1047 
1048 	/*
1049 	 * Advance down one in current direction and go down as far as possible
1050 	 * in the opposite direction.
1051 	 */
1052 	self = self->rb_nodes[direction];
1053 	KASSERT(!RB_SENTINEL_P(self));
1054 	while (!RB_SENTINEL_P(self->rb_nodes[other]))
1055 		self = self->rb_nodes[other];
1056 	return self;
1057 }
1058 
1059 static unsigned int
rb_tree_count_black(const struct rb_node * self)1060 rb_tree_count_black(const struct rb_node *self)
1061 {
1062 	unsigned int left, right;
1063 
1064 	if (RB_SENTINEL_P(self))
1065 		return 0;
1066 
1067 	left = rb_tree_count_black(self->rb_left);
1068 	right = rb_tree_count_black(self->rb_right);
1069 
1070 	KASSERT(left == right);
1071 
1072 	return left + RB_BLACK_P(self);
1073 }
1074 
1075 static bool
rb_tree_check_node(const struct rb_tree * rbt,const struct rb_node * self,const struct rb_node * prev,bool red_check)1076 rb_tree_check_node(const struct rb_tree *rbt, const struct rb_node *self,
1077 	const struct rb_node *prev, bool red_check)
1078 {
1079 	const rb_tree_ops_t *rbto = rbt->rbt_ops;
1080 	rbto_compare_nodes_fn compare_nodes = rbto->rbto_compare_nodes;
1081 
1082 	KASSERT(!RB_SENTINEL_P(self));
1083 	KASSERT(prev == NULL || (*compare_nodes)(rbto->rbto_context,
1084 	    RB_NODETOITEM(rbto, prev), RB_NODETOITEM(rbto, self)) < 0);
1085 
1086 	/*
1087 	 * Verify our relationship to our parent.
1088 	 */
1089 	if (RB_ROOT_P(rbt, self)) {
1090 		KASSERT(self == rbt->rbt_root);
1091 		KASSERT(RB_POSITION(self) == RB_DIR_LEFT);
1092 		KASSERT(RB_FATHER(self)->rb_nodes[RB_DIR_LEFT] == self);
1093 		KASSERT(RB_FATHER(self) == (const struct rb_node *) &rbt->rbt_root);
1094 	} else {
1095 		int diff = (*compare_nodes)(rbto->rbto_context,
1096 		    RB_NODETOITEM(rbto, self),
1097 		    RB_NODETOITEM(rbto, RB_FATHER(self)));
1098 
1099 		KASSERT(self != rbt->rbt_root);
1100 		KASSERT(!RB_FATHER_SENTINEL_P(self));
1101 		if (RB_POSITION(self) == RB_DIR_LEFT) {
1102 			KASSERT(diff < 0);
1103 			KASSERT(RB_FATHER(self)->rb_nodes[RB_DIR_LEFT] == self);
1104 		} else {
1105 			KASSERT(diff > 0);
1106 			KASSERT(RB_FATHER(self)->rb_nodes[RB_DIR_RIGHT] == self);
1107 		}
1108 	}
1109 
1110 	/*
1111 	 * Verify our position in the linked list against the tree itself.
1112 	 */
1113 	{
1114 		const struct rb_node *prev0 = rb_tree_iterate_const(rbt, self, RB_DIR_LEFT);
1115 		const struct rb_node *next0 = rb_tree_iterate_const(rbt, self, RB_DIR_RIGHT);
1116 		KASSERT(prev0 == TAILQ_PREV(self, rb_node_qh, rb_link));
1117 		KASSERT(next0 == TAILQ_NEXT(self, rb_link));
1118 #ifndef RBSMALL
1119 		KASSERT(prev0 != NULL || self == rbt->rbt_minmax[RB_DIR_LEFT]);
1120 		KASSERT(next0 != NULL || self == rbt->rbt_minmax[RB_DIR_RIGHT]);
1121 #endif
1122 	}
1123 
1124 	/*
1125 	 * The root must be black.
1126 	 * There can never be two adjacent red nodes.
1127 	 */
1128 	if (red_check) {
1129 		KASSERT(!RB_ROOT_P(rbt, self) || RB_BLACK_P(self));
1130 		(void) rb_tree_count_black(self);
1131 		if (RB_RED_P(self)) {
1132 			const struct rb_node *brother;
1133 			KASSERT(!RB_ROOT_P(rbt, self));
1134 			brother = RB_FATHER(self)->rb_nodes[RB_POSITION(self) ^ RB_DIR_OTHER];
1135 			KASSERT(RB_BLACK_P(RB_FATHER(self)));
1136 			/*
1137 			 * I'm red and have no children, then I must either
1138 			 * have no brother or my brother also be red and
1139 			 * also have no children.  (black count == 0)
1140 			 */
1141 			KASSERT(!RB_CHILDLESS_P(self)
1142 				|| RB_SENTINEL_P(brother)
1143 				|| RB_RED_P(brother)
1144 				|| RB_CHILDLESS_P(brother));
1145 			/*
1146 			 * If I'm not childless, I must have two children
1147 			 * and they must be both be black.
1148 			 */
1149 			KASSERT(RB_CHILDLESS_P(self)
1150 				|| (RB_TWOCHILDREN_P(self)
1151 				    && RB_BLACK_P(self->rb_left)
1152 				    && RB_BLACK_P(self->rb_right)));
1153 			/*
1154 			 * If I'm not childless, thus I have black children,
1155 			 * then my brother must either be black or have two
1156 			 * black children.
1157 			 */
1158 			KASSERT(RB_CHILDLESS_P(self)
1159 				|| RB_BLACK_P(brother)
1160 				|| (RB_TWOCHILDREN_P(brother)
1161 				    && RB_BLACK_P(brother->rb_left)
1162 				    && RB_BLACK_P(brother->rb_right)));
1163 		} else {
1164 			/*
1165 			 * If I'm black and have one child, that child must
1166 			 * be red and childless.
1167 			 */
1168 			KASSERT(RB_CHILDLESS_P(self)
1169 				|| RB_TWOCHILDREN_P(self)
1170 				|| (!RB_LEFT_SENTINEL_P(self)
1171 				    && RB_RIGHT_SENTINEL_P(self)
1172 				    && RB_RED_P(self->rb_left)
1173 				    && RB_CHILDLESS_P(self->rb_left))
1174 				|| (!RB_RIGHT_SENTINEL_P(self)
1175 				    && RB_LEFT_SENTINEL_P(self)
1176 				    && RB_RED_P(self->rb_right)
1177 				    && RB_CHILDLESS_P(self->rb_right)));
1178 
1179 			/*
1180 			 * If I'm a childless black node and my parent is
1181 			 * black, my 2nd closet relative away from my parent
1182 			 * is either red or has a red parent or red children.
1183 			 */
1184 			if (!RB_ROOT_P(rbt, self)
1185 			    && RB_CHILDLESS_P(self)
1186 			    && RB_BLACK_P(RB_FATHER(self))) {
1187 				const unsigned int which = RB_POSITION(self);
1188 				const unsigned int other = which ^ RB_DIR_OTHER;
1189 				const struct rb_node *relative0, *relative;
1190 
1191 				relative0 = rb_tree_iterate_const(rbt,
1192 				    self, other);
1193 				KASSERT(relative0 != NULL);
1194 				relative = rb_tree_iterate_const(rbt,
1195 				    relative0, other);
1196 				KASSERT(relative != NULL);
1197 				KASSERT(RB_SENTINEL_P(relative->rb_nodes[which]));
1198 #if 0
1199 				KASSERT(RB_RED_P(relative)
1200 					|| RB_RED_P(relative->rb_left)
1201 					|| RB_RED_P(relative->rb_right)
1202 					|| RB_RED_P(RB_FATHER(relative)));
1203 #endif
1204 			}
1205 		}
1206 		/*
1207 		 * A grandparent's children must be real nodes and not
1208 		 * sentinels.  First check out grandparent.
1209 		 */
1210 		KASSERT(RB_ROOT_P(rbt, self)
1211 			|| RB_ROOT_P(rbt, RB_FATHER(self))
1212 			|| RB_TWOCHILDREN_P(RB_FATHER(RB_FATHER(self))));
1213 		/*
1214 		 * If we are have grandchildren on our left, then
1215 		 * we must have a child on our right.
1216 		 */
1217 		KASSERT(RB_LEFT_SENTINEL_P(self)
1218 			|| RB_CHILDLESS_P(self->rb_left)
1219 			|| !RB_RIGHT_SENTINEL_P(self));
1220 		/*
1221 		 * If we are have grandchildren on our right, then
1222 		 * we must have a child on our left.
1223 		 */
1224 		KASSERT(RB_RIGHT_SENTINEL_P(self)
1225 			|| RB_CHILDLESS_P(self->rb_right)
1226 			|| !RB_LEFT_SENTINEL_P(self));
1227 
1228 		/*
1229 		 * If we have a child on the left and it doesn't have two
1230 		 * children make sure we don't have great-great-grandchildren on
1231 		 * the right.
1232 		 */
1233 		KASSERT(RB_TWOCHILDREN_P(self->rb_left)
1234 			|| RB_CHILDLESS_P(self->rb_right)
1235 			|| RB_CHILDLESS_P(self->rb_right->rb_left)
1236 			|| RB_CHILDLESS_P(self->rb_right->rb_left->rb_left)
1237 			|| RB_CHILDLESS_P(self->rb_right->rb_left->rb_right)
1238 			|| RB_CHILDLESS_P(self->rb_right->rb_right)
1239 			|| RB_CHILDLESS_P(self->rb_right->rb_right->rb_left)
1240 			|| RB_CHILDLESS_P(self->rb_right->rb_right->rb_right));
1241 
1242 		/*
1243 		 * If we have a child on the right and it doesn't have two
1244 		 * children make sure we don't have great-great-grandchildren on
1245 		 * the left.
1246 		 */
1247 		KASSERT(RB_TWOCHILDREN_P(self->rb_right)
1248 			|| RB_CHILDLESS_P(self->rb_left)
1249 			|| RB_CHILDLESS_P(self->rb_left->rb_left)
1250 			|| RB_CHILDLESS_P(self->rb_left->rb_left->rb_left)
1251 			|| RB_CHILDLESS_P(self->rb_left->rb_left->rb_right)
1252 			|| RB_CHILDLESS_P(self->rb_left->rb_right)
1253 			|| RB_CHILDLESS_P(self->rb_left->rb_right->rb_left)
1254 			|| RB_CHILDLESS_P(self->rb_left->rb_right->rb_right));
1255 
1256 		/*
1257 		 * If we are fully interior node, then our predecessors and
1258 		 * successors must have no children in our direction.
1259 		 */
1260 		if (RB_TWOCHILDREN_P(self)) {
1261 			const struct rb_node *prev0;
1262 			const struct rb_node *next0;
1263 
1264 			prev0 = rb_tree_iterate_const(rbt, self, RB_DIR_LEFT);
1265 			KASSERT(prev0 != NULL);
1266 			KASSERT(RB_RIGHT_SENTINEL_P(prev0));
1267 
1268 			next0 = rb_tree_iterate_const(rbt, self, RB_DIR_RIGHT);
1269 			KASSERT(next0 != NULL);
1270 			KASSERT(RB_LEFT_SENTINEL_P(next0));
1271 		}
1272 	}
1273 
1274 	return true;
1275 }
1276 
1277 void
rb_tree_check(const struct rb_tree * rbt,bool red_check)1278 rb_tree_check(const struct rb_tree *rbt, bool red_check)
1279 {
1280 	const struct rb_node *self;
1281 	const struct rb_node *prev;
1282 #ifdef RBSTATS
1283 	unsigned int count = 0;
1284 #endif
1285 
1286 	KASSERT(rbt->rbt_root != NULL);
1287 	KASSERT(RB_LEFT_P(rbt->rbt_root));
1288 
1289 #if defined(RBSTATS) && !defined(RBSMALL)
1290 	KASSERT(rbt->rbt_count > 1
1291 	    || rbt->rbt_minmax[RB_DIR_LEFT] == rbt->rbt_minmax[RB_DIR_RIGHT]);
1292 #endif
1293 
1294 	prev = NULL;
1295 	TAILQ_FOREACH(self, &rbt->rbt_nodes, rb_link) {
1296 		rb_tree_check_node(rbt, self, prev, false);
1297 #ifdef RBSTATS
1298 		count++;
1299 #endif
1300 	}
1301 #ifdef RBSTATS
1302 	KASSERT(rbt->rbt_count == count);
1303 #endif
1304 	if (red_check) {
1305 		KASSERT(RB_BLACK_P(rbt->rbt_root));
1306 		KASSERT(RB_SENTINEL_P(rbt->rbt_root)
1307 			|| rb_tree_count_black(rbt->rbt_root));
1308 
1309 		/*
1310 		 * The root must be black.
1311 		 * There can never be two adjacent red nodes.
1312 		 */
1313 		TAILQ_FOREACH(self, &rbt->rbt_nodes, rb_link) {
1314 			rb_tree_check_node(rbt, self, NULL, true);
1315 		}
1316 	}
1317 }
1318 #endif /* RBDEBUG */
1319 
1320 #ifdef RBSTATS
1321 static void
rb_tree_mark_depth(const struct rb_tree * rbt,const struct rb_node * self,size_t * depths,size_t depth)1322 rb_tree_mark_depth(const struct rb_tree *rbt, const struct rb_node *self,
1323 	size_t *depths, size_t depth)
1324 {
1325 	if (RB_SENTINEL_P(self))
1326 		return;
1327 
1328 	if (RB_TWOCHILDREN_P(self)) {
1329 		rb_tree_mark_depth(rbt, self->rb_left, depths, depth + 1);
1330 		rb_tree_mark_depth(rbt, self->rb_right, depths, depth + 1);
1331 		return;
1332 	}
1333 	depths[depth]++;
1334 	if (!RB_LEFT_SENTINEL_P(self)) {
1335 		rb_tree_mark_depth(rbt, self->rb_left, depths, depth + 1);
1336 	}
1337 	if (!RB_RIGHT_SENTINEL_P(self)) {
1338 		rb_tree_mark_depth(rbt, self->rb_right, depths, depth + 1);
1339 	}
1340 }
1341 
1342 void
rb_tree_depths(const struct rb_tree * rbt,size_t * depths)1343 rb_tree_depths(const struct rb_tree *rbt, size_t *depths)
1344 {
1345 	rb_tree_mark_depth(rbt, rbt->rbt_root, depths, 1);
1346 }
1347 #endif /* RBSTATS */
1348