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