1/*
2  Red Black Trees
3  (C) 1999  Andrea Arcangeli <andrea@suse.de>
4  (C) 2002  David Woodhouse <dwmw2@infradead.org>
5  (C) 2012  Michel Lespinasse <walken@google.com>
6
7  This program is free software; you can redistribute it and/or modify
8  it under the terms of the GNU General Public License as published by
9  the Free Software Foundation; either version 2 of the License, or
10  (at your option) any later version.
11
12  This program is distributed in the hope that it will be useful,
13  but WITHOUT ANY WARRANTY; without even the implied warranty of
14  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15  GNU General Public License for more details.
16
17  You should have received a copy of the GNU General Public License
18  along with this program; if not, write to the Free Software
19  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
20
21  linux/include/linux/rbtree_augmented.h
22*/
23
24#ifndef _LINUX_RBTREE_AUGMENTED_H
25#define _LINUX_RBTREE_AUGMENTED_H
26
27#include <linux/compiler.h>
28#include <linux/rbtree.h>
29
30/*
31 * Please note - only struct rb_augment_callbacks and the prototypes for
32 * rb_insert_augmented() and rb_erase_augmented() are intended to be public.
33 * The rest are implementation details you are not expected to depend on.
34 *
35 * See Documentation/rbtree.txt for documentation and samples.
36 */
37
38struct rb_augment_callbacks {
39	void (*propagate)(struct rb_node *node, struct rb_node *stop);
40	void (*copy)(struct rb_node *old, struct rb_node *new);
41	void (*rotate)(struct rb_node *old, struct rb_node *new);
42};
43
44extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
45	void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
46static inline void
47rb_insert_augmented(struct rb_node *node, struct rb_root *root,
48		    const struct rb_augment_callbacks *augment)
49{
50	__rb_insert_augmented(node, root, augment->rotate);
51}
52
53#define RB_DECLARE_CALLBACKS(rbstatic, rbname, rbstruct, rbfield,	\
54			     rbtype, rbaugmented, rbcompute)		\
55static inline void							\
56rbname ## _propagate(struct rb_node *rb, struct rb_node *stop)		\
57{									\
58	while (rb != stop) {						\
59		rbstruct *node = rb_entry(rb, rbstruct, rbfield);	\
60		rbtype augmented = rbcompute(node);			\
61		if (node->rbaugmented == augmented)			\
62			break;						\
63		node->rbaugmented = augmented;				\
64		rb = rb_parent(&node->rbfield);				\
65	}								\
66}									\
67static inline void							\
68rbname ## _copy(struct rb_node *rb_old, struct rb_node *rb_new)		\
69{									\
70	rbstruct *old = rb_entry(rb_old, rbstruct, rbfield);		\
71	rbstruct *new = rb_entry(rb_new, rbstruct, rbfield);		\
72	new->rbaugmented = old->rbaugmented;				\
73}									\
74static void								\
75rbname ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)	\
76{									\
77	rbstruct *old = rb_entry(rb_old, rbstruct, rbfield);		\
78	rbstruct *new = rb_entry(rb_new, rbstruct, rbfield);		\
79	new->rbaugmented = old->rbaugmented;				\
80	old->rbaugmented = rbcompute(old);				\
81}									\
82rbstatic const struct rb_augment_callbacks rbname = {			\
83	rbname ## _propagate, rbname ## _copy, rbname ## _rotate	\
84};
85
86
87#define	RB_RED		0
88#define	RB_BLACK	1
89
90#define __rb_parent(pc)    ((struct rb_node *)(pc & ~3))
91
92#define __rb_color(pc)     ((pc) & 1)
93#define __rb_is_black(pc)  __rb_color(pc)
94#define __rb_is_red(pc)    (!__rb_color(pc))
95#define rb_color(rb)       __rb_color((rb)->__rb_parent_color)
96#define rb_is_red(rb)      __rb_is_red((rb)->__rb_parent_color)
97#define rb_is_black(rb)    __rb_is_black((rb)->__rb_parent_color)
98
99static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
100{
101	rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
102}
103
104static inline void rb_set_parent_color(struct rb_node *rb,
105				       struct rb_node *p, int color)
106{
107	rb->__rb_parent_color = (unsigned long)p | color;
108}
109
110static inline void
111__rb_change_child(struct rb_node *old, struct rb_node *new,
112		  struct rb_node *parent, struct rb_root *root)
113{
114	if (parent) {
115		if (parent->rb_left == old)
116			parent->rb_left = new;
117		else
118			parent->rb_right = new;
119	} else
120		root->rb_node = new;
121}
122
123extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
124	void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
125
126static __always_inline struct rb_node *
127__rb_erase_augmented(struct rb_node *node, struct rb_root *root,
128		     const struct rb_augment_callbacks *augment)
129{
130	struct rb_node *child = node->rb_right, *tmp = node->rb_left;
131	struct rb_node *parent, *rebalance;
132	unsigned long pc;
133
134	if (!tmp) {
135		/*
136		 * Case 1: node to erase has no more than 1 child (easy!)
137		 *
138		 * Note that if there is one child it must be red due to 5)
139		 * and node must be black due to 4). We adjust colors locally
140		 * so as to bypass __rb_erase_color() later on.
141		 */
142		pc = node->__rb_parent_color;
143		parent = __rb_parent(pc);
144		__rb_change_child(node, child, parent, root);
145		if (child) {
146			child->__rb_parent_color = pc;
147			rebalance = NULL;
148		} else
149			rebalance = __rb_is_black(pc) ? parent : NULL;
150		tmp = parent;
151	} else if (!child) {
152		/* Still case 1, but this time the child is node->rb_left */
153		tmp->__rb_parent_color = pc = node->__rb_parent_color;
154		parent = __rb_parent(pc);
155		__rb_change_child(node, tmp, parent, root);
156		rebalance = NULL;
157		tmp = parent;
158	} else {
159		struct rb_node *successor = child, *child2;
160		tmp = child->rb_left;
161		if (!tmp) {
162			/*
163			 * Case 2: node's successor is its right child
164			 *
165			 *    (n)          (s)
166			 *    / \          / \
167			 *  (x) (s)  ->  (x) (c)
168			 *        \
169			 *        (c)
170			 */
171			parent = successor;
172			child2 = successor->rb_right;
173			augment->copy(node, successor);
174		} else {
175			/*
176			 * Case 3: node's successor is leftmost under
177			 * node's right child subtree
178			 *
179			 *    (n)          (s)
180			 *    / \          / \
181			 *  (x) (y)  ->  (x) (y)
182			 *      /            /
183			 *    (p)          (p)
184			 *    /            /
185			 *  (s)          (c)
186			 *    \
187			 *    (c)
188			 */
189			do {
190				parent = successor;
191				successor = tmp;
192				tmp = tmp->rb_left;
193			} while (tmp);
194			parent->rb_left = child2 = successor->rb_right;
195			successor->rb_right = child;
196			rb_set_parent(child, successor);
197			augment->copy(node, successor);
198			augment->propagate(parent, successor);
199		}
200
201		successor->rb_left = tmp = node->rb_left;
202		rb_set_parent(tmp, successor);
203
204		pc = node->__rb_parent_color;
205		tmp = __rb_parent(pc);
206		__rb_change_child(node, successor, tmp, root);
207		if (child2) {
208			successor->__rb_parent_color = pc;
209			rb_set_parent_color(child2, parent, RB_BLACK);
210			rebalance = NULL;
211		} else {
212			unsigned long pc2 = successor->__rb_parent_color;
213			successor->__rb_parent_color = pc;
214			rebalance = __rb_is_black(pc2) ? parent : NULL;
215		}
216		tmp = successor;
217	}
218
219	augment->propagate(tmp, NULL);
220	return rebalance;
221}
222
223static __always_inline void
224rb_erase_augmented(struct rb_node *node, struct rb_root *root,
225		   const struct rb_augment_callbacks *augment)
226{
227	struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
228	if (rebalance)
229		__rb_erase_color(rebalance, root, augment->rotate);
230}
231
232#endif	/* _LINUX_RBTREE_AUGMENTED_H */
233