1/*	$OpenBSD: subr_tree.c,v 1.10 2018/10/09 08:28:43 dlg Exp $ */
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/tree.h>
45
46static inline struct rb_entry *
47rb_n2e(const struct rb_type *t, void *node)
48{
49	unsigned long addr = (unsigned long)node;
50
51	return ((struct rb_entry *)(addr + t->t_offset));
52}
53
54static inline void *
55rb_e2n(const struct rb_type *t, struct rb_entry *rbe)
56{
57	unsigned long addr = (unsigned long)rbe;
58
59	return ((void *)(addr - t->t_offset));
60}
61
62#define RBE_LEFT(_rbe)		(_rbe)->rbt_left
63#define RBE_RIGHT(_rbe)		(_rbe)->rbt_right
64#define RBE_PARENT(_rbe)	(_rbe)->rbt_parent
65#define RBE_COLOR(_rbe)		(_rbe)->rbt_color
66
67#define RBH_ROOT(_rbt)		(_rbt)->rbt_root
68
69static inline void
70rbe_set(struct rb_entry *rbe, struct rb_entry *parent)
71{
72	RBE_PARENT(rbe) = parent;
73	RBE_LEFT(rbe) = RBE_RIGHT(rbe) = NULL;
74	RBE_COLOR(rbe) = RB_RED;
75}
76
77static inline void
78rbe_set_blackred(struct rb_entry *black, struct rb_entry *red)
79{
80	RBE_COLOR(black) = RB_BLACK;
81	RBE_COLOR(red) = RB_RED;
82}
83
84static inline void
85rbe_augment(const struct rb_type *t, struct rb_entry *rbe)
86{
87	(*t->t_augment)(rb_e2n(t, rbe));
88}
89
90static inline void
91rbe_if_augment(const struct rb_type *t, struct rb_entry *rbe)
92{
93	if (t->t_augment != NULL)
94		rbe_augment(t, rbe);
95}
96
97static inline void
98rbe_rotate_left(const struct rb_type *t, struct rb_tree *rbt,
99    struct rb_entry *rbe)
100{
101	struct rb_entry *parent;
102	struct rb_entry *tmp;
103
104	tmp = RBE_RIGHT(rbe);
105	RBE_RIGHT(rbe) = RBE_LEFT(tmp);
106	if (RBE_RIGHT(rbe) != NULL)
107		RBE_PARENT(RBE_LEFT(tmp)) = rbe;
108
109	parent = RBE_PARENT(rbe);
110	RBE_PARENT(tmp) = parent;
111	if (parent != NULL) {
112		if (rbe == RBE_LEFT(parent))
113			RBE_LEFT(parent) = tmp;
114		else
115			RBE_RIGHT(parent) = tmp;
116	} else
117		RBH_ROOT(rbt) = tmp;
118
119	RBE_LEFT(tmp) = rbe;
120	RBE_PARENT(rbe) = tmp;
121
122	if (t->t_augment != NULL) {
123		rbe_augment(t, rbe);
124		rbe_augment(t, tmp);
125		parent = RBE_PARENT(tmp);
126		if (parent != NULL)
127			rbe_augment(t, parent);
128	}
129}
130
131static inline void
132rbe_rotate_right(const struct rb_type *t, struct rb_tree *rbt,
133    struct rb_entry *rbe)
134{
135	struct rb_entry *parent;
136	struct rb_entry *tmp;
137
138	tmp = RBE_LEFT(rbe);
139	RBE_LEFT(rbe) = RBE_RIGHT(tmp);
140	if (RBE_LEFT(rbe) != NULL)
141		RBE_PARENT(RBE_RIGHT(tmp)) = rbe;
142
143	parent = RBE_PARENT(rbe);
144	RBE_PARENT(tmp) = parent;
145	if (parent != NULL) {
146		if (rbe == RBE_LEFT(parent))
147			RBE_LEFT(parent) = tmp;
148		else
149			RBE_RIGHT(parent) = tmp;
150	} else
151		RBH_ROOT(rbt) = tmp;
152
153	RBE_RIGHT(tmp) = rbe;
154	RBE_PARENT(rbe) = tmp;
155
156	if (t->t_augment != NULL) {
157		rbe_augment(t, rbe);
158		rbe_augment(t, tmp);
159		parent = RBE_PARENT(tmp);
160		if (parent != NULL)
161			rbe_augment(t, parent);
162	}
163}
164
165static inline void
166rbe_insert_color(const struct rb_type *t, struct rb_tree *rbt,
167    struct rb_entry *rbe)
168{
169	struct rb_entry *parent, *gparent, *tmp;
170
171	while ((parent = RBE_PARENT(rbe)) != NULL &&
172	    RBE_COLOR(parent) == RB_RED) {
173		gparent = RBE_PARENT(parent);
174
175		if (parent == RBE_LEFT(gparent)) {
176			tmp = RBE_RIGHT(gparent);
177			if (tmp != NULL && RBE_COLOR(tmp) == RB_RED) {
178				RBE_COLOR(tmp) = RB_BLACK;
179				rbe_set_blackred(parent, gparent);
180				rbe = gparent;
181				continue;
182			}
183
184			if (RBE_RIGHT(parent) == rbe) {
185				rbe_rotate_left(t, rbt, parent);
186				tmp = parent;
187				parent = rbe;
188				rbe = tmp;
189			}
190
191			rbe_set_blackred(parent, gparent);
192			rbe_rotate_right(t, rbt, gparent);
193		} else {
194			tmp = RBE_LEFT(gparent);
195			if (tmp != NULL && RBE_COLOR(tmp) == RB_RED) {
196				RBE_COLOR(tmp) = RB_BLACK;
197				rbe_set_blackred(parent, gparent);
198				rbe = gparent;
199				continue;
200			}
201
202			if (RBE_LEFT(parent) == rbe) {
203				rbe_rotate_right(t, rbt, parent);
204				tmp = parent;
205				parent = rbe;
206				rbe = tmp;
207			}
208
209			rbe_set_blackred(parent, gparent);
210			rbe_rotate_left(t, rbt, gparent);
211		}
212	}
213
214	RBE_COLOR(RBH_ROOT(rbt)) = RB_BLACK;
215}
216
217static inline void
218rbe_remove_color(const struct rb_type *t, struct rb_tree *rbt,
219    struct rb_entry *parent, struct rb_entry *rbe)
220{
221	struct rb_entry *tmp;
222
223	while ((rbe == NULL || RBE_COLOR(rbe) == RB_BLACK) &&
224	    rbe != RBH_ROOT(rbt)) {
225		if (RBE_LEFT(parent) == rbe) {
226			tmp = RBE_RIGHT(parent);
227			if (RBE_COLOR(tmp) == RB_RED) {
228				rbe_set_blackred(tmp, parent);
229				rbe_rotate_left(t, rbt, parent);
230				tmp = RBE_RIGHT(parent);
231			}
232			if ((RBE_LEFT(tmp) == NULL ||
233			     RBE_COLOR(RBE_LEFT(tmp)) == RB_BLACK) &&
234			    (RBE_RIGHT(tmp) == NULL ||
235			     RBE_COLOR(RBE_RIGHT(tmp)) == RB_BLACK)) {
236				RBE_COLOR(tmp) = RB_RED;
237				rbe = parent;
238				parent = RBE_PARENT(rbe);
239			} else {
240				if (RBE_RIGHT(tmp) == NULL ||
241				    RBE_COLOR(RBE_RIGHT(tmp)) == RB_BLACK) {
242					struct rb_entry *oleft;
243
244					oleft = RBE_LEFT(tmp);
245					if (oleft != NULL)
246						RBE_COLOR(oleft) = RB_BLACK;
247
248					RBE_COLOR(tmp) = RB_RED;
249					rbe_rotate_right(t, rbt, tmp);
250					tmp = RBE_RIGHT(parent);
251				}
252
253				RBE_COLOR(tmp) = RBE_COLOR(parent);
254				RBE_COLOR(parent) = RB_BLACK;
255				if (RBE_RIGHT(tmp))
256					RBE_COLOR(RBE_RIGHT(tmp)) = RB_BLACK;
257
258				rbe_rotate_left(t, rbt, parent);
259				rbe = RBH_ROOT(rbt);
260				break;
261			}
262		} else {
263			tmp = RBE_LEFT(parent);
264			if (RBE_COLOR(tmp) == RB_RED) {
265				rbe_set_blackred(tmp, parent);
266				rbe_rotate_right(t, rbt, parent);
267				tmp = RBE_LEFT(parent);
268			}
269
270			if ((RBE_LEFT(tmp) == NULL ||
271			     RBE_COLOR(RBE_LEFT(tmp)) == RB_BLACK) &&
272			    (RBE_RIGHT(tmp) == NULL ||
273			     RBE_COLOR(RBE_RIGHT(tmp)) == RB_BLACK)) {
274				RBE_COLOR(tmp) = RB_RED;
275				rbe = parent;
276				parent = RBE_PARENT(rbe);
277			} else {
278				if (RBE_LEFT(tmp) == NULL ||
279				    RBE_COLOR(RBE_LEFT(tmp)) == RB_BLACK) {
280					struct rb_entry *oright;
281
282					oright = RBE_RIGHT(tmp);
283					if (oright != NULL)
284						RBE_COLOR(oright) = RB_BLACK;
285
286					RBE_COLOR(tmp) = RB_RED;
287					rbe_rotate_left(t, rbt, tmp);
288					tmp = RBE_LEFT(parent);
289				}
290
291				RBE_COLOR(tmp) = RBE_COLOR(parent);
292				RBE_COLOR(parent) = RB_BLACK;
293				if (RBE_LEFT(tmp) != NULL)
294					RBE_COLOR(RBE_LEFT(tmp)) = RB_BLACK;
295
296				rbe_rotate_right(t, rbt, parent);
297				rbe = RBH_ROOT(rbt);
298				break;
299			}
300		}
301	}
302
303	if (rbe != NULL)
304		RBE_COLOR(rbe) = RB_BLACK;
305}
306
307static inline struct rb_entry *
308rbe_remove(const struct rb_type *t, struct rb_tree *rbt, struct rb_entry *rbe)
309{
310	struct rb_entry *child, *parent, *old = rbe;
311	unsigned int color;
312
313	if (RBE_LEFT(rbe) == NULL)
314		child = RBE_RIGHT(rbe);
315	else if (RBE_RIGHT(rbe) == NULL)
316		child = RBE_LEFT(rbe);
317	else {
318		struct rb_entry *tmp;
319
320		rbe = RBE_RIGHT(rbe);
321		while ((tmp = RBE_LEFT(rbe)) != NULL)
322			rbe = tmp;
323
324		child = RBE_RIGHT(rbe);
325		parent = RBE_PARENT(rbe);
326		color = RBE_COLOR(rbe);
327		if (child != NULL)
328			RBE_PARENT(child) = parent;
329		if (parent != NULL) {
330			if (RBE_LEFT(parent) == rbe)
331				RBE_LEFT(parent) = child;
332			else
333				RBE_RIGHT(parent) = child;
334
335			rbe_if_augment(t, parent);
336		} else
337			RBH_ROOT(rbt) = child;
338		if (RBE_PARENT(rbe) == old)
339			parent = rbe;
340		*rbe = *old;
341
342		tmp = RBE_PARENT(old);
343		if (tmp != NULL) {
344			if (RBE_LEFT(tmp) == old)
345				RBE_LEFT(tmp) = rbe;
346			else
347				RBE_RIGHT(tmp) = rbe;
348
349			rbe_if_augment(t, tmp);
350		} else
351			RBH_ROOT(rbt) = rbe;
352
353		RBE_PARENT(RBE_LEFT(old)) = rbe;
354		if (RBE_RIGHT(old))
355			RBE_PARENT(RBE_RIGHT(old)) = rbe;
356
357		if (t->t_augment != NULL && parent != NULL) {
358			tmp = parent;
359			do {
360				rbe_augment(t, tmp);
361				tmp = RBE_PARENT(tmp);
362			} while (tmp != NULL);
363		}
364
365		goto color;
366	}
367
368	parent = RBE_PARENT(rbe);
369	color = RBE_COLOR(rbe);
370
371	if (child != NULL)
372		RBE_PARENT(child) = parent;
373	if (parent != NULL) {
374		if (RBE_LEFT(parent) == rbe)
375			RBE_LEFT(parent) = child;
376		else
377			RBE_RIGHT(parent) = child;
378
379		rbe_if_augment(t, parent);
380	} else
381		RBH_ROOT(rbt) = child;
382color:
383	if (color == RB_BLACK)
384		rbe_remove_color(t, rbt, parent, child);
385
386	return (old);
387}
388
389void *
390_rb_remove(const struct rb_type *t, struct rb_tree *rbt, void *elm)
391{
392	struct rb_entry *rbe = rb_n2e(t, elm);
393	struct rb_entry *old;
394
395	old = rbe_remove(t, rbt, rbe);
396
397	return (old == NULL ? NULL : rb_e2n(t, old));
398}
399
400void *
401_rb_insert(const struct rb_type *t, struct rb_tree *rbt, void *elm)
402{
403	struct rb_entry *rbe = rb_n2e(t, elm);
404	struct rb_entry *tmp;
405	struct rb_entry *parent = NULL;
406	void *node;
407	int comp = 0;
408
409	tmp = RBH_ROOT(rbt);
410	while (tmp != NULL) {
411		parent = tmp;
412
413		node = rb_e2n(t, tmp);
414		comp = (*t->t_compare)(elm, node);
415		if (comp < 0)
416			tmp = RBE_LEFT(tmp);
417		else if (comp > 0)
418			tmp = RBE_RIGHT(tmp);
419		else
420			return (node);
421	}
422
423	rbe_set(rbe, parent);
424
425	if (parent != NULL) {
426		if (comp < 0)
427			RBE_LEFT(parent) = rbe;
428		else
429			RBE_RIGHT(parent) = rbe;
430
431		rbe_if_augment(t, parent);
432	} else
433		RBH_ROOT(rbt) = rbe;
434
435	rbe_insert_color(t, rbt, rbe);
436
437	return (NULL);
438}
439
440/* Finds the node with the same key as elm */
441void *
442_rb_find(const struct rb_type *t, struct rb_tree *rbt, const void *key)
443{
444	struct rb_entry *tmp = RBH_ROOT(rbt);
445	void *node;
446	int comp;
447
448	while (tmp != NULL) {
449		node = rb_e2n(t, tmp);
450		comp = (*t->t_compare)(key, node);
451		if (comp < 0)
452			tmp = RBE_LEFT(tmp);
453		else if (comp > 0)
454			tmp = RBE_RIGHT(tmp);
455		else
456			return (node);
457	}
458
459	return (NULL);
460}
461
462/* Finds the first node greater than or equal to the search key */
463void *
464_rb_nfind(const struct rb_type *t, struct rb_tree *rbt, const void *key)
465{
466	struct rb_entry *tmp = RBH_ROOT(rbt);
467	void *node;
468	void *res = NULL;
469	int comp;
470
471	while (tmp != NULL) {
472		node = rb_e2n(t, tmp);
473		comp = (*t->t_compare)(key, node);
474		if (comp < 0) {
475			res = node;
476			tmp = RBE_LEFT(tmp);
477		} else if (comp > 0)
478			tmp = RBE_RIGHT(tmp);
479		else
480			return (node);
481	}
482
483	return (res);
484}
485
486void *
487_rb_next(const struct rb_type *t, void *elm)
488{
489	struct rb_entry *rbe = rb_n2e(t, elm);
490
491	if (RBE_RIGHT(rbe) != NULL) {
492		rbe = RBE_RIGHT(rbe);
493		while (RBE_LEFT(rbe) != NULL)
494			rbe = RBE_LEFT(rbe);
495	} else {
496		if (RBE_PARENT(rbe) &&
497		    (rbe == RBE_LEFT(RBE_PARENT(rbe))))
498			rbe = RBE_PARENT(rbe);
499		else {
500			while (RBE_PARENT(rbe) &&
501			    (rbe == RBE_RIGHT(RBE_PARENT(rbe))))
502				rbe = RBE_PARENT(rbe);
503			rbe = RBE_PARENT(rbe);
504		}
505	}
506
507	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
508}
509
510void *
511_rb_prev(const struct rb_type *t, void *elm)
512{
513	struct rb_entry *rbe = rb_n2e(t, elm);
514
515	if (RBE_LEFT(rbe)) {
516		rbe = RBE_LEFT(rbe);
517		while (RBE_RIGHT(rbe))
518			rbe = RBE_RIGHT(rbe);
519	} else {
520		if (RBE_PARENT(rbe) &&
521		    (rbe == RBE_RIGHT(RBE_PARENT(rbe))))
522			rbe = RBE_PARENT(rbe);
523		else {
524			while (RBE_PARENT(rbe) &&
525			    (rbe == RBE_LEFT(RBE_PARENT(rbe))))
526				rbe = RBE_PARENT(rbe);
527			rbe = RBE_PARENT(rbe);
528		}
529	}
530
531	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
532}
533
534void *
535_rb_root(const struct rb_type *t, struct rb_tree *rbt)
536{
537	struct rb_entry *rbe = RBH_ROOT(rbt);
538
539	return (rbe == NULL ? rbe : rb_e2n(t, rbe));
540}
541
542void *
543_rb_min(const struct rb_type *t, struct rb_tree *rbt)
544{
545	struct rb_entry *rbe = RBH_ROOT(rbt);
546	struct rb_entry *parent = NULL;
547
548	while (rbe != NULL) {
549		parent = rbe;
550		rbe = RBE_LEFT(rbe);
551	}
552
553	return (parent == NULL ? NULL : rb_e2n(t, parent));
554}
555
556void *
557_rb_max(const struct rb_type *t, struct rb_tree *rbt)
558{
559	struct rb_entry *rbe = RBH_ROOT(rbt);
560	struct rb_entry *parent = NULL;
561
562	while (rbe != NULL) {
563		parent = rbe;
564		rbe = RBE_RIGHT(rbe);
565	}
566
567	return (parent == NULL ? NULL : rb_e2n(t, parent));
568}
569
570void *
571_rb_left(const struct rb_type *t, void *node)
572{
573	struct rb_entry *rbe = rb_n2e(t, node);
574	rbe = RBE_LEFT(rbe);
575	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
576}
577
578void *
579_rb_right(const struct rb_type *t, void *node)
580{
581	struct rb_entry *rbe = rb_n2e(t, node);
582	rbe = RBE_RIGHT(rbe);
583	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
584}
585
586void *
587_rb_parent(const struct rb_type *t, void *node)
588{
589	struct rb_entry *rbe = rb_n2e(t, node);
590	rbe = RBE_PARENT(rbe);
591	return (rbe == NULL ? NULL : rb_e2n(t, rbe));
592}
593
594void
595_rb_set_left(const struct rb_type *t, void *node, void *left)
596{
597	struct rb_entry *rbe = rb_n2e(t, node);
598	struct rb_entry *rbl = (left == NULL) ? NULL : rb_n2e(t, left);
599
600	RBE_LEFT(rbe) = rbl;
601}
602
603void
604_rb_set_right(const struct rb_type *t, void *node, void *right)
605{
606	struct rb_entry *rbe = rb_n2e(t, node);
607	struct rb_entry *rbr = (right == NULL) ? NULL : rb_n2e(t, right);
608
609	RBE_RIGHT(rbe) = rbr;
610}
611
612void
613_rb_set_parent(const struct rb_type *t, void *node, void *parent)
614{
615	struct rb_entry *rbe = rb_n2e(t, node);
616	struct rb_entry *rbp = (parent == NULL) ? NULL : rb_n2e(t, parent);
617
618	RBE_PARENT(rbe) = rbp;
619}
620
621void
622_rb_poison(const struct rb_type *t, void *node, unsigned long poison)
623{
624	struct rb_entry *rbe = rb_n2e(t, node);
625
626	RBE_PARENT(rbe) = RBE_LEFT(rbe) = RBE_RIGHT(rbe) =
627	    (struct rb_entry *)poison;
628}
629
630int
631_rb_check(const struct rb_type *t, void *node, unsigned long poison)
632{
633	struct rb_entry *rbe = rb_n2e(t, node);
634
635	return ((unsigned long)RBE_PARENT(rbe) == poison &&
636	    (unsigned long)RBE_LEFT(rbe) == poison &&
637	    (unsigned long)RBE_RIGHT(rbe) == poison);
638}
639