From 14b94af0b251a2c80885b60538166fb7d04a642e Mon Sep 17 00:00:00 2001 From: Michel Lespinasse Date: Mon, 8 Oct 2012 16:31:17 -0700 Subject: 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 Acked-by: Rik van Riel Cc: Peter Zijlstra Cc: Andrea Arcangeli Cc: David Woodhouse Signed-off-by: Andrew Morton Signed-off-by: Linus Torvalds --- lib/rbtree.c | 83 ++++++++++++++++++++++++++++++++++++++++++++++++++----- lib/rbtree_test.c | 58 ++++++++++++++++++++++++++++---------- 2 files changed, 120 insertions(+), 21 deletions(-) (limited to 'lib') 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, __rb_change_child(old, new, parent, root); } -void rb_insert_color(struct rb_node *node, struct rb_root *root) +static __always_inline void +__rb_insert(struct rb_node *node, struct rb_root *root, + void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) { struct rb_node *parent = rb_red_parent(node), *gparent, *tmp; @@ -169,6 +171,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root) rb_set_parent_color(tmp, parent, RB_BLACK); rb_set_parent_color(parent, node, RB_RED); + augment_rotate(parent, node); parent = node; tmp = node->rb_right; } @@ -187,6 +190,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root) if (tmp) rb_set_parent_color(tmp, gparent, RB_BLACK); __rb_rotate_set_parents(gparent, parent, root, RB_RED); + augment_rotate(gparent, parent); break; } else { tmp = gparent->rb_left; @@ -209,6 +213,7 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root) rb_set_parent_color(tmp, parent, RB_BLACK); rb_set_parent_color(parent, node, RB_RED); + augment_rotate(parent, node); parent = node; tmp = node->rb_left; } @@ -219,13 +224,15 @@ void rb_insert_color(struct rb_node *node, struct rb_root *root) if (tmp) rb_set_parent_color(tmp, gparent, RB_BLACK); __rb_rotate_set_parents(gparent, parent, root, RB_RED); + augment_rotate(gparent, parent); break; } } } -EXPORT_SYMBOL(rb_insert_color); -static void __rb_erase_color(struct rb_node *parent, struct rb_root *root) +static __always_inline void +__rb_erase_color(struct rb_node *parent, struct rb_root *root, + const struct rb_augment_callbacks *augment) { struct rb_node *node = NULL, *sibling, *tmp1, *tmp2; @@ -254,6 +261,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root) rb_set_parent_color(tmp1, parent, RB_BLACK); __rb_rotate_set_parents(parent, sibling, root, RB_RED); + augment->rotate(parent, sibling); sibling = tmp1; } tmp1 = sibling->rb_right; @@ -305,6 +313,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root) if (tmp1) rb_set_parent_color(tmp1, sibling, RB_BLACK); + augment->rotate(sibling, tmp2); tmp1 = sibling; sibling = tmp2; } @@ -327,6 +336,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root) rb_set_parent(tmp2, parent); __rb_rotate_set_parents(parent, sibling, root, RB_BLACK); + augment->rotate(parent, sibling); break; } else { sibling = parent->rb_left; @@ -337,6 +347,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root) rb_set_parent_color(tmp1, parent, RB_BLACK); __rb_rotate_set_parents(parent, sibling, root, RB_RED); + augment->rotate(parent, sibling); sibling = tmp1; } tmp1 = sibling->rb_left; @@ -363,6 +374,7 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root) if (tmp1) rb_set_parent_color(tmp1, sibling, RB_BLACK); + augment->rotate(sibling, tmp2); tmp1 = sibling; sibling = tmp2; } @@ -374,12 +386,15 @@ static void __rb_erase_color(struct rb_node *parent, struct rb_root *root) rb_set_parent(tmp2, parent); __rb_rotate_set_parents(parent, sibling, root, RB_BLACK); + augment->rotate(parent, sibling); break; } } } -void rb_erase(struct rb_node *node, struct rb_root *root) +static __always_inline void +__rb_erase(struct rb_node *node, struct rb_root *root, + const struct rb_augment_callbacks *augment) { struct rb_node *child = node->rb_right, *tmp = node->rb_left; struct rb_node *parent, *rebalance; @@ -401,12 +416,14 @@ void rb_erase(struct rb_node *node, struct rb_root *root) rebalance = NULL; } else rebalance = __rb_is_black(pc) ? parent : NULL; + tmp = parent; } else if (!child) { /* Still case 1, but this time the child is node->rb_left */ tmp->__rb_parent_color = pc = node->__rb_parent_color; parent = __rb_parent(pc); __rb_change_child(node, tmp, parent, root); rebalance = NULL; + tmp = parent; } else { struct rb_node *successor = child, *child2; tmp = child->rb_left; @@ -420,8 +437,9 @@ void rb_erase(struct rb_node *node, struct rb_root *root) * \ * (c) */ - parent = child; - child2 = child->rb_right; + parent = successor; + child2 = successor->rb_right; + augment->copy(node, successor); } else { /* * Case 3: node's successor is leftmost under @@ -445,6 +463,8 @@ void rb_erase(struct rb_node *node, struct rb_root *root) parent->rb_left = child2 = successor->rb_right; successor->rb_right = child; rb_set_parent(child, successor); + augment->copy(node, successor); + augment->propagate(parent, successor); } successor->rb_left = tmp = node->rb_left; @@ -462,13 +482,62 @@ void rb_erase(struct rb_node *node, struct rb_root *root) successor->__rb_parent_color = pc; rebalance = __rb_is_black(pc2) ? parent : NULL; } + tmp = successor; } + augment->propagate(tmp, NULL); if (rebalance) - __rb_erase_color(rebalance, root); + __rb_erase_color(rebalance, root, augment); +} + +/* + * Non-augmented rbtree manipulation functions. + * + * We use dummy augmented callbacks here, and have the compiler optimize them + * out of the rb_insert_color() and rb_erase() function definitions. + */ + +static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {} +static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {} +static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {} + +static const struct rb_augment_callbacks dummy_callbacks = { + dummy_propagate, dummy_copy, dummy_rotate +}; + +void rb_insert_color(struct rb_node *node, struct rb_root *root) +{ + __rb_insert(node, root, dummy_rotate); +} +EXPORT_SYMBOL(rb_insert_color); + +void rb_erase(struct rb_node *node, struct rb_root *root) +{ + __rb_erase(node, root, &dummy_callbacks); } EXPORT_SYMBOL(rb_erase); +/* + * Augmented rbtree manipulation functions. + * + * This instantiates the same __always_inline functions as in the non-augmented + * case, but this time with user-defined callbacks. + */ + +void __rb_insert_augmented(struct rb_node *node, struct rb_root *root, + void (*augment_rotate)(struct rb_node *old, struct rb_node *new)) +{ + __rb_insert(node, root, augment_rotate); +} +EXPORT_SYMBOL(__rb_insert_augmented); + +void rb_erase_augmented(struct rb_node *node, struct rb_root *root, + const struct rb_augment_callbacks *augment) +{ + __rb_erase(node, root, augment); +} +EXPORT_SYMBOL(rb_erase_augmented); + static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data) { 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) return max; } -static void augment_callback(struct rb_node *rb, void *unused) +static void augment_propagate(struct rb_node *rb, struct rb_node *stop) { - struct test_node *node = rb_entry(rb, struct test_node, rb); - node->augmented = augment_recompute(node); + while (rb != stop) { + struct test_node *node = rb_entry(rb, struct test_node, rb); + u32 augmented = augment_recompute(node); + if (node->augmented == augmented) + break; + node->augmented = augmented; + rb = rb_parent(&node->rb); + } +} + +static void augment_copy(struct rb_node *rb_old, struct rb_node *rb_new) +{ + struct test_node *old = rb_entry(rb_old, struct test_node, rb); + struct test_node *new = rb_entry(rb_new, struct test_node, rb); + new->augmented = old->augmented; } +static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new) +{ + struct test_node *old = rb_entry(rb_old, struct test_node, rb); + struct test_node *new = rb_entry(rb_new, struct test_node, rb); + + /* Rotation doesn't change subtree's augmented value */ + new->augmented = old->augmented; + old->augmented = augment_recompute(old); +} + +static const struct rb_augment_callbacks augment_callbacks = { + augment_propagate, augment_copy, augment_rotate +}; + static void insert_augmented(struct test_node *node, struct rb_root *root) { - struct rb_node **new = &root->rb_node, *parent = NULL; + struct rb_node **new = &root->rb_node, *rb_parent = NULL; u32 key = node->key; + u32 val = node->val; + struct test_node *parent; while (*new) { - parent = *new; - if (key < rb_entry(parent, struct test_node, rb)->key) - new = &parent->rb_left; + rb_parent = *new; + parent = rb_entry(rb_parent, struct test_node, rb); + if (parent->augmented < val) + parent->augmented = val; + if (key < parent->key) + new = &parent->rb.rb_left; else - new = &parent->rb_right; + new = &parent->rb.rb_right; } - rb_link_node(&node->rb, parent, new); - rb_insert_color(&node->rb, root); - rb_augment_insert(&node->rb, augment_callback, NULL); + node->augmented = val; + rb_link_node(&node->rb, rb_parent, new); + rb_insert_augmented(&node->rb, root, &augment_callbacks); } static void erase_augmented(struct test_node *node, struct rb_root *root) { - struct rb_node *deepest = rb_augment_erase_begin(&node->rb); - rb_erase(&node->rb, root); - rb_augment_erase_end(deepest, augment_callback, NULL); + rb_erase_augmented(&node->rb, root, &augment_callbacks); } static void init(void) -- cgit v1.2.3