+++ b/lib/rbtree.c
@@ -88,7 +88,8 @@ __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
root->rb_node = new;
}
-void rb_insert_color(struct rb_node *node, struct rb_root *root)
+inline void rb_insert_augmented(struct rb_node *node, struct rb_root *root,
+ rb_augment_rotate *augment)
{
struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
@@ -152,6 +153,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(parent, node);
+static inline void dummy(struct rb_node *old, struct rb_node *new) {}
+
+void rb_insert_color(struct rb_node *node, struct rb_root *root) {
+ rb_insert_augmented(node, root, dummy);
+}
EXPORT_SYMBOL(rb_insert_color);
static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c
index 2dfafe4..5ace332 100644
--- a/lib/rbtree_test.c
+++ b/lib/rbtree_test.c
@@ -67,22 +67,37 @@ static void augment_callback(struct rb_node *rb, void *unused)
node->augmented = augment_recompute(node);
}
+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);
+}