insert() and insert_cached() open-code exactly what the generic rb_add() and rb_add_cached() helpers do. Use the helpers instead. Besides removing duplicated code, this makes rbtree_test cover the helpers themselves, which is what most in-tree rbtree users call nowadays. Signed-off-by: Yiwei Lin Assisted-by: LLM --- lib/rbtree_test.c | 37 ++++++++----------------------------- 1 file changed, 8 insertions(+), 29 deletions(-) diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c index 768c5e6453f37..c1386d328cfb1 100644 --- a/lib/rbtree_test.c +++ b/lib/rbtree_test.c @@ -30,41 +30,20 @@ static struct test_node *nodes = NULL; static struct rnd_state rnd; -static void insert(struct test_node *node, struct rb_root_cached *root) +static inline bool less(struct rb_node *a, const struct rb_node *b) { - struct rb_node **new = &root->rb_root.rb_node, *parent = NULL; - u32 key = node->key; - - while (*new) { - parent = *new; - if (key < rb_entry(parent, struct test_node, rb)->key) - new = &parent->rb_left; - else - new = &parent->rb_right; - } + return rb_entry(a, struct test_node, rb)->key < + rb_entry(b, struct test_node, rb)->key; +} - rb_link_node(&node->rb, parent, new); - rb_insert_color(&node->rb, &root->rb_root); +static void insert(struct test_node *node, struct rb_root_cached *root) +{ + rb_add(&node->rb, &root->rb_root, less); } static void insert_cached(struct test_node *node, struct rb_root_cached *root) { - struct rb_node **new = &root->rb_root.rb_node, *parent = NULL; - u32 key = node->key; - bool leftmost = true; - - while (*new) { - parent = *new; - if (key < rb_entry(parent, struct test_node, rb)->key) - new = &parent->rb_left; - else { - new = &parent->rb_right; - leftmost = false; - } - } - - rb_link_node(&node->rb, parent, new); - rb_insert_color_cached(&node->rb, root, leftmost); + rb_add_cached(&node->rb, root, less); } static inline void erase(struct test_node *node, struct rb_root_cached *root) -- 2.34.1 From: "Peter Zijlstra (Intel)" RB_DECLARE_CALLBACKS_MULTI() asks its user for a function that copies the augmented fields and one that recomputes them from the children, with the early-exit protocol of ->propagate() hand-coded in the latter. sched/eevdf, its only user, needs five helpers for three fields. Describe each augmented field instead: RB_AUG(val, aug, fold) names the per-node value, the member holding the subtree aggregate and how two aggregates combine (min, max, a sum, ...). RB_AUG_FUNC() takes a function for the per-node value. The template then generates, per field, a recompute that works on a local and stores once, and a copy, and combines them into the callbacks; the early exit is the AND of the per-field results. RB_DECLARE_CALLBACKS() is the only template left. RB_DECLARE_CALLBACKS_MAX() becomes RB_AUG_FUNC(RBVALUE, RBAUGMENTED, max) in a wrapper, with its RBCOMPUTE argument renamed to RBVALUE, since it returns the per-node scalar and is a different thing from the RBCOMPUTE of the generic template it used to build on. RB_DECLARE_CALLBACKS_MULTI() goes away. sched/eevdf shrinks to a wrapping-safe min() for min_vruntime and three RB_AUG() lines. net/tipc's service range tree uses RB_AUG() directly. No functional change intended. Signed-off-by: Peter Zijlstra (Intel) Link: https://lore.kernel.org/r/20260928213736.GA2947991@noisy.programming.kicks-ass.net [yiwei: take a fold(a, b) instead of a "replace?" compare so that sums and counts can be expressed too, which also lets min()/max() replace RB_MIN()/RB_MAX(); wrapped the lines over 100 columns] Signed-off-by: Yiwei Lin Assisted-by: LLM --- include/linux/rbtree_augmented.h | 167 ++++++++++++++++++++----------- kernel/sched/fair.c | 70 ++----------- net/tipc/name_table.c | 7 +- 3 files changed, 122 insertions(+), 122 deletions(-) diff --git a/include/linux/rbtree_augmented.h b/include/linux/rbtree_augmented.h index d2fa1c41bfd2b..fbfa675dce5fb 100644 --- a/include/linux/rbtree_augmented.h +++ b/include/linux/rbtree_augmented.h @@ -15,6 +15,8 @@ #include #include #include +#include +#include /* * Please note - only struct rb_augment_callbacks and the prototypes for @@ -86,6 +88,88 @@ rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree, return leftmost ? node : NULL; } +#define RB_FOR_EACH_1(what, RBNAME, RBSTRUCT, RBFIELD, x) \ + what(1, RBNAME, RBSTRUCT, RBFIELD, x) +#define RB_FOR_EACH_2(what, RBNAME, RBSTRUCT, RBFIELD, x, ...) \ + what(2, RBNAME, RBSTRUCT, RBFIELD, x) \ + RB_FOR_EACH_1(what, RBNAME, RBSTRUCT, RBFIELD, __VA_ARGS__) +#define RB_FOR_EACH_3(what, RBNAME, RBSTRUCT, RBFIELD, x, ...) \ + what(3, RBNAME, RBSTRUCT, RBFIELD, x) \ + RB_FOR_EACH_2(what, RBNAME, RBSTRUCT, RBFIELD, __VA_ARGS__) +#define RB_FOR_EACH_4(what, RBNAME, RBSTRUCT, RBFIELD, x, ...) \ + what(4, RBNAME, RBSTRUCT, RBFIELD, x) \ + RB_FOR_EACH_3(what, RBNAME, RBSTRUCT, RBFIELD, __VA_ARGS__) +#define RB_FOR_EACH_5(what, RBNAME, RBSTRUCT, RBFIELD, x, ...) \ + what(5, RBNAME, RBSTRUCT, RBFIELD, x) \ + RB_FOR_EACH_4(what, RBNAME, RBSTRUCT, RBFIELD, __VA_ARGS__) +#define RB_FOR_EACH_6(what, RBNAME, RBSTRUCT, RBFIELD, x, ...) \ + what(6, RBNAME, RBSTRUCT, RBFIELD, x) \ + RB_FOR_EACH_5(what, RBNAME, RBSTRUCT, RBFIELD, __VA_ARGS__) +#define RB_FOR_EACH_7(what, RBNAME, RBSTRUCT, RBFIELD, x, ...) \ + what(7, RBNAME, RBSTRUCT, RBFIELD, x) \ + RB_FOR_EACH_6(what, RBNAME, RBSTRUCT, RBFIELD, __VA_ARGS__) +#define RB_FOR_EACH_8(what, RBNAME, RBSTRUCT, RBFIELD, x, ...) \ + what(8, RBNAME, RBSTRUCT, RBFIELD, x) \ + RB_FOR_EACH_7(what, RBNAME, RBSTRUCT, RBFIELD, __VA_ARGS__) + +#define RB_FOR_EACH(action, RBNAME, RBSTRUCT, RBFIELD, ...) \ + CONCATENATE(RB_FOR_EACH_, COUNT_ARGS(__VA_ARGS__)) \ + (action, RBNAME, RBSTRUCT, RBFIELD, __VA_ARGS__) + +/* + * One augmented field: @val is the node's own contribution (a member for + * RB_AUG(), a function of the node for RB_AUG_FUNC()), @aug the member + * holding the aggregate over the subtree, and @fold(a, b) combines two + * aggregates: min, max, a sum, ... It must be commutative and associative. + */ +#define RB_AUG_FUNC(val, aug, fold) (val(s), aug, fold) +#define RB_AUG(val, aug, fold) (s->val, aug, fold) +#define RB_UNPACK(...) __VA_ARGS__ + +#define __RB_INST(n, RBNAME, RBSTRUCT, RBFIELD, val, aug, fold) \ +static inline void \ +RBNAME ## _copy_ ## n(RBSTRUCT *old, RBSTRUCT *new) \ +{ \ + new->aug = old->aug; \ +} \ +static inline bool \ +RBNAME ## _compute_ ## n(RBSTRUCT *s, bool exit) \ +{ \ + TYPEOF_UNQUAL(s->aug) _old_aug = s->aug; \ + TYPEOF_UNQUAL(s->aug) _val = val; \ + struct rb_node *_node = &s->RBFIELD; \ + if (_node->rb_right) { \ + RBSTRUCT *_c = container_of(_node->rb_right, typeof(*s), RBFIELD); \ + _val = fold(_val, _c->aug); \ + } \ + if (_node->rb_left) { \ + RBSTRUCT *_c = container_of(_node->rb_left, typeof(*s), RBFIELD); \ + _val = fold(_val, _c->aug); \ + } \ + if (exit && _old_aug == _val) \ + return true; \ + s->aug = _val; \ + return false; \ +} +#define _RB_INST(n, RBNAME, RBSTRUCT, RBFIELD, args) \ + __RB_INST(n, RBNAME, RBSTRUCT, RBFIELD, args) +#define RB_INST(n, RBNAME, RBSTRUCT, RBFIELD, x) \ + _RB_INST(n, RBNAME, RBSTRUCT, RBFIELD, RB_UNPACK x) + +#define __RB_COPY(n, RBNAME, RBSTRUCT, RBFIELD, val, aug, fold) \ + RBNAME ## _copy_ ## n(old, new); +#define _RB_COPY(n, RBNAME, RBSTRUCT, RBFIELD, args) \ + __RB_COPY(n, RBNAME, RBSTRUCT, RBFIELD, args) +#define RB_COPY(n, RBNAME, RBSTRUCT, RBFIELD, x) \ + _RB_COPY(n, RBNAME, RBSTRUCT, RBFIELD, RB_UNPACK x) + +#define __RB_COMPUTE(n, RBNAME, RBSTRUCT, RBFIELD, val, aug, fold) \ + ret &= RBNAME ## _compute_ ## n(node, exit); +#define _RB_COMPUTE(n, RBNAME, RBSTRUCT, RBFIELD, args) \ + __RB_COMPUTE(n, RBNAME, RBSTRUCT, RBFIELD, args) +#define RB_COMPUTE(n, RBNAME, RBSTRUCT, RBFIELD, x) \ + _RB_COMPUTE(n, RBNAME, RBSTRUCT, RBFIELD, RB_UNPACK x) + /* * Template for declaring augmented rbtree callbacks (generic multi fields) * @@ -93,18 +177,29 @@ rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree, * RBNAME: name of the rb_augment_callbacks structure * RBSTRUCT: struct type of the tree nodes * RBFIELD: name of struct rb_node field within RBSTRUCT - * RBCOPY: name of function that copies the RBAUGMENTED datas - * RBCOMPUTE: name of function that recomputes the RBAUGMENTED datas + * RBAUG...: list of RB_AUG() describing the augmented data */ - -#define RB_DECLARE_CALLBACKS_MULTI(RBSTATIC, RBNAME, \ - RBSTRUCT, RBFIELD, RBCOPY, RBCOMPUTE) \ +#define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME, \ + RBSTRUCT, RBFIELD, RBAUG...) \ +RB_FOR_EACH(RB_INST, RBNAME, RBSTRUCT, RBFIELD, RBAUG) \ +static inline void \ +RBNAME ## __copy(RBSTRUCT *old, RBSTRUCT *new) \ +{ \ + RB_FOR_EACH(RB_COPY, RBNAME, RBSTRUCT, RBFIELD, RBAUG); \ +} \ +static inline bool \ +RBNAME ## __compute(RBSTRUCT *node, bool exit) \ +{ \ + bool ret = true; \ + RB_FOR_EACH(RB_COMPUTE, RBNAME, RBSTRUCT, RBFIELD, RBAUG); \ + return ret; \ +} \ static inline void \ RBNAME ## _propagate(struct rb_node *rb, struct rb_node *stop) \ { \ while (rb != stop) { \ RBSTRUCT *node = rb_entry(rb, RBSTRUCT, RBFIELD); \ - if (RBCOMPUTE(node, true)) \ + if (RBNAME ## __compute(node, true)) \ break; \ rb = rb_parent(&node->RBFIELD); \ } \ @@ -114,15 +209,15 @@ RBNAME ## _copy(struct rb_node *rb_old, struct rb_node *rb_new) \ { \ RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD); \ RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD); \ - RBCOPY(new, old); \ + RBNAME ## __copy(old, new); \ } \ static void \ RBNAME ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new) \ { \ RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD); \ RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD); \ - RBCOPY(new, old); \ - RBCOMPUTE(old, false); \ + RBNAME ## __copy(old, new); \ + RBNAME ## __compute(old, false); \ } \ RBSTATIC const struct rb_augment_callbacks RBNAME = { \ .propagate = RBNAME ## _propagate, \ @@ -130,27 +225,6 @@ RBSTATIC const struct rb_augment_callbacks RBNAME = { \ .rotate = RBNAME ## _rotate \ }; -/* - * Template for declaring augmented rbtree callbacks (generic single field) - * - * RBSTATIC: 'static' or empty - * RBNAME: name of the rb_augment_callbacks structure - * RBSTRUCT: struct type of the tree nodes - * RBFIELD: name of struct rb_node field within RBSTRUCT - * RBAUGMENTED: name of field within RBSTRUCT holding data for subtree - * RBCOMPUTE: name of function that recomputes the RBAUGMENTED data - */ - -#define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME, \ - RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE) \ -static inline void \ -RBNAME ## _copy_single(RBSTRUCT *new, RBSTRUCT *old) \ -{ \ - new->RBAUGMENTED = old->RBAUGMENTED; \ -} \ -RB_DECLARE_CALLBACKS_MULTI(RBSTATIC, RBNAME, \ - RBSTRUCT, RBFIELD, RBNAME ## _copy_single, RBCOMPUTE) - /* * Template for declaring augmented rbtree callbacks, * computing RBAUGMENTED scalar as max(RBCOMPUTE(node)) for all subtree nodes. @@ -159,34 +233,15 @@ RB_DECLARE_CALLBACKS_MULTI(RBSTATIC, RBNAME, \ * RBNAME: name of the rb_augment_callbacks structure * RBSTRUCT: struct type of the tree nodes * RBFIELD: name of struct rb_node field within RBSTRUCT - * RBTYPE: type of the RBAUGMENTED field - * RBAUGMENTED: name of RBTYPE field within RBSTRUCT holding data for subtree - * RBCOMPUTE: name of function that returns the per-node RBTYPE scalar + * RBTYPE: type of the RBAUGMENTED field -- unused, assumed typeof(RBAUGMENTED) + * RBAUGMENTED: name of field within RBSTRUCT holding data for subtree + * RBVALUE: name of function that returns the per-node RBTYPE scalar */ -#define RB_DECLARE_CALLBACKS_MAX(RBSTATIC, RBNAME, RBSTRUCT, RBFIELD, \ - RBTYPE, RBAUGMENTED, RBCOMPUTE) \ -static inline bool RBNAME ## _compute_max(RBSTRUCT *node, bool exit) \ -{ \ - RBSTRUCT *child; \ - RBTYPE max = RBCOMPUTE(node); \ - if (node->RBFIELD.rb_left) { \ - child = rb_entry(node->RBFIELD.rb_left, RBSTRUCT, RBFIELD); \ - if (child->RBAUGMENTED > max) \ - max = child->RBAUGMENTED; \ - } \ - if (node->RBFIELD.rb_right) { \ - child = rb_entry(node->RBFIELD.rb_right, RBSTRUCT, RBFIELD); \ - if (child->RBAUGMENTED > max) \ - max = child->RBAUGMENTED; \ - } \ - if (exit && node->RBAUGMENTED == max) \ - return true; \ - node->RBAUGMENTED = max; \ - return false; \ -} \ -RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME, \ - RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max) +#define RB_DECLARE_CALLBACKS_MAX(RBSTATIC, RBNAME, RBSTRUCT, RBFIELD, \ + RBTYPE, RBAUGMENTED, RBVALUE) \ +RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME, RBSTRUCT, RBFIELD, \ + RB_AUG_FUNC(RBVALUE, RBAUGMENTED, max)) #define RB_RED 0 diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c index 7455a83a6a990..fa7f01159b493 100644 --- a/kernel/sched/fair.c +++ b/kernel/sched/fair.c @@ -1004,71 +1004,17 @@ static inline bool __entity_less(struct rb_node *a, const struct rb_node *b) return entity_before(__node_2_se(a), __node_2_se(b)); } -static inline void __min_vruntime_update(struct sched_entity *se, struct rb_node *node) +/* min() for wrapping vruntimes */ +static inline u64 __min_vruntime(u64 a, u64 b) { - if (node) { - struct sched_entity *rse = __node_2_se(node); - - if (vruntime_cmp(se->min_vruntime, ">", rse->min_vruntime)) - se->min_vruntime = rse->min_vruntime; - } -} - -static inline void __min_slice_update(struct sched_entity *se, struct rb_node *node) -{ - if (node) { - struct sched_entity *rse = __node_2_se(node); - if (rse->min_slice < se->min_slice) - se->min_slice = rse->min_slice; - } -} - -static inline void __max_slice_update(struct sched_entity *se, struct rb_node *node) -{ - if (node) { - struct sched_entity *rse = __node_2_se(node); - if (rse->max_slice > se->max_slice) - se->max_slice = rse->max_slice; - } -} - -static inline void min_vruntime_copy(struct sched_entity *new, struct sched_entity *old) -{ - new->min_vruntime = old->min_vruntime; - new->min_slice = old->min_slice; - new->max_slice = old->max_slice; + return vruntime_cmp(a, "<", b) ? a : b; } -/* - * se->min_vruntime = min(se->vruntime, {left,right}->min_vruntime) - */ -static inline bool min_vruntime_update(struct sched_entity *se, bool exit) -{ - u64 old_min_vruntime = se->min_vruntime; - u64 old_min_slice = se->min_slice; - u64 old_max_slice = se->max_slice; - struct rb_node *node = &se->run_node; - - se->min_vruntime = se->vruntime; - __min_vruntime_update(se, node->rb_right); - __min_vruntime_update(se, node->rb_left); - - se->min_slice = se->slice; - __min_slice_update(se, node->rb_right); - __min_slice_update(se, node->rb_left); - - se->max_slice = se->slice; - __max_slice_update(se, node->rb_right); - __max_slice_update(se, node->rb_left); - - return se->min_vruntime == old_min_vruntime && - se->min_slice == old_min_slice && - se->max_slice == old_max_slice; -} - - -RB_DECLARE_CALLBACKS_MULTI(static, min_vruntime_cb, struct sched_entity, - run_node, min_vruntime_copy, min_vruntime_update); +RB_DECLARE_CALLBACKS(static, min_vruntime_cb, + struct sched_entity, run_node, + RB_AUG(vruntime, min_vruntime, __min_vruntime), + RB_AUG(slice, min_slice, min), + RB_AUG(slice, max_slice, max)); /* * Enqueue an entity into the rb-tree: diff --git a/net/tipc/name_table.c b/net/tipc/name_table.c index 6fda36ab17669..45189012a0f94 100644 --- a/net/tipc/name_table.c +++ b/net/tipc/name_table.c @@ -88,10 +88,9 @@ struct tipc_service { struct rcu_head rcu; }; -#define service_range_upper(sr) ((sr)->upper) -RB_DECLARE_CALLBACKS_MAX(static, sr_callbacks, - struct service_range, tree_node, u32, max, - service_range_upper) +RB_DECLARE_CALLBACKS(static, sr_callbacks, + struct service_range, tree_node, + RB_AUG(upper, max, max)); #define service_range_entry(rbtree_node) \ (container_of(rbtree_node, struct service_range, tree_node)) -- 2.34.1 rb_add_augmented_cached() links the new node and then calls ->propagate() from the parent up to the root to bring the augmented data of the ancestors up to date. It marks that as "suboptimal": every augmented user that open-codes its insertion (interval trees, vmalloc, drm_mm, gpu buddy, ...) instead folds the new node's value into each ancestor while walking down, which touches no sibling and needs no second pass. ->propagate() has to read both children of every ancestor until it finds one whose value did not change. Add a ->merge() callback to struct rb_augment_callbacks that folds the augmented data of a node being inserted into one of its future ancestors, and use it on the descent. RB_DECLARE_CALLBACKS() derives it from the RB_AUG() list, so every user of the template gets it for free, and sched/eevdf, the only caller of the helper, needs no change. The new node's augmented data must describe the node alone before the call. That was already required, since ->propagate() read it, and sched/eevdf sets it up that way. Measured with rbtree_test's cached augmented test switched to rb_add_augmented_cached() (next patch), nnodes cached insert+delete, median of 5 runs, for reference also the open-coded descent used by augmented test 1. Raspberry Pi 4 Model B (Cortex-A72, 1.8 GHz), arch timer ticks (54 MHz): nnodes open-coded before after 100 459 500 442 (-11.6%) 10000 182046 186144 180208 (-3.2%) 100000 5533679 5711647 5579444 (-2.3%) x86-64 KVM guest (Intel i7-13800H), cycles: nnodes open-coded before after 100 7657 7622 7059 (-7.4%) 10000 4738075 5075847 4627462 (-8.8%) 100000 75618429 79840335 71324700 (-10.7%) The other rbtree_test numbers are unchanged and the augmented invariant checks still pass on both. Signed-off-by: Yiwei Lin Assisted-by: LLM --- Documentation/core-api/rbtree.rst | 27 ++++++++++++++++++++++++--- include/linux/rbtree_augmented.h | 27 +++++++++++++++++++++++++-- 2 files changed, 49 insertions(+), 5 deletions(-) diff --git a/Documentation/core-api/rbtree.rst b/Documentation/core-api/rbtree.rst index cce80e19087b4..210045f128d65 100644 --- a/Documentation/core-api/rbtree.rst +++ b/Documentation/core-api/rbtree.rst @@ -257,8 +257,14 @@ When erasing a node, the user must call rb_erase_augmented() instead of rb_erase(). rb_erase_augmented() calls back into user provided functions to update the augmented information on affected subtrees. -In both cases, the callbacks are provided through struct rb_augment_callbacks. -3 callbacks must be defined: +Alternatively, rb_add_augmented_cached() performs the whole insertion into +a cached tree: it walks down to the insertion point, merges the new node's +augmented information into every node on that path, links the node and +rebalances. The new node's augmented information must already describe +the node alone when it is called. + +In all cases, the callbacks are provided through struct rb_augment_callbacks. +3 callbacks must be defined, plus a fourth one for rb_add_augmented_cached(): - A propagation callback, which updates the augmented value for a given node and its ancestors, up to a given stop point (or NULL to update @@ -271,6 +277,10 @@ In both cases, the callbacks are provided through struct rb_augment_callbacks. subtree to a newly assigned subtree root AND recomputes the augmented information for the former subtree root. +- A merge callback, which folds the augmented value of a node being inserted + into the augmented value of one of its future ancestors. It is only used + by rb_add_augmented_cached(). + The compiled code for rb_erase_augmented() may inline the propagation and copy callbacks, which results in a large function, so each augmented rbtree user should have a single rb_erase_augmented() call site in order to limit @@ -395,8 +405,19 @@ Insertion/removal are defined using the following augmented callbacks:: old->__subtree_last = compute_subtree_last(old); } + static void augment_merge(struct rb_node *rb, struct rb_node *rb_new) + { + struct interval_tree_node *node = + rb_entry(rb, struct interval_tree_node, rb); + struct interval_tree_node *new = + rb_entry(rb_new, struct interval_tree_node, rb); + + if (node->__subtree_last < new->__subtree_last) + node->__subtree_last = new->__subtree_last; + } + static const struct rb_augment_callbacks augment_callbacks = { - augment_propagate, augment_copy, augment_rotate + augment_propagate, augment_copy, augment_rotate, augment_merge }; void interval_tree_insert(struct interval_tree_node *node, diff --git a/include/linux/rbtree_augmented.h b/include/linux/rbtree_augmented.h index fbfa675dce5fb..9ac0f216cc301 100644 --- a/include/linux/rbtree_augmented.h +++ b/include/linux/rbtree_augmented.h @@ -30,6 +30,7 @@ struct rb_augment_callbacks { void (*propagate)(struct rb_node *node, struct rb_node *stop); void (*copy)(struct rb_node *old, struct rb_node *new); void (*rotate)(struct rb_node *old, struct rb_node *new); + void (*merge)(struct rb_node *node, struct rb_node *new); }; extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root, @@ -62,6 +63,12 @@ rb_insert_augmented_cached(struct rb_node *node, rb_insert_augmented(node, &root->rb_root, augment); } +/* + * Insert @node into the leftmost cached augmented tree @tree. + * + * The augmented data of @node must already describe @node alone; it is + * merged into every ancestor on the way down through augment->merge(). + */ static __always_inline struct rb_node * rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree, bool (*less)(struct rb_node *, const struct rb_node *), @@ -73,6 +80,7 @@ rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree, while (*link) { parent = *link; + augment->merge(parent, node); if (less(node, parent)) { link = &parent->rb_left; } else { @@ -82,7 +90,6 @@ rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree, } rb_link_node(node, parent, link); - augment->propagate(parent, NULL); /* suboptimal */ rb_insert_augmented_cached(node, tree, leftmost, augment); return leftmost ? node : NULL; @@ -170,6 +177,14 @@ RBNAME ## _compute_ ## n(RBSTRUCT *s, bool exit) \ #define RB_COMPUTE(n, RBNAME, RBSTRUCT, RBFIELD, x) \ _RB_COMPUTE(n, RBNAME, RBSTRUCT, RBFIELD, RB_UNPACK x) +/* fold @new, about to become a descendant of @node, into @node */ +#define __RB_MERGE(n, RBNAME, RBSTRUCT, RBFIELD, val, aug, fold) \ + node->aug = fold(node->aug, new->aug); +#define _RB_MERGE(n, RBNAME, RBSTRUCT, RBFIELD, args) \ + __RB_MERGE(n, RBNAME, RBSTRUCT, RBFIELD, args) +#define RB_MERGE(n, RBNAME, RBSTRUCT, RBFIELD, x) \ + _RB_MERGE(n, RBNAME, RBSTRUCT, RBFIELD, RB_UNPACK x) + /* * Template for declaring augmented rbtree callbacks (generic multi fields) * @@ -219,10 +234,18 @@ RBNAME ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new) \ RBNAME ## __copy(old, new); \ RBNAME ## __compute(old, false); \ } \ +static inline void \ +RBNAME ## _merge(struct rb_node *rb, struct rb_node *rb_new) \ +{ \ + RBSTRUCT *node = rb_entry(rb, RBSTRUCT, RBFIELD); \ + RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD); \ + RB_FOR_EACH(RB_MERGE, RBNAME, RBSTRUCT, RBFIELD, RBAUG); \ +} \ RBSTATIC const struct rb_augment_callbacks RBNAME = { \ .propagate = RBNAME ## _propagate, \ .copy = RBNAME ## _copy, \ - .rotate = RBNAME ## _rotate \ + .rotate = RBNAME ## _rotate, \ + .merge = RBNAME ## _merge \ }; /* -- 2.34.1 insert_augmented_cached() open-codes the descent that the generic rb_add_augmented_cached() helper provides. Use the helper instead, so that rbtree_test covers it: it is what sched/eevdf relies on and nothing in lib/ exercises it today. Signed-off-by: Yiwei Lin Assisted-by: LLM --- lib/rbtree_test.c | 25 ++----------------------- 1 file changed, 2 insertions(+), 23 deletions(-) diff --git a/lib/rbtree_test.c b/lib/rbtree_test.c index c1386d328cfb1..00f62e4f9d10e 100644 --- a/lib/rbtree_test.c +++ b/lib/rbtree_test.c @@ -89,29 +89,8 @@ static void insert_augmented(struct test_node *node, static void insert_augmented_cached(struct test_node *node, struct rb_root_cached *root) { - struct rb_node **new = &root->rb_root.rb_node, *rb_parent = NULL; - u32 key = node->key; - u32 val = node->val; - struct test_node *parent; - bool leftmost = true; - - while (*new) { - 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.rb_right; - leftmost = false; - } - } - - node->augmented = val; - rb_link_node(&node->rb, rb_parent, new); - rb_insert_augmented_cached(&node->rb, root, - leftmost, &augment_callbacks); + node->augmented = node->val; + rb_add_augmented_cached(&node->rb, root, less, &augment_callbacks); } -- 2.34.1