1/*
2 * Copyright (C) 2012 by Darren Reed.
3 *
4 * See the IPFILTER.LICENCE file for details on licencing.
5 *
6 */
7typedef enum rbcolour_e {
8	C_BLACK = 0,
9	C_RED = 1
10} rbcolour_t;
11
12#define	RBI_LINK(_n, _t)							\
13	struct _n##_rb_link {						\
14		struct _t	*left;					\
15		struct _t	*right;					\
16		struct _t	*parent;				\
17		rbcolour_t	colour;					\
18	}
19
20#define	RBI_HEAD(_n, _t)						\
21struct _n##_rb_head {							\
22	struct _t	top;						\
23	int		count;						\
24	int		(* compare)(struct _t *, struct _t *);		\
25}
26
27#define	RBI_CODE(_n, _t, _f, _cmp)					\
28									\
29typedef	void	(*_n##_rb_walker_t)(_t *, void *);			\
30									\
31_t *	_n##_rb_delete(struct _n##_rb_head *, _t *);			\
32void	_n##_rb_init(struct _n##_rb_head *);				\
33void	_n##_rb_insert(struct _n##_rb_head *, _t *);			\
34_t *	_n##_rb_search(struct _n##_rb_head *, void *);			\
35void	_n##_rb_walktree(struct _n##_rb_head *, _n##_rb_walker_t, void *);\
36									\
37static void								\
38rotate_left(struct _n##_rb_head *head, _t *node)			\
39{									\
40	_t *parent, *tmp1, *tmp2;					\
41									\
42	parent = node->_f.parent;					\
43	tmp1 = node->_f.right;						\
44	tmp2 = tmp1->_f.left;						\
45	node->_f.right = tmp2;						\
46	if (tmp2 != & _n##_rb_zero)					\
47		tmp2->_f.parent = node;					\
48	if (parent == & _n##_rb_zero)					\
49		head->top._f.right = tmp1;				\
50	else if (parent->_f.right == node)				\
51		parent->_f.right = tmp1;				\
52	else								\
53		parent->_f.left = tmp1;					\
54	tmp1->_f.left = node;						\
55	tmp1->_f.parent = parent;					\
56	node->_f.parent = tmp1;						\
57}									\
58									\
59static void								\
60rotate_right(struct _n##_rb_head *head, _t *node)			\
61{									\
62	_t *parent, *tmp1, *tmp2;					\
63									\
64	parent = node->_f.parent;					\
65	tmp1 = node->_f.left;						\
66	tmp2 = tmp1->_f.right;						\
67	node->_f.left = tmp2;						\
68	if (tmp2 != &_n##_rb_zero)					\
69		tmp2->_f.parent = node;					\
70	if (parent == &_n##_rb_zero)					\
71		head->top._f.right = tmp1;				\
72	else if (parent->_f.right == node)				\
73		parent->_f.right = tmp1;				\
74	else								\
75		parent->_f.left = tmp1;					\
76	tmp1->_f.right = node;						\
77	tmp1->_f.parent = parent;					\
78	node->_f.parent = tmp1;						\
79}									\
80									\
81void									\
82_n##_rb_insert(struct _n##_rb_head *head, _t *node)			\
83{									\
84	_t *n, *parent, **p, *tmp1, *gparent;				\
85									\
86	parent = &head->top;						\
87	node->_f.left = &_n##_rb_zero;					\
88	node->_f.right = &_n##_rb_zero;					\
89	p = &head->top._f.right;					\
90	while ((n = *p) != &_n##_rb_zero) {				\
91		if (_cmp(node, n) < 0)					\
92			p = &n->_f.left;				\
93		else							\
94			p = &n->_f.right;				\
95		parent = n;						\
96	}								\
97	*p = node;							\
98	node->_f.colour = C_RED;					\
99	node->_f.parent = parent;					\
100									\
101	while ((node != &_n##_rb_zero) && (parent->_f.colour == C_RED)){\
102		gparent = parent->_f.parent;				\
103		if (parent == gparent->_f.left) {			\
104			tmp1 = gparent->_f.right;			\
105			if (tmp1->_f.colour == C_RED) {			\
106				parent->_f.colour = C_BLACK;		\
107				tmp1->_f.colour = C_BLACK;		\
108				gparent->_f.colour = C_RED;		\
109				node = gparent;				\
110			} else {					\
111				if (node == parent->_f.right) {		\
112					node = parent;			\
113					rotate_left(head, node);	\
114					parent = node->_f.parent;	\
115				}					\
116				parent->_f.colour = C_BLACK;		\
117				gparent->_f.colour = C_RED;		\
118				rotate_right(head, gparent);		\
119			}						\
120		} else {						\
121			tmp1 = gparent->_f.left;			\
122			if (tmp1->_f.colour == C_RED) {			\
123				parent->_f.colour = C_BLACK;		\
124				tmp1->_f.colour = C_BLACK;		\
125				gparent->_f.colour = C_RED;		\
126				node = gparent;				\
127			} else {					\
128				if (node == parent->_f.left) {		\
129					node = parent;			\
130					rotate_right(head, node);	\
131					parent = node->_f.parent;	\
132				}					\
133				parent->_f.colour = C_BLACK;		\
134				gparent->_f.colour = C_RED;		\
135				rotate_left(head, parent->_f.parent);	\
136			}						\
137		}							\
138		parent = node->_f.parent;				\
139	}								\
140	head->top._f.right->_f.colour = C_BLACK;			\
141	head->count++;						\
142}									\
143									\
144static void								\
145deleteblack(struct _n##_rb_head *head, _t *parent, _t *node)		\
146{									\
147	_t *tmp;							\
148									\
149	while ((node == &_n##_rb_zero || node->_f.colour == C_BLACK) &&	\
150	       node != &head->top) {					\
151		if (parent->_f.left == node) {				\
152			tmp = parent->_f.right;				\
153			if (tmp->_f.colour == C_RED) {			\
154				tmp->_f.colour = C_BLACK;		\
155				parent->_f.colour = C_RED;		\
156				rotate_left(head, parent);		\
157				tmp = parent->_f.right;			\
158			}						\
159			if ((tmp->_f.left == &_n##_rb_zero ||		\
160			     tmp->_f.left->_f.colour == C_BLACK) &&	\
161			    (tmp->_f.right == &_n##_rb_zero ||		\
162			     tmp->_f.right->_f.colour == C_BLACK)) {	\
163				tmp->_f.colour = C_RED;			\
164				node = parent;				\
165				parent = node->_f.parent;		\
166			} else {					\
167				if (tmp->_f.right == &_n##_rb_zero ||	\
168				    tmp->_f.right->_f.colour == C_BLACK) {\
169					_t *tmp2 = tmp->_f.left;	\
170									\
171					if (tmp2 != &_n##_rb_zero)	\
172						tmp2->_f.colour = C_BLACK;\
173					tmp->_f.colour = C_RED;		\
174					rotate_right(head, tmp);	\
175					tmp = parent->_f.right;		\
176				}					\
177				tmp->_f.colour = parent->_f.colour;	\
178				parent->_f.colour = C_BLACK;		\
179				if (tmp->_f.right != &_n##_rb_zero)	\
180					tmp->_f.right->_f.colour = C_BLACK;\
181				rotate_left(head, parent);		\
182				node = head->top._f.right;		\
183			}						\
184		} else {						\
185			tmp = parent->_f.left;				\
186			if (tmp->_f.colour == C_RED) {			\
187				tmp->_f.colour = C_BLACK;		\
188				parent->_f.colour = C_RED;		\
189				rotate_right(head, parent);		\
190				tmp = parent->_f.left;			\
191			}						\
192			if ((tmp->_f.left == &_n##_rb_zero ||		\
193			     tmp->_f.left->_f.colour == C_BLACK) &&	\
194			    (tmp->_f.right == &_n##_rb_zero ||		\
195			     tmp->_f.right->_f.colour == C_BLACK)) {	\
196				tmp->_f.colour = C_RED;			\
197				node = parent;				\
198				parent = node->_f.parent;		\
199			} else {					\
200				if (tmp->_f.left == &_n##_rb_zero ||	\
201				    tmp->_f.left->_f.colour == C_BLACK) {\
202					_t *tmp2 = tmp->_f.right;	\
203									\
204					if (tmp2 != &_n##_rb_zero)	\
205						tmp2->_f.colour = C_BLACK;\
206					tmp->_f.colour = C_RED;		\
207					rotate_left(head, tmp);		\
208					tmp = parent->_f.left;		\
209				}					\
210				tmp->_f.colour = parent->_f.colour;	\
211				parent->_f.colour = C_BLACK;		\
212				if (tmp->_f.left != &_n##_rb_zero)	\
213					tmp->_f.left->_f.colour = C_BLACK;\
214				rotate_right(head, parent);		\
215				node = head->top._f.right;		\
216				break;					\
217			}						\
218		}							\
219	}								\
220	if (node != &_n##_rb_zero)					\
221		node->_f.colour = C_BLACK;				\
222}									\
223									\
224_t *									\
225_n##_rb_delete(struct _n##_rb_head *head, _t *node)			\
226{									\
227	_t *child, *parent, *old = node, *left;				\
228	rbcolour_t color;						\
229									\
230	if (node->_f.left == &_n##_rb_zero) {				\
231		child = node->_f.right;					\
232	} else if (node->_f.right == &_n##_rb_zero) {			\
233		child = node->_f.left;					\
234	} else {							\
235		node = node->_f.right;					\
236		while ((left = node->_f.left) != &_n##_rb_zero)		\
237			node = left;					\
238		child = node->_f.right;					\
239		parent = node->_f.parent;				\
240		color = node->_f.colour;				\
241		if (child != &_n##_rb_zero)				\
242			child->_f.parent = parent;			\
243		if (parent != &_n##_rb_zero) {				\
244			if (parent->_f.left == node)			\
245				parent->_f.left = child;		\
246			else						\
247				parent->_f.right = child;		\
248		} else {						\
249			head->top._f.right = child;			\
250		}							\
251		if (node->_f.parent == old)				\
252			parent = node;					\
253		*node = *old;						\
254		if (old->_f.parent != &_n##_rb_zero) {			\
255			if (old->_f.parent->_f.left == old)		\
256				old->_f.parent->_f.left = node;		\
257			else						\
258				old->_f.parent->_f.right = node;	\
259		} else {						\
260			head->top._f.right = child;			\
261		}							\
262		old->_f.left->_f.parent = node;				\
263		if (old->_f.right != &_n##_rb_zero)			\
264			old->_f.right->_f.parent = node;		\
265		if (parent != &_n##_rb_zero) {				\
266			left = parent;					\
267		}							\
268		goto colour;						\
269	}								\
270	parent = node->_f.parent;					\
271	color= node->_f.colour;						\
272	if (child != &_n##_rb_zero)					\
273		child->_f.parent = parent;				\
274	if (parent != &_n##_rb_zero) {					\
275		if (parent->_f.left == node)				\
276			parent->_f.left = child;			\
277		else							\
278			parent->_f.right = child;			\
279	} else {							\
280		head->top._f.right = child;				\
281	}								\
282colour:									\
283	if (color == C_BLACK)						\
284		deleteblack(head, parent, node);			\
285	head->count--;							\
286	return old;							\
287}									\
288									\
289void									\
290_n##_rb_init(struct _n##_rb_head *head)					\
291{									\
292	memset(head, 0, sizeof(*head));					\
293	memset(&_n##_rb_zero, 0, sizeof(_n##_rb_zero));			\
294	head->top._f.left = &_n##_rb_zero;				\
295	head->top._f.right = &_n##_rb_zero;				\
296	head->top._f.parent = &head->top;				\
297	_n##_rb_zero._f.left = &_n##_rb_zero;				\
298	_n##_rb_zero._f.right = &_n##_rb_zero;				\
299	_n##_rb_zero._f.parent = &_n##_rb_zero;				\
300}									\
301									\
302void									\
303_n##_rb_walktree(struct _n##_rb_head *head, _n##_rb_walker_t func, void *arg)\
304{									\
305	_t *prev;							\
306	_t *next;							\
307	_t *node = head->top._f.right;					\
308	_t *base;							\
309									\
310	while (node != &_n##_rb_zero)					\
311		node = node->_f.left;					\
312									\
313	for (;;) {							\
314		base = node;						\
315		prev = node;						\
316		while ((node->_f.parent->_f.right == node) &&		\
317		       (node != &_n##_rb_zero))	{			\
318			prev = node;					\
319			node = node->_f.parent;				\
320		}							\
321									\
322		node = prev;						\
323		for (node = node->_f.parent->_f.right; node != &_n##_rb_zero;\
324		     node = node->_f.left)				\
325			prev = node;					\
326		next = prev;						\
327									\
328		if (node != &_n##_rb_zero)				\
329			func(node, arg);				\
330									\
331		node = next;						\
332		if (node == &_n##_rb_zero)				\
333			break;						\
334	}								\
335}									\
336									\
337_t *									\
338_n##_rb_search(struct _n##_rb_head *head, void *key)			\
339{									\
340	int	match;							\
341	_t	*node;							\
342	node = head->top._f.right;					\
343	while (node != &_n##_rb_zero) {					\
344		match = _cmp(key, node);				\
345		if (match == 0)						\
346			break;						\
347		if (match< 0)						\
348			node = node->_f.left;				\
349		else							\
350			node = node->_f.right;				\
351	}								\
352	if (node == &_n##_rb_zero || match != 0)			\
353		return (NULL);						\
354	return (node);							\
355}
356
357#define	RBI_DELETE(_n, _h, _v)		_n##_rb_delete(_h, _v)
358#define	RBI_FIELD(_n)			struct _n##_rb_link
359#define	RBI_INIT(_n, _h)		_n##_rb_init(_h)
360#define	RBI_INSERT(_n, _h, _v)		_n##_rb_insert(_h, _v)
361#define	RBI_ISEMPTY(_h)			((_h)->count == 0)
362#define	RBI_SEARCH(_n, _h, _k)		_n##_rb_search(_h, _k)
363#define	RBI_WALK(_n, _h, _w, _a)	_n##_rb_walktree(_h, _w, _a)
364#define	RBI_ZERO(_n)			_n##_rb_zero
365