git: 7ba58de6b6d0 - stable/13 - rb_tree: update augmentation after element change

From: Doug Moore <dougm_at_FreeBSD.org>
Date: Mon, 22 Aug 2022 09:22:20 UTC
The branch stable/13 has been updated by dougm:

URL: https://cgit.FreeBSD.org/src/commit/?id=7ba58de6b6d015094f4153030dabcc441a179c54

commit 7ba58de6b6d015094f4153030dabcc441a179c54
Author:     Doug Moore <dougm@FreeBSD.org>
AuthorDate: 2022-08-02 16:19:46 +0000
Commit:     Doug Moore <dougm@FreeBSD.org>
CommitDate: 2022-08-22 09:21:37 +0000

    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
    
    (cherry picked from commit 35557a0d9169172b90df903b0daf389c2839b749)
---
 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 c45ce3d3ceb6..86d6569d1955 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.
@@ -607,12 +610,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 5d215d0babe9..b66fa88829b2 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);							\
 }