git: 35557a0d9169 - main - rb_tree: update augmentation after element change

From: Doug Moore <dougm_at_FreeBSD.org>
Date: Tue, 02 Aug 2022 16:23:24 UTC
The branch main has been updated by dougm:

URL: https://cgit.FreeBSD.org/src/commit/?id=35557a0d9169172b90df903b0daf389c2839b749

commit 35557a0d9169172b90df903b0daf389c2839b749
Author:     Doug Moore <dougm@FreeBSD.org>
AuthorDate: 2022-08-02 16:19:46 +0000
Commit:     Doug Moore <dougm@FreeBSD.org>
CommitDate: 2022-08-02 16:23:12 +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
---
 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);							\
 }