From 35557a0d9169172b90df903b0daf389c2839b749 Mon Sep 17 00:00:00 2001
From: Doug Moore <dougm@FreeBSD.org>
Date: Tue, 2 Aug 2022 11:19:46 -0500
Subject: [PATCH] rb_tree: update augmentation after element change

For an augmented rb_tree, allow a faster alternative to removing an
element from the tree, tweaking it slightly, and inserting it back
into the tree, knowing that its relative position in the tree is
unchanged. Instead, just change the element and invoke
RB_UPDATE_AUGMENT to fix the augmentation data for all the nodes in
the tree.

Reviewed by:	kib
MFC after:	1 week
Differential Revision:	https://reviews.freebsd.org/D36010
---
 share/man/man3/tree.3 | 24 ++++++++++++++++++++----
 sys/sys/tree.h        | 18 ++++++++++--------
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/share/man/man3/tree.3 b/share/man/man3/tree.3
index c68c71fff85b..cbc741fe642d 100644
--- a/share/man/man3/tree.3
+++ b/share/man/man3/tree.3
@@ -100,6 +100,7 @@
 .Nm RB_REMOVE ,
 .Nm RB_REINSERT ,
 .Nm RB_AUGMENT
+.Nm RB_UPDATE_AUGMENT
 .Nd "implementations of splay and rank-balanced (wavl) trees"
 .Sh SYNOPSIS
 .In sys/tree.h
@@ -196,6 +197,8 @@
 .Fn RB_REINSERT NAME "RB_HEAD *head" "struct TYPE *elm"
 .Ft "void"
 .Fn RB_AUGMENT NAME "struct TYPE *elm"
+.Ft "void"
+.Fn RB_UPDATE_AUGMENT NAME "struct TYPE *elm"
 .Sh DESCRIPTION
 These macros define data structures for different types of trees:
 splay trees and rank-balanced (wavl) trees.
@@ -592,12 +595,25 @@ macro updates augmentation data of the element
 in the tree.
 By default, it has no effect.
 It is not meant to be invoked by the RB user.
-If RB_AUGMENT is defined by the RB user, then when an element is
-inserted or removed from the tree, it is invoked for every element in
-the tree that is the root of an altered subtree, working from the
-bottom of the tree up to the top.
+If
+.Fn RB_AUGMENT
+is defined by the RB user, then when an element is inserted or removed
+from the tree, it is invoked for every element in the tree that is the
+root of an altered subtree, working from the bottom of the tree up to
+the top.
 It is typically used to maintain some associative accumulation of tree
 elements, such as sums, minima, maxima, and the like.
+.Pp
+The
+.Fn RB_UPDATE_AUGMENT
+macro updates augmentation data of the element
+.Fa elm
+and its ancestors in the tree.
+If RB_AUGMENT is defined by the RB user, then when an element in the
+tree is changed, without changing the order of items in the tree,
+invoking this function on that element restores consistency of the
+augmentation state of the tree as if the element had been removed and
+inserted again.
 .Sh EXAMPLES
 The following example demonstrates how to declare a rank-balanced tree
 holding integers.
diff --git a/sys/sys/tree.h b/sys/sys/tree.h
index 1529b07dfcdb..c0d21b5f8b73 100644
--- a/sys/sys/tree.h
+++ b/sys/sys/tree.h
@@ -371,6 +371,13 @@ struct {								\
 #define RB_AUGMENT(x)	break
 #endif
 
+#define RB_UPDATE_AUGMENT(elm, field) do {				\
+	__typeof(elm) tmp = (elm);					\
+	do {								\
+		RB_AUGMENT(tmp);					\
+	} while ((tmp = RB_PARENT(tmp, field)) != NULL);		\
+} while (0)
+
 #define RB_SWAP_CHILD(head, out, in, field) do {			\
 	if (RB_PARENT(out, field) == NULL)				\
 		RB_ROOT(head) = (in);					\
@@ -678,11 +685,9 @@ name##_RB_REMOVE(struct name *head, struct type *elm)			\
 	RB_SWAP_CHILD(head, old, elm, field);				\
 	if (child != NULL)						\
 		RB_SET_PARENT(child, parent, field);			\
-	if (parent != NULL)						\
+	if (parent != NULL) {						\
 		name##_RB_REMOVE_COLOR(head, parent, child);		\
-	while (parent != NULL) {					\
-		RB_AUGMENT(parent);					\
-		parent = RB_PARENT(parent, field);			\
+		RB_UPDATE_AUGMENT(parent, field);			\
 	}								\
 	return (old);							\
 }
@@ -714,10 +719,7 @@ name##_RB_INSERT(struct name *head, struct type *elm)			\
 	else								\
 		RB_RIGHT(parent, field) = elm;				\
 	name##_RB_INSERT_COLOR(head, elm);				\
-	while (elm != NULL) {						\
-		RB_AUGMENT(elm);					\
-		elm = RB_PARENT(elm, field);				\
-	}								\
+	RB_UPDATE_AUGMENT(elm, field);					\
 	return (NULL);							\
 }