aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorMichel Lespinasse <walken@google.com>2012-10-08 19:31:17 -0400
committerLinus Torvalds <torvalds@linux-foundation.org>2012-10-09 03:22:37 -0400
commit14b94af0b251a2c80885b60538166fb7d04a642e (patch)
treeef447d340435c441f8c3e54eb8f26f747aa73108
parentdadf93534f125b9eda486b471446a8456a603d27 (diff)
rbtree: faster augmented rbtree manipulation
Introduce new augmented rbtree APIs that allow minimal recalculation of augmented node information. A new callback is added to the rbtree insertion and erase rebalancing functions, to be called on each tree rotations. Such rotations preserve the subtree's root augmented value, but require recalculation of the one child that was previously located at the subtree root. In the insertion case, the handcoded search phase must be updated to maintain the augmented information on insertion, and then the rbtree coloring/rebalancing algorithms keep it up to date. In the erase case, things are more complicated since it is library code that manipulates the rbtree in order to remove internal nodes. This requires a couple additional callbacks to copy a subtree's augmented value when a new root is stitched in, and to recompute augmented values down the ancestry path when a node is removed from the tree. In order to preserve maximum speed for the non-augmented case, we provide two versions of each tree manipulation function. rb_insert_augmented() is the augmented equivalent of rb_insert_color(), and rb_erase_augmented() is the augmented equivalent of rb_erase(). Signed-off-by: Michel Lespinasse <walken@google.com> Acked-by: Rik van Riel <riel@redhat.com> Cc: Peter Zijlstra <a.p.zijlstra@chello.nl> Cc: Andrea Arcangeli <aarcange@redhat.com> Cc: David Woodhouse <dwmw2@infradead.org> Signed-off-by: Andrew Morton <akpm@linux-foundation.org> Signed-off-by: Linus Torvalds <torvalds@linux-foundation.org>
-rw-r--r--Documentation/rbtree.txt190
-rw-r--r--include/linux/rbtree.h19
-rw-r--r--lib/rbtree.c83
-rw-r--r--lib/rbtree_test.c58
4 files changed, 296 insertions, 54 deletions
diff --git a/Documentation/rbtree.txt b/Documentation/rbtree.txt
index 8d32d85a5234..0a0b6dce3e08 100644
--- a/Documentation/rbtree.txt
+++ b/Documentation/rbtree.txt
@@ -193,24 +193,42 @@ Example:
193Support for Augmented rbtrees 193Support for Augmented rbtrees
194----------------------------- 194-----------------------------
195 195
196Augmented rbtree is an rbtree with "some" additional data stored in each node. 196Augmented rbtree is an rbtree with "some" additional data stored in
197This data can be used to augment some new functionality to rbtree. 197each node, where the additional data for node N must be a function of
198Augmented rbtree is an optional feature built on top of basic rbtree 198the contents of all nodes in the subtree rooted at N. This data can
199infrastructure. An rbtree user who wants this feature will have to call the 199be used to augment some new functionality to rbtree. Augmented rbtree
200augmentation functions with the user provided augmentation callback 200is an optional feature built on top of basic rbtree infrastructure.
201when inserting and erasing nodes. 201An rbtree user who wants this feature will have to call the augmentation
202functions with the user provided augmentation callback when inserting
203and erasing nodes.
202 204
203On insertion, the user must call rb_augment_insert() once the new node is in 205On insertion, the user must update the augmented information on the path
204place. This will cause the augmentation function callback to be called for 206leading to the inserted node, then call rb_link_node() as usual and
205each node between the new node and the root which has been affected by the 207rb_augment_inserted() instead of the usual rb_insert_color() call.
206insertion. 208If rb_augment_inserted() rebalances the rbtree, it will callback into
209a user provided function to update the augmented information on the
210affected subtrees.
207 211
208When erasing a node, the user must call rb_augment_erase_begin() first to 212When erasing a node, the user must call rb_erase_augmented() instead of
209retrieve the deepest node on the rebalance path. Then, after erasing the 213rb_erase(). rb_erase_augmented() calls back into user provided functions
210original node, the user must call rb_augment_erase_end() with the deepest 214to updated the augmented information on affected subtrees.
211node found earlier. This will cause the augmentation function to be called
212for each affected node between the deepest node and the root.
213 215
216In both cases, the callbacks are provided through struct rb_augment_callbacks.
2173 callbacks must be defined:
218
219- A propagation callback, which updates the augmented value for a given
220 node and its ancestors, up to a given stop point (or NULL to update
221 all the way to the root).
222
223- A copy callback, which copies the augmented value for a given subtree
224 to a newly assigned subtree root.
225
226- A tree rotation callback, which copies the augmented value for a given
227 subtree to a newly assigned subtree root AND recomputes the augmented
228 information for the former subtree root.
229
230
231Sample usage:
214 232
215Interval tree is an example of augmented rb tree. Reference - 233Interval tree is an example of augmented rb tree. Reference -
216"Introduction to Algorithms" by Cormen, Leiserson, Rivest and Stein. 234"Introduction to Algorithms" by Cormen, Leiserson, Rivest and Stein.
@@ -230,26 +248,132 @@ and its immediate children. And this will be used in O(log n) lookup
230for lowest match (lowest start address among all possible matches) 248for lowest match (lowest start address among all possible matches)
231with something like: 249with something like:
232 250
233find_lowest_match(lo, hi, node) 251struct interval_tree_node *
252interval_tree_first_match(struct rb_root *root,
253 unsigned long start, unsigned long last)
234{ 254{
235 lowest_match = NULL; 255 struct interval_tree_node *node;
236 while (node) { 256
237 if (max_hi(node->left) > lo) { 257 if (!root->rb_node)
238 // Lowest overlap if any must be on left side 258 return NULL;
239 node = node->left; 259 node = rb_entry(root->rb_node, struct interval_tree_node, rb);
240 } else if (overlap(lo, hi, node)) { 260
241 lowest_match = node; 261 while (true) {
242 break; 262 if (node->rb.rb_left) {
243 } else if (lo > node->lo) { 263 struct interval_tree_node *left =
244 // Lowest overlap if any must be on right side 264 rb_entry(node->rb.rb_left,
245 node = node->right; 265 struct interval_tree_node, rb);
246 } else { 266 if (left->__subtree_last >= start) {
247 break; 267 /*
268 * Some nodes in left subtree satisfy Cond2.
269 * Iterate to find the leftmost such node N.
270 * If it also satisfies Cond1, that's the match
271 * we are looking for. Otherwise, there is no
272 * matching interval as nodes to the right of N
273 * can't satisfy Cond1 either.
274 */
275 node = left;
276 continue;
277 }
248 } 278 }
279 if (node->start <= last) { /* Cond1 */
280 if (node->last >= start) /* Cond2 */
281 return node; /* node is leftmost match */
282 if (node->rb.rb_right) {
283 node = rb_entry(node->rb.rb_right,
284 struct interval_tree_node, rb);
285 if (node->__subtree_last >= start)
286 continue;
287 }
288 }
289 return NULL; /* No match */
290 }
291}
292
293Insertion/removal are defined using the following augmented callbacks:
294
295static inline unsigned long
296compute_subtree_last(struct interval_tree_node *node)
297{
298 unsigned long max = node->last, subtree_last;
299 if (node->rb.rb_left) {
300 subtree_last = rb_entry(node->rb.rb_left,
301 struct interval_tree_node, rb)->__subtree_last;
302 if (max < subtree_last)
303 max = subtree_last;
304 }
305 if (node->rb.rb_right) {
306 subtree_last = rb_entry(node->rb.rb_right,
307 struct interval_tree_node, rb)->__subtree_last;
308 if (max < subtree_last)
309 max = subtree_last;
310 }
311 return max;
312}
313
314static void augment_propagate(struct rb_node *rb, struct rb_node *stop)
315{
316 while (rb != stop) {
317 struct interval_tree_node *node =
318 rb_entry(rb, struct interval_tree_node, rb);
319 unsigned long subtree_last = compute_subtree_last(node);
320 if (node->__subtree_last == subtree_last)
321 break;
322 node->__subtree_last = subtree_last;
323 rb = rb_parent(&node->rb);
324 }
325}
326
327static void augment_copy(struct rb_node *rb_old, struct rb_node *rb_new)
328{
329 struct interval_tree_node *old =
330 rb_entry(rb_old, struct interval_tree_node, rb);
331 struct interval_tree_node *new =
332 rb_entry(rb_new, struct interval_tree_node, rb);
333
334 new->__subtree_last = old->__subtree_last;
335}
336
337static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
338{
339 struct interval_tree_node *old =
340 rb_entry(rb_old, struct interval_tree_node, rb);
341 struct interval_tree_node *new =
342 rb_entry(rb_new, struct interval_tree_node, rb);
343
344 new->__subtree_last = old->__subtree_last;
345 old->__subtree_last = compute_subtree_last(old);
346}
347
348static const struct rb_augment_callbacks augment_callbacks = {
349 augment_propagate, augment_copy, augment_rotate
350};
351
352void interval_tree_insert(struct interval_tree_node *node,
353 struct rb_root *root)
354{
355 struct rb_node **link = &root->rb_node, *rb_parent = NULL;
356 unsigned long start = node->start, last = node->last;
357 struct interval_tree_node *parent;
358
359 while (*link) {
360 rb_parent = *link;
361 parent = rb_entry(rb_parent, struct interval_tree_node, rb);
362 if (parent->__subtree_last < last)
363 parent->__subtree_last = last;
364 if (start < parent->start)
365 link = &parent->rb.rb_left;
366 else
367 link = &parent->rb.rb_right;
249 } 368 }
250 return lowest_match; 369
370 node->__subtree_last = last;
371 rb_link_node(&node->rb, rb_parent, link);
372 rb_insert_augmented(&node->rb, root, &augment_callbacks);
251} 373}
252 374
253Finding exact match will be to first find lowest match and then to follow 375void interval_tree_remove(struct interval_tree_node *node,
254successor nodes looking for exact match, until the start of a node is beyond 376 struct rb_root *root)
255the hi value we are looking for. 377{
378 rb_erase_augmented(&node->rb, root, &augment_callbacks);
379}
diff --git a/include/linux/rbtree.h b/include/linux/rbtree.h
index bf836a2c6a32..c902eb9d6506 100644
--- a/include/linux/rbtree.h
+++ b/include/linux/rbtree.h
@@ -61,6 +61,25 @@ struct rb_root {
61extern void rb_insert_color(struct rb_node *, struct rb_root *); 61extern void rb_insert_color(struct rb_node *, struct rb_root *);
62extern void rb_erase(struct rb_node *, struct rb_root *); 62extern void rb_erase(struct rb_node *, struct rb_root *);
63 63
64
65struct rb_augment_callbacks {
66 void (*propagate)(struct rb_node *node, struct rb_node *stop);
67 void (*copy)(struct rb_node *old, struct rb_node *new);
68 void (*rotate)(struct rb_node *old, struct rb_node *new);
69};
70
71extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
72 void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
73extern void rb_erase_augmented(struct rb_node *node, struct rb_root *root,
74 const struct rb_augment_callbacks *augment);
75static inline void
76rb_insert_augmented(struct rb_node *node, struct rb_root *root,
77 const struct rb_augment_callbacks *augment)
78{
79 __rb_insert_augmented(node, root, augment->rotate);
80}
81
82
64typedef void (*rb_augment_f)(struct rb_node *node, void *data); 83typedef void (*rb_augment_f)(struct rb_node *node, void *data);
65 84
66extern void rb_augment_insert(struct rb_node *node, 85extern void rb_augment_insert(struct rb_node *node,
diff --git a/lib/rbtree.c b/lib/rbtree.c
index 938061ecbe61..a37ee7954b8f 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -105,7 +105,9 @@ __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
105 __rb_change_child(old, new, parent, root); 105 __rb_change_child(old, new, parent, root);
106} 106}
107 107
108void rb_insert_color(struct rb_node *node, struct rb_root *root) 108static __always_inline void
109__rb_insert(struct rb_node *node, struct rb_root *root,
110 void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
109{ 111{
110 struct rb_node *parent = rb_red_parent(node), *gparent, *tmp; 112 struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
111 113
@@ -169,6 +171,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
169 rb_set_parent_color(tmp, parent, 171 rb_set_parent_color(tmp, parent,
170 RB_BLACK); 172 RB_BLACK);
171 rb_set_parent_color(parent, node, RB_RED); 173 rb_set_parent_color(parent, node, RB_RED);
174 augment_rotate(parent, node);
172 parent = node; 175 parent = node;
173 tmp = node->rb_right; 176 tmp = node->rb_right;
174 } 177 }
@@ -187,6 +190,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
187 if (tmp) 190 if (tmp)
188 rb_set_parent_color(tmp, gparent, RB_BLACK); 191 rb_set_parent_color(tmp, gparent, RB_BLACK);
189 __rb_rotate_set_parents(gparent, parent, root, RB_RED); 192 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
193 augment_rotate(gparent, parent);
190 break; 194 break;
191 } else { 195 } else {
192 tmp = gparent->rb_left; 196 tmp = gparent->rb_left;
@@ -209,6 +213,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
209 rb_set_parent_color(tmp, parent, 213 rb_set_parent_color(tmp, parent,
210 RB_BLACK); 214 RB_BLACK);
211 rb_set_parent_color(parent, node, RB_RED); 215 rb_set_parent_color(parent, node, RB_RED);
216 augment_rotate(parent, node);
212 parent = node; 217 parent = node;
213 tmp = node->rb_left; 218 tmp = node->rb_left;
214 } 219 }
@@ -219,13 +224,15 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root)
219 if (tmp) 224 if (tmp)
220 rb_set_parent_color(tmp, gparent, RB_BLACK); 225 rb_set_parent_color(tmp, gparent, RB_BLACK);
221 __rb_rotate_set_parents(gparent, parent, root, RB_RED); 226 __rb_rotate_set_parents(gparent, parent, root, RB_RED);
227 augment_rotate(gparent, parent);
222 break; 228 break;
223 } 229 }
224 } 230 }
225} 231}
226EXPORT_SYMBOL(rb_insert_color);
227 232
228static void __rb_erase_color(struct rb_node *parent, struct rb_root *root) 233static __always_inline void
234__rb_erase_color(struct rb_node *parent, struct rb_root *root,
235 const struct rb_augment_callbacks *augment)
229{ 236{
230 struct rb_node *node = NULL, *sibling, *tmp1, *tmp2; 237 struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
231 238
@@ -254,6 +261,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root)
254 rb_set_parent_color(tmp1, parent, RB_BLACK); 261 rb_set_parent_color(tmp1, parent, RB_BLACK);
255 __rb_rotate_set_parents(parent, sibling, root, 262 __rb_rotate_set_parents(parent, sibling, root,
256 RB_RED); 263 RB_RED);
264 augment->rotate(parent, sibling);
257 sibling = tmp1; 265 sibling = tmp1;
258 } 266 }
259 tmp1 = sibling->rb_right; 267 tmp1 = sibling->rb_right;
@@ -305,6 +313,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root)
305 if (tmp1) 313 if (tmp1)
306 rb_set_parent_color(tmp1, sibling, 314 rb_set_parent_color(tmp1, sibling,
307 RB_BLACK); 315 RB_BLACK);
316 augment->rotate(sibling, tmp2);
308 tmp1 = sibling; 317 tmp1 = sibling;
309 sibling = tmp2; 318 sibling = tmp2;
310 } 319 }
@@ -327,6 +336,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root)
327 rb_set_parent(tmp2, parent); 336 rb_set_parent(tmp2, parent);
328 __rb_rotate_set_parents(parent, sibling, root, 337 __rb_rotate_set_parents(parent, sibling, root,
329 RB_BLACK); 338 RB_BLACK);
339 augment->rotate(parent, sibling);
330 break; 340 break;
331 } else { 341 } else {
332 sibling = parent->rb_left; 342 sibling = parent->rb_left;
@@ -337,6 +347,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root)
337 rb_set_parent_color(tmp1, parent, RB_BLACK); 347 rb_set_parent_color(tmp1, parent, RB_BLACK);
338 __rb_rotate_set_parents(parent, sibling, root, 348 __rb_rotate_set_parents(parent, sibling, root,
339 RB_RED); 349 RB_RED);
350 augment->rotate(parent, sibling);
340 sibling = tmp1; 351 sibling = tmp1;
341 } 352 }
342 tmp1 = sibling->rb_left; 353 tmp1 = sibling->rb_left;
@@ -363,6 +374,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root)
363 if (tmp1) 374 if (tmp1)
364 rb_set_parent_color(tmp1, sibling, 375 rb_set_parent_color(tmp1, sibling,
365 RB_BLACK); 376 RB_BLACK);
377 augment->rotate(sibling, tmp2);
366 tmp1 = sibling; 378 tmp1 = sibling;
367 sibling = tmp2; 379 sibling = tmp2;
368 } 380 }
@@ -374,12 +386,15 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root)
374 rb_set_parent(tmp2, parent); 386 rb_set_parent(tmp2, parent);
375 __rb_rotate_set_parents(parent, sibling, root, 387 __rb_rotate_set_parents(parent, sibling, root,
376 RB_BLACK); 388 RB_BLACK);
389 augment->rotate(parent, sibling);
377 break; 390 break;
378 } 391 }
379 } 392 }
380} 393}
381 394
382void rb_erase(struct rb_node *node, struct rb_root *root) 395static __always_inline void
396__rb_erase(struct rb_node *node, struct rb_root *root,
397 const struct rb_augment_callbacks *augment)
383{ 398{
384 struct rb_node *child = node->rb_right, *tmp = node->rb_left; 399 struct rb_node *child = node->rb_right, *tmp = node->rb_left;
385 struct rb_node *parent, *rebalance; 400 struct rb_node *parent, *rebalance;
@@ -401,12 +416,14 @@ void rb_erase(struct rb_node *node, struct rb_root *root)
401 rebalance = NULL; 416 rebalance = NULL;
402 } else 417 } else
403 rebalance = __rb_is_black(pc) ? parent : NULL; 418 rebalance = __rb_is_black(pc) ? parent : NULL;
419 tmp = parent;
404 } else if (!child) { 420 } else if (!child) {
405 /* Still case 1, but this time the child is node->rb_left */ 421 /* Still case 1, but this time the child is node->rb_left */
406 tmp->__rb_parent_color = pc = node->__rb_parent_color; 422 tmp->__rb_parent_color = pc = node->__rb_parent_color;
407 parent = __rb_parent(pc); 423 parent = __rb_parent(pc);
408 __rb_change_child(node, tmp, parent, root); 424 __rb_change_child(node, tmp, parent, root);
409 rebalance = NULL; 425 rebalance = NULL;
426 tmp = parent;
410 } else { 427 } else {
411 struct rb_node *successor = child, *child2; 428 struct rb_node *successor = child, *child2;
412 tmp = child->rb_left; 429 tmp = child->rb_left;
@@ -420,8 +437,9 @@ void rb_erase(struct rb_node *node, struct rb_root *root)
420 * \ 437 * \
421 * (c) 438 * (c)
422 */ 439 */
423 parent = child; 440 parent = successor;
424 child2 = child->rb_right; 441 child2 = successor->rb_right;
442 augment->copy(node, successor);
425 } else { 443 } else {
426 /* 444 /*
427 * Case 3: node's successor is leftmost under 445 * Case 3: node's successor is leftmost under
@@ -445,6 +463,8 @@ void rb_erase(struct rb_node *node, struct rb_root *root)
445 parent->rb_left = child2 = successor->rb_right; 463 parent->rb_left = child2 = successor->rb_right;
446 successor->rb_right = child; 464 successor->rb_right = child;
447 rb_set_parent(child, successor); 465 rb_set_parent(child, successor);
466 augment->copy(node, successor);
467 augment->propagate(parent, successor);
448 } 468 }
449 469
450 successor->rb_left = tmp = node->rb_left; 470 successor->rb_left = tmp = node->rb_left;
@@ -462,13 +482,62 @@ void rb_erase(struct rb_node *node, struct rb_root *root)
462 successor->__rb_parent_color = pc; 482 successor->__rb_parent_color = pc;
463 rebalance = __rb_is_black(pc2) ? parent : NULL; 483 rebalance = __rb_is_black(pc2) ? parent : NULL;
464 } 484 }
485 tmp = successor;
465 } 486 }
466 487
488 augment->propagate(tmp, NULL);
467 if (rebalance) 489 if (rebalance)
468 __rb_erase_color(rebalance, root); 490 __rb_erase_color(rebalance, root, augment);
491}
492
493/*
494 * Non-augmented rbtree manipulation functions.
495 *
496 * We use dummy augmented callbacks here, and have the compiler optimize them
497 * out of the rb_insert_color() and rb_erase() function definitions.
498 */
499
500static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
501static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
502static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
503
504static const struct rb_augment_callbacks dummy_callbacks = {
505 dummy_propagate, dummy_copy, dummy_rotate
506};
507
508void rb_insert_color(struct rb_node *node, struct rb_root *root)
509{
510 __rb_insert(node, root, dummy_rotate);
511}
512EXPORT_SYMBOL(rb_insert_color);
513
514void rb_erase(struct rb_node *node, struct rb_root *root)
515{
516 __rb_erase(node, root, &dummy_callbacks);
469} 517}
470EXPORT_SYMBOL(rb_erase); 518EXPORT_SYMBOL(rb_erase);
471 519
520/*
521 * Augmented rbtree manipulation functions.
522 *
523 * This instantiates the same __always_inline functions as in the non-augmented
524 * case, but this time with user-defined callbacks.
525 */
526
527void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
528 void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
529{
530 __rb_insert(node, root, augment_rotate);
531}
532EXPORT_SYMBOL(__rb_insert_augmented);
533
534void rb_erase_augmented(struct rb_node *node, struct rb_root *root,
535 const struct rb_augment_callbacks *augment)
536{
537 __rb_erase(node, root, augment);
538}
539EXPORT_SYMBOL(rb_erase_augmented);
540
472static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data) 541static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
473{ 542{
474 struct rb_node *parent; 543 struct rb_node *parent;
diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c
index 66e41d4bfc39..e28345df09bf 100644
--- a/lib/rbtree_test.c
+++ b/lib/rbtree_test.c
@@ -61,35 +61,65 @@ static inline u32 augment_recompute(struct test_node *node)
61 return max; 61 return max;
62} 62}
63 63
64static void augment_callback(struct rb_node *rb, void *unused) 64static void augment_propagate(struct rb_node *rb, struct rb_node *stop)
65{ 65{
66 struct test_node *node = rb_entry(rb, struct test_node, rb); 66 while (rb != stop) {
67 node->augmented = augment_recompute(node); 67 struct test_node *node = rb_entry(rb, struct test_node, rb);
68 u32 augmented = augment_recompute(node);
69 if (node->augmented == augmented)
70 break;
71 node->augmented = augmented;
72 rb = rb_parent(&node->rb);
73 }
74}
75
76static void augment_copy(struct rb_node *rb_old, struct rb_node *rb_new)
77{
78 struct test_node *old = rb_entry(rb_old, struct test_node, rb);
79 struct test_node *new = rb_entry(rb_new, struct test_node, rb);
80 new->augmented = old->augmented;
68} 81}
69 82
83static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
84{
85 struct test_node *old = rb_entry(rb_old, struct test_node, rb);
86 struct test_node *new = rb_entry(rb_new, struct test_node, rb);
87
88 /* Rotation doesn't change subtree's augmented value */
89 new->augmented = old->augmented;
90 old->augmented = augment_recompute(old);
91}
92
93static const struct rb_augment_callbacks augment_callbacks = {
94 augment_propagate, augment_copy, augment_rotate
95};
96
70static void insert_augmented(struct test_node *node, struct rb_root *root) 97static void insert_augmented(struct test_node *node, struct rb_root *root)
71{ 98{
72 struct rb_node **new = &root->rb_node, *parent = NULL; 99 struct rb_node **new = &root->rb_node, *rb_parent = NULL;
73 u32 key = node->key; 100 u32 key = node->key;
101 u32 val = node->val;
102 struct test_node *parent;
74 103
75 while (*new) { 104 while (*new) {
76 parent = *new; 105 rb_parent = *new;
77 if (key < rb_entry(parent, struct test_node, rb)->key) 106 parent = rb_entry(rb_parent, struct test_node, rb);
78 new = &parent->rb_left; 107 if (parent->augmented < val)
108 parent->augmented = val;
109 if (key < parent->key)
110 new = &parent->rb.rb_left;
79 else 111 else
80 new = &parent->rb_right; 112 new = &parent->rb.rb_right;
81 } 113 }
82 114
83 rb_link_node(&node->rb, parent, new); 115 node->augmented = val;
84 rb_insert_color(&node->rb, root); 116 rb_link_node(&node->rb, rb_parent, new);
85 rb_augment_insert(&node->rb, augment_callback, NULL); 117 rb_insert_augmented(&node->rb, root, &augment_callbacks);
86} 118}
87 119
88static void erase_augmented(struct test_node *node, struct rb_root *root) 120static void erase_augmented(struct test_node *node, struct rb_root *root)
89{ 121{
90 struct rb_node *deepest = rb_augment_erase_begin(&node->rb); 122 rb_erase_augmented(&node->rb, root, &augment_callbacks);
91 rb_erase(&node->rb, root);
92 rb_augment_erase_end(deepest, augment_callback, NULL);
93} 123}
94 124
95static void init(void) 125static void init(void)