[PATCH 0/3] rbtree: fix rb_add_augmented_cached() descent and exercise rb_add*() helpers in rbtree_test

[PATCH 0/3] rbtree: fix rb_add_augmented_cached() descent and exercise rb_add*() helpers in rbtree_test

[Yiwei Lin]

rbtree_test open-codes the insertion of every flavour of rbtree it
exercises, while the generic rb_add*() helpers have been the way most
users insert nodes for years now.  Patches 1 and 3 make the test use the
helpers where the helper does exactly what the test did, so the helpers
themselves get covered.

Converting the cached augmented test exposed the cost of the
"suboptimal" propagate-from-parent path in rb_add_augmented_cached():
2-3% on a Raspberry Pi 4 and 5-7% on an x86-64 KVM guest on the
augmented insert+delete benchmark, against the documented
update-on-the-way-down pattern.  Patch 2 adds a ->merge() callback to
struct rb_augment_callbacks and uses it during the descent, which gets
the helper to parity before the test starts relying on it.
sched/eevdf, its only user, supplies the callback from its existing
per-field helpers; the resulting kernel boots and runs on the Pi.

The augmented invariant checks pass at every step.

Yiwei Lin (3):
  rbtree_test: use rb_add() and rb_add_cached() for the basic tests
  rbtree: update augmented data on the way down in
    rb_add_augmented_cached()
  rbtree_test: use rb_add_augmented_cached() for the cached augmented
    test

 Documentation/core-api/rbtree.rst | 27 ++++++++++++--
 include/linux/rbtree_augmented.h  | 40 +++++++++++++++++---
 kernel/sched/fair.c               | 13 ++++++-
 lib/rbtree_test.c                 | 62 +++++--------------------------
 4 files changed, 80 insertions(+), 62 deletions(-)

-- 
2.34.1

[PATCH 1/3] rbtree_test: use rb_add() and rb_add_cached() for the basic tests

[Yiwei Lin]

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 <s921975628@gmail.com>
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

[PATCH 2/3] rbtree: update augmented data on the way down in rb_add_augmented_cached()

[Yiwei Lin]

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.

The new node's augmented data must describe the node alone before the
call. That was already required, since ->propagate() read it, and both
callers set it up that way.

Measured with rbtree_test's cached augmented test switched to
rb_add_augmented_cached() (next patch) on a Raspberry Pi 4 Model B
(Cortex-A72, 1.8 GHz), nnodes cached insert+delete, median of 5 runs,
arch timer ticks (54 MHz), for reference also the open-coded descent
used by augmented test 1. This brings the helper to the level of
the open-coded descent:

    nnodes    open-coded      before       after
       100           468         486         475  (-2.3%)
     10000        182772      188644      185810  (-1.5%)
    100000       5551237     5700029     5518725  (-3.2%)

Signed-off-by: Yiwei Lin <s921975628@gmail.com>
Assisted-by: LLM
---
 Documentation/core-api/rbtree.rst | 27 ++++++++++++++++++---
 include/linux/rbtree_augmented.h  | 40 ++++++++++++++++++++++++++-----
 kernel/sched/fair.c               | 13 +++++++++-
 3 files changed, 70 insertions(+), 10 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 d2fa1c41bfd2b..6a1b78f6cb1e2 100644
--- a/include/linux/rbtree_augmented.h
+++ b/include/linux/rbtree_augmented.h
@@ -28,6 +28,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,
@@ -60,6 +61,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 *),
@@ -71,6 +78,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 {
@@ -80,7 +88,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;
@@ -95,10 +102,13 @@ rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree,
  * 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
+ * RBMERGE:     name of function that merges a new node's RBAUGMENTED datas
+ *              into an ancestor's, used on the insertion path
  */
 
 #define RB_DECLARE_CALLBACKS_MULTI(RBSTATIC, RBNAME,			\
-			     RBSTRUCT, RBFIELD, RBCOPY, RBCOMPUTE)	\
+			     RBSTRUCT, RBFIELD, RBCOPY, RBCOMPUTE,	\
+			     RBMERGE)					\
 static inline void							\
 RBNAME ## _propagate(struct rb_node *rb, struct rb_node *stop)		\
 {									\
@@ -124,10 +134,18 @@ RBNAME ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)	\
 	RBCOPY(new, old);						\
 	RBCOMPUTE(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);		\
+	RBMERGE(node, new);						\
+}									\
 RBSTATIC const struct rb_augment_callbacks RBNAME = {			\
 	.propagate = RBNAME ## _propagate,				\
 	.copy = RBNAME ## _copy,					\
-	.rotate = RBNAME ## _rotate					\
+	.rotate = RBNAME ## _rotate,					\
+	.merge = RBNAME ## _merge					\
 };
 
 /*
@@ -139,17 +157,21 @@ RBSTATIC const struct rb_augment_callbacks RBNAME = {			\
  * 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
+ * RBMERGE:     name of function that merges a new node's RBAUGMENTED data
+ *              into an ancestor's, used on the insertion path
  */
 
 #define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,				\
-			     RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE)	\
+			     RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE,	\
+			     RBMERGE)					\
 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)
+		     RBSTRUCT, RBFIELD, RBNAME ## _copy_single,		\
+		     RBCOMPUTE, RBMERGE)
 
 /*
  * Template for declaring augmented rbtree callbacks,
@@ -185,8 +207,14 @@ static inline bool RBNAME ## _compute_max(RBSTRUCT *node, bool exit)	      \
 	node->RBAUGMENTED = max;					      \
 	return false;							      \
 }									      \
+static inline void RBNAME ## _merge_max(RBSTRUCT *node, RBSTRUCT *new)	      \
+{									      \
+	if (node->RBAUGMENTED < new->RBAUGMENTED)			      \
+		node->RBAUGMENTED = new->RBAUGMENTED;			      \
+}									      \
 RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,					      \
-		     RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max)
+		     RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max,  \
+		     RBNAME ## _merge_max)
 
 
 #define	RB_RED		0
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 7455a83a6a990..60db4624b9897 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -1066,9 +1066,20 @@ static inline bool min_vruntime_update(struct sched_entity *se, bool exit)
 	       se->max_slice == old_max_slice;
 }
 
+/*
+ * Fold @new's subtree data into @se, for each @se on @new's insertion path.
+ */
+static inline void
+min_vruntime_merge(struct sched_entity *se, struct sched_entity *new)
+{
+	__min_vruntime_update(se, &new->run_node);
+	__min_slice_update(se, &new->run_node);
+	__max_slice_update(se, &new->run_node);
+}
 
 RB_DECLARE_CALLBACKS_MULTI(static, min_vruntime_cb, struct sched_entity,
-		     run_node, min_vruntime_copy, min_vruntime_update);
+		     run_node, min_vruntime_copy, min_vruntime_update,
+		     min_vruntime_merge);
 
 /*
  * Enqueue an entity into the rb-tree:
-- 
2.34.1

Re: [PATCH 2/3] rbtree: update augmented data on the way down in rb_add_augmented_cached()

[Peter Zijlstra]

On Mon, Sep 28, 2026 at 08:26:10PM +0800, Yiwei Lin wrote:

> diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
> index 7455a83a6a990..60db4624b9897 100644
> --- a/kernel/sched/fair.c
> +++ b/kernel/sched/fair.c
> @@ -1066,9 +1066,20 @@ static inline bool min_vruntime_update(struct sched_entity *se, bool exit)
>  	       se->max_slice == old_max_slice;
>  }
>  
> +/*
> + * Fold @new's subtree data into @se, for each @se on @new's insertion path.
> + */
> +static inline void
> +min_vruntime_merge(struct sched_entity *se, struct sched_entity *new)
> +{
> +	__min_vruntime_update(se, &new->run_node);
> +	__min_slice_update(se, &new->run_node);
> +	__max_slice_update(se, &new->run_node);
> +}
>  

So min_vruntime_update() can be written in terms of this helper like:

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;
	se->min_slice = se->slice;
	se->max_slice = se->slice;

	min_vruntime_merge(se, node->rb_right);
	min_vruntime_merge(se, node->rb_left);

	return se->min_vruntime == old_min_vruntime &&
	       se->min_slice == old_min_slice &&
	       se->max_slice == old_max_slice;
}

And that is *very* close to being generalizable, obviating the need for
RBCOMPUTE. Does your LLM see a way to make that happen?

That is, for RB_DECLARE_CALLBACKS() and RD_DECLARE_CALLBACKS_MAX() this
is trivially doable, but I'm not sure I see a clean way to make this
happen for RB_DECLARE_CALLBACKS_MULTI().

Re: [PATCH 2/3] rbtree: update augmented data on the way down in rb_add_augmented_cached()

[Peter Zijlstra]

On Mon, Sep 28, 2026 at 03:37:33PM +0200, Peter Zijlstra wrote:
> On Mon, Sep 28, 2026 at 08:26:10PM +0800, Yiwei Lin wrote:
> 
> > diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
> > index 7455a83a6a990..60db4624b9897 100644
> > --- a/kernel/sched/fair.c
> > +++ b/kernel/sched/fair.c
> > @@ -1066,9 +1066,20 @@ static inline bool min_vruntime_update(struct sched_entity *se, bool exit)
> >  	       se->max_slice == old_max_slice;
> >  }
> >  
> > +/*
> > + * Fold @new's subtree data into @se, for each @se on @new's insertion path.
> > + */
> > +static inline void
> > +min_vruntime_merge(struct sched_entity *se, struct sched_entity *new)
> > +{
> > +	__min_vruntime_update(se, &new->run_node);
> > +	__min_slice_update(se, &new->run_node);
> > +	__max_slice_update(se, &new->run_node);
> > +}
> >  
> 
> So min_vruntime_update() can be written in terms of this helper like:
> 
> 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;
> 	se->min_slice = se->slice;
> 	se->max_slice = se->slice;
> 
> 	min_vruntime_merge(se, node->rb_right);
> 	min_vruntime_merge(se, node->rb_left);
> 
> 	return se->min_vruntime == old_min_vruntime &&
> 	       se->min_slice == old_min_slice &&
> 	       se->max_slice == old_max_slice;
> }
> 
> And that is *very* close to being generalizable, obviating the need for
> RBCOMPUTE. Does your LLM see a way to make that happen?

Perhaps by doing something like so?

diff --git a/include/linux/rbtree_augmented.h b/include/linux/rbtree_augmented.h
index d2fa1c41bfd2..c4cf5160aa17 100644
--- a/include/linux/rbtree_augmented.h
+++ b/include/linux/rbtree_augmented.h
@@ -15,6 +15,7 @@
 #include <linux/compiler.h>
 #include <linux/rbtree.h>
 #include <linux/rcupdate.h>
+#include <linux/args.h>
 
 /*
  * Please note - only struct rb_augment_callbacks and the prototypes for
@@ -86,6 +87,20 @@ rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree,
 	return leftmost ? node : NULL;
 }
 
+#define FOR_EACH_1(what, x)       what(x)
+#define FOR_EACH_2(what, x, ...)  what(x) FOR_EACH_1(what, __VA_ARGS__)
+#define FOR_EACH_3(what, x, ...)  what(x) FOR_EACH_2(what, __VA_ARGS__)
+#define FOR_EACH_4(what, x, ...)  what(x) FOR_EACH_3(what, __VA_ARGS__)
+#define FOR_EACH_5(what, x, ...)  what(x) FOR_EACH_4(what, __VA_ARGS__)
+#define FOR_EACH_6(what, x, ...)  what(x) FOR_EACH_5(what, __VA_ARGS__)
+#define FOR_EACH_7(what, x, ...)  what(x) FOR_EACH_6(what, __VA_ARGS__)
+#define FOR_EACH_8(what, x, ...)  what(x) FOR_EACH_7(what, __VA_ARGS__)
+
+#define FOR_EACH(action, ...) \
+	CONCATENATE(FOR_EACH_, COUNT_ARGS(__VA_ARGS__))(action, __VA_ARGS__)
+
+#define COPY_VAL(x) new->x = old->x;
+
 /*
  * Template for declaring augmented rbtree callbacks (generic multi fields)
  *
@@ -93,12 +108,12 @@ 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...:	field names within RBSTRUCT holding data for the subtree
  */
 
 #define RB_DECLARE_CALLBACKS_MULTI(RBSTATIC, RBNAME,			\
-			     RBSTRUCT, RBFIELD, RBCOPY, RBCOMPUTE)	\
+			     RBSTRUCT, RBFIELD, RBCOMPUTE, RBAUG...)	\
 static inline void							\
 RBNAME ## _propagate(struct rb_node *rb, struct rb_node *stop)		\
 {									\
@@ -110,18 +125,23 @@ RBNAME ## _propagate(struct rb_node *rb, struct rb_node *stop)		\
 	}								\
 }									\
 static inline void							\
+RBNAME ## __copy(RBSTRUCT *new, RBSTRUCT *old)				\
+{									\
+	FOR_EACH(COPY_VAL, RBAUG);					\
+}									\
+static inline void							\
 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(new, old);					\
 }									\
 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);						\
+	RBNAME ## __copy(new, old);					\
 	RBCOMPUTE(old, false);						\
 }									\
 RBSTATIC const struct rb_augment_callbacks RBNAME = {			\
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index e707da7177df..cc4cd3565268 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -1061,9 +1061,8 @@ static inline bool min_vruntime_update(struct sched_entity *se, bool exit)
 	       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);
+		     run_node, min_vruntime_update, min_vruntime, min_slice, max_slice)
 
 /*
  * Enqueue an entity into the rb-tree:

Re: [PATCH 2/3] rbtree: update augmented data on the way down in rb_add_augmented_cached()

[Peter Zijlstra]

On Mon, Sep 28, 2026 at 04:05:13PM +0200, Peter Zijlstra wrote:
> On Mon, Sep 28, 2026 at 03:37:33PM +0200, Peter Zijlstra wrote:
> > On Mon, Sep 28, 2026 at 08:26:10PM +0800, Yiwei Lin wrote:
> > 
> > > diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
> > > index 7455a83a6a990..60db4624b9897 100644
> > > --- a/kernel/sched/fair.c
> > > +++ b/kernel/sched/fair.c
> > > @@ -1066,9 +1066,20 @@ static inline bool min_vruntime_update(struct sched_entity *se, bool exit)
> > >  	       se->max_slice == old_max_slice;
> > >  }
> > >  
> > > +/*
> > > + * Fold @new's subtree data into @se, for each @se on @new's insertion path.
> > > + */
> > > +static inline void
> > > +min_vruntime_merge(struct sched_entity *se, struct sched_entity *new)
> > > +{
> > > +	__min_vruntime_update(se, &new->run_node);
> > > +	__min_slice_update(se, &new->run_node);
> > > +	__max_slice_update(se, &new->run_node);
> > > +}
> > >  
> > 
> > So min_vruntime_update() can be written in terms of this helper like:
> > 
> > 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;
> > 	se->min_slice = se->slice;
> > 	se->max_slice = se->slice;
> > 
> > 	min_vruntime_merge(se, node->rb_right);
> > 	min_vruntime_merge(se, node->rb_left);
> > 
> > 	return se->min_vruntime == old_min_vruntime &&
> > 	       se->min_slice == old_min_slice &&
> > 	       se->max_slice == old_max_slice;
> > }
> > 
> > And that is *very* close to being generalizable, obviating the need for
> > RBCOMPUTE. Does your LLM see a way to make that happen?

Also, since you're poking at things here, can you please include a patch
to rename RB_DECLARE_CALLBACKS_MAX()'s RBCOMPUTE, it is totally
different from the earlier RBCOMPUTE. Them sharing a name is highly
confusing.

Perhaps rename to something like RBVALUE?

Re: [PATCH 2/3] rbtree: update augmented data on the way down in rb_add_augmented_cached()

[Yiwei Lin]

On 2026/09/28 22:05, Peter Zijlstra wrote:
> > And that is *very* close to being generalizable, obviating the need for
> > RBCOMPUTE. Does your LLM see a way to make that happen?
>
> Perhaps by doing something like so?
[...]
>  #define RB_DECLARE_CALLBACKS_MULTI(RBSTATIC, RBNAME,			\
> -			     RBSTRUCT, RBFIELD, RBCOPY, RBCOMPUTE)	\
> +			     RBSTRUCT, RBFIELD, RBCOMPUTE, RBAUG...)	\

Yes, and since RBCOMPUTE is always "reset the fields to the node's own contribution,
then fold in each child". I think we can also generalize it with RBMERGE, by
adding one extra RBINIT:

        static inline void min_vruntime_init(struct sched_entity *se)
        {
                se->min_vruntime = se->vruntime;
                se->min_slice = se->slice;
                se->max_slice = se->slice;
        }

So the original RBCOMPUTE can be generated by the template instead, which allow us to
remove min_vruntime_update.

        static inline bool RBNAME ## _compute(RBSTRUCT *node)                         \
        {                                                                             \
                typeof(node->RBAUGMENTED) aug;                                        \
                RBSTRUCT *child;                                                      \
                                                                                      \
                RBINIT(node, &aug);                                                   \
                if (node->RBFIELD.rb_left) {                                          \
                        child = rb_entry(node->RBFIELD.rb_left, RBSTRUCT, RBFIELD);   \
                        RBMERGE(&aug, &child->RBAUGMENTED);                           \
                }                                                                     \
                if (node->RBFIELD.rb_right) {                                         \
                        child = rb_entry(node->RBFIELD.rb_right, RBSTRUCT, RBFIELD);  \
                        RBMERGE(&aug, &child->RBAUGMENTED);                           \
                }                                                                     \
                if (!memcmp(&aug, &node->RBAUGMENTED, sizeof(aug)))                   \
                        return true;                                                  \
                node->RBAUGMENTED = aug;                                              \
                return false;                                                         \
        }                                                                             \

Base on this I have another idea: can we let augmented data becomes one member and
init/merge just work on its type by value. Taking sched for example:

        struct_group_tagged(sched_aug, aug,
                u64 min_vruntime;
                u64 min_slice;
                u64 max_slice;
        );

So fair.c's se->min_vruntime etc. can stay as they are, we don't have to rely on the
FOR_EACH machinery. Do you think this can be better or do you prefer the field-list form?

Thanks,
Yiwei Lin

Re: [PATCH 2/3] rbtree: update augmented data on the way down in rb_add_augmented_cached()

[Peter Zijlstra]

On Mon, Sep 28, 2026 at 04:05:13PM +0200, Peter Zijlstra wrote:

> > And that is *very* close to being generalizable, obviating the need for
> > RBCOMPUTE. Does your LLM see a way to make that happen?
> 
> Perhaps by doing something like so?

This builds and boot, so it must be perfect, right ;-)

Assisted-by: Brain
Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 include/linux/rbtree_augmented.h | 153 +++++++++++++++++++++++++--------------
 kernel/sched/fair.c              |  69 ++----------------
 net/tipc/name_table.c            |   7 +-
 3 files changed, 108 insertions(+), 121 deletions(-)

diff --git a/include/linux/rbtree_augmented.h b/include/linux/rbtree_augmented.h
index d2fa1c41bfd2..9c1707aa83b2 100644
--- a/include/linux/rbtree_augmented.h
+++ b/include/linux/rbtree_augmented.h
@@ -15,6 +15,7 @@
 #include <linux/compiler.h>
 #include <linux/rbtree.h>
 #include <linux/rcupdate.h>
+#include <linux/args.h>
 
 /*
  * Please note - only struct rb_augment_callbacks and the prototypes for
@@ -86,6 +87,74 @@ 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__)
+
+#define RB_AUG_FUNC(val, aug, cmp) (val(s), aug, cmp)
+#define RB_AUG(val, aug, cmp) (s->val, aug, cmp)
+#define RB_UNPACK(...) __VA_ARGS__
+
+#define __RB_INST(n, RBNAME, RBSTRUCT, RBFIELD, val, aug, cmp)		\
+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); \
+		if (cmp(_c->aug, _val))				\
+			_val = _c->aug;					\
+	}								\
+	if (_node->rb_left) {						\
+		RBSTRUCT *_c = container_of(_node->rb_left, typeof(*s), RBFIELD); \
+		if (cmp(_c->aug, _val))				\
+			_val = _c->aug;					\
+	}								\
+	s->aug = _val;							\
+	return _old_aug == _val;					\
+}
+#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, cmp)		\
+	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, cmp)	\
+	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 +162,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 aughmented 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 +194,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,26 +210,8 @@ 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)
+#define RB_MAX(a, b) (a > b)
+#define RB_MIN(a, b) (a < b)
 
 /*
  * Template for declaring augmented rbtree callbacks,
@@ -159,34 +221,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, RB_MAX))
 
 
 #define	RB_RED		0
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index e707da7177df..967e71c15a24 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -999,71 +999,16 @@ 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)
+static inline bool __min_vruntime_cmp(u64 child_vruntime, u64 min_vruntime)
 {
-	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;
-	}
+	return vruntime_cmp(min_vruntime, ">", child_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;
-}
-
-/*
- * 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_cmp),
+		     RB_AUG(slice, min_slice, RB_MIN),
+		     RB_AUG(slice, max_slice, RB_MAX));
 
 /*
  * Enqueue an entity into the rb-tree:
diff --git a/net/tipc/name_table.c b/net/tipc/name_table.c
index 6fda36ab1766..bed04a1f0fbd 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, RB_MAX));
 
 #define service_range_entry(rbtree_node)				\
 	(container_of(rbtree_node, struct service_range, tree_node))

[PATCH 3/3] rbtree_test: use rb_add_augmented_cached() for the cached augmented test

[Yiwei Lin]

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 <s921975628@gmail.com>
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