include/linux/rbtree_augmented.h

0001 /* SPDX-License-Identifier: GPL-2.0-or-later */
0002 /*
0003   Red Black Trees
0004   (C) 1999  Andrea Arcangeli <andrea@suse.de>
0005   (C) 2002  David Woodhouse <dwmw2@infradead.org>
0006   (C) 2012  Michel Lespinasse <walken@google.com>
0007
0008
0009   linux/include/linux/rbtree_augmented.h
0010 */
0011
0012 #ifndef _LINUX_RBTREE_AUGMENTED_H
0013 #define _LINUX_RBTREE_AUGMENTED_H
0014
0015 #include <linux/compiler.h>
0016 #include <linux/rbtree.h>
0017 #include <linux/rcupdate.h>
0018
0019 /*
0020  * Please note - only struct rb_augment_callbacks and the prototypes for
0021  * rb_insert_augmented() and rb_erase_augmented() are intended to be public.
0022  * The rest are implementation details you are not expected to depend on.
0023  *
0024  * See Documentation/core-api/rbtree.rst for documentation and samples.
0025  */
0026
0027 struct rb_augment_callbacks {
0028     void (*propagate)(struct rb_node *node, struct rb_node *stop);
0029     void (*copy)(struct rb_node *old, struct rb_node *new);
0030     void (*rotate)(struct rb_node *old, struct rb_node *new);
0031 };
0032
0033 extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
0034     void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
0035
0036 /*
0037  * Fixup the rbtree and update the augmented information when rebalancing.
0038  *
0039  * On insertion, the user must update the augmented information on the path
0040  * leading to the inserted node, then call rb_link_node() as usual and
0041  * rb_insert_augmented() instead of the usual rb_insert_color() call.
0042  * If rb_insert_augmented() rebalances the rbtree, it will callback into
0043  * a user provided function to update the augmented information on the
0044  * affected subtrees.
0045  */
0046 static inline void
0047 rb_insert_augmented(struct rb_node *node, struct rb_root *root,
0048             const struct rb_augment_callbacks *augment)
0049 {
0050     __rb_insert_augmented(node, root, augment->rotate);
0051 }
0052
0053 static inline void
0054 rb_insert_augmented_cached(struct rb_node *node,
0055                struct rb_root_cached *root, bool newleft,
0056                const struct rb_augment_callbacks *augment)
0057 {
0058     if (newleft)
0059         root->rb_leftmost = node;
0060     rb_insert_augmented(node, &root->rb_root, augment);
0061 }
0062
0063 /*
0064  * Template for declaring augmented rbtree callbacks (generic case)
0065  *
0066  * RBSTATIC:    'static' or empty
0067  * RBNAME:      name of the rb_augment_callbacks structure
0068  * RBSTRUCT:    struct type of the tree nodes
0069  * RBFIELD:     name of struct rb_node field within RBSTRUCT
0070  * RBAUGMENTED: name of field within RBSTRUCT holding data for subtree
0071  * RBCOMPUTE:   name of function that recomputes the RBAUGMENTED data
0072  */
0073
0074 #define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,              \
0075                  RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE) \
0076 static inline void                          \
0077 RBNAME ## _propagate(struct rb_node *rb, struct rb_node *stop)      \
0078 {                                   \
0079     while (rb != stop) {                        \
0080         RBSTRUCT *node = rb_entry(rb, RBSTRUCT, RBFIELD);   \
0081         if (RBCOMPUTE(node, true))              \
0082             break;                      \
0083         rb = rb_parent(&node->RBFIELD);             \
0084     }                               \
0085 }                                   \
0086 static inline void                          \
0087 RBNAME ## _copy(struct rb_node *rb_old, struct rb_node *rb_new)     \
0088 {                                   \
0089     RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD);        \
0090     RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD);        \
0091     new->RBAUGMENTED = old->RBAUGMENTED;                \
0092 }                                   \
0093 static void                             \
0094 RBNAME ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)   \
0095 {                                   \
0096     RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD);        \
0097     RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD);        \
0098     new->RBAUGMENTED = old->RBAUGMENTED;                \
0099     RBCOMPUTE(old, false);                      \
0100 }                                   \
0101 RBSTATIC const struct rb_augment_callbacks RBNAME = {           \
0102     .propagate = RBNAME ## _propagate,              \
0103     .copy = RBNAME ## _copy,                    \
0104     .rotate = RBNAME ## _rotate                 \
0105 };
0106
0107 /*
0108  * Template for declaring augmented rbtree callbacks,
0109  * computing RBAUGMENTED scalar as max(RBCOMPUTE(node)) for all subtree nodes.
0110  *
0111  * RBSTATIC:    'static' or empty
0112  * RBNAME:      name of the rb_augment_callbacks structure
0113  * RBSTRUCT:    struct type of the tree nodes
0114  * RBFIELD:     name of struct rb_node field within RBSTRUCT
0115  * RBTYPE:      type of the RBAUGMENTED field
0116  * RBAUGMENTED: name of RBTYPE field within RBSTRUCT holding data for subtree
0117  * RBCOMPUTE:   name of function that returns the per-node RBTYPE scalar
0118  */
0119
0120 #define RB_DECLARE_CALLBACKS_MAX(RBSTATIC, RBNAME, RBSTRUCT, RBFIELD,         \
0121                  RBTYPE, RBAUGMENTED, RBCOMPUTE)          \
0122 static inline bool RBNAME ## _compute_max(RBSTRUCT *node, bool exit)          \
0123 {                                         \
0124     RBSTRUCT *child;                              \
0125     RBTYPE max = RBCOMPUTE(node);                         \
0126     if (node->RBFIELD.rb_left) {                          \
0127         child = rb_entry(node->RBFIELD.rb_left, RBSTRUCT, RBFIELD);   \
0128         if (child->RBAUGMENTED > max)                     \
0129             max = child->RBAUGMENTED;                 \
0130     }                                     \
0131     if (node->RBFIELD.rb_right) {                         \
0132         child = rb_entry(node->RBFIELD.rb_right, RBSTRUCT, RBFIELD);  \
0133         if (child->RBAUGMENTED > max)                     \
0134             max = child->RBAUGMENTED;                 \
0135     }                                     \
0136     if (exit && node->RBAUGMENTED == max)                     \
0137         return true;                              \
0138     node->RBAUGMENTED = max;                          \
0139     return false;                                 \
0140 }                                         \
0141 RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,                        \
0142              RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max)
0143
0144
0145 #define RB_RED      0
0146 #define RB_BLACK    1
0147
0148 #define __rb_parent(pc)    ((struct rb_node *)(pc & ~3))
0149
0150 #define __rb_color(pc)     ((pc) & 1)
0151 #define __rb_is_black(pc)  __rb_color(pc)
0152 #define __rb_is_red(pc)    (!__rb_color(pc))
0153 #define rb_color(rb)       __rb_color((rb)->__rb_parent_color)
0154 #define rb_is_red(rb)      __rb_is_red((rb)->__rb_parent_color)
0155 #define rb_is_black(rb)    __rb_is_black((rb)->__rb_parent_color)
0156
0157 static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
0158 {
0159     rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
0160 }
0161
0162 static inline void rb_set_parent_color(struct rb_node *rb,
0163                        struct rb_node *p, int color)
0164 {
0165     rb->__rb_parent_color = (unsigned long)p | color;
0166 }
0167
0168 static inline void
0169 __rb_change_child(struct rb_node *old, struct rb_node *new,
0170           struct rb_node *parent, struct rb_root *root)
0171 {
0172     if (parent) {
0173         if (parent->rb_left == old)
0174             WRITE_ONCE(parent->rb_left, new);
0175         else
0176             WRITE_ONCE(parent->rb_right, new);
0177     } else
0178         WRITE_ONCE(root->rb_node, new);
0179 }
0180
0181 static inline void
0182 __rb_change_child_rcu(struct rb_node *old, struct rb_node *new,
0183               struct rb_node *parent, struct rb_root *root)
0184 {
0185     if (parent) {
0186         if (parent->rb_left == old)
0187             rcu_assign_pointer(parent->rb_left, new);
0188         else
0189             rcu_assign_pointer(parent->rb_right, new);
0190     } else
0191         rcu_assign_pointer(root->rb_node, new);
0192 }
0193
0194 extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
0195     void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
0196
0197 static __always_inline struct rb_node *
0198 __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
0199              const struct rb_augment_callbacks *augment)
0200 {
0201     struct rb_node *child = node->rb_right;
0202     struct rb_node *tmp = node->rb_left;
0203     struct rb_node *parent, *rebalance;
0204     unsigned long pc;
0205
0206     if (!tmp) {
0207         /*
0208          * Case 1: node to erase has no more than 1 child (easy!)
0209          *
0210          * Note that if there is one child it must be red due to 5)
0211          * and node must be black due to 4). We adjust colors locally
0212          * so as to bypass __rb_erase_color() later on.
0213          */
0214         pc = node->__rb_parent_color;
0215         parent = __rb_parent(pc);
0216         __rb_change_child(node, child, parent, root);
0217         if (child) {
0218             child->__rb_parent_color = pc;
0219             rebalance = NULL;
0220         } else
0221             rebalance = __rb_is_black(pc) ? parent : NULL;
0222         tmp = parent;
0223     } else if (!child) {
0224         /* Still case 1, but this time the child is node->rb_left */
0225         tmp->__rb_parent_color = pc = node->__rb_parent_color;
0226         parent = __rb_parent(pc);
0227         __rb_change_child(node, tmp, parent, root);
0228         rebalance = NULL;
0229         tmp = parent;
0230     } else {
0231         struct rb_node *successor = child, *child2;
0232
0233         tmp = child->rb_left;
0234         if (!tmp) {
0235             /*
0236              * Case 2: node's successor is its right child
0237              *
0238              *    (n)          (s)
0239              *    / \          / \
0240              *  (x) (s)  ->  (x) (c)
0241              *        \
0242              *        (c)
0243              */
0244             parent = successor;
0245             child2 = successor->rb_right;
0246
0247             augment->copy(node, successor);
0248         } else {
0249             /*
0250              * Case 3: node's successor is leftmost under
0251              * node's right child subtree
0252              *
0253              *    (n)          (s)
0254              *    / \          / \
0255              *  (x) (y)  ->  (x) (y)
0256              *      /            /
0257              *    (p)          (p)
0258              *    /            /
0259              *  (s)          (c)
0260              *    \
0261              *    (c)
0262              */
0263             do {
0264                 parent = successor;
0265                 successor = tmp;
0266                 tmp = tmp->rb_left;
0267             } while (tmp);
0268             child2 = successor->rb_right;
0269             WRITE_ONCE(parent->rb_left, child2);
0270             WRITE_ONCE(successor->rb_right, child);
0271             rb_set_parent(child, successor);
0272
0273             augment->copy(node, successor);
0274             augment->propagate(parent, successor);
0275         }
0276
0277         tmp = node->rb_left;
0278         WRITE_ONCE(successor->rb_left, tmp);
0279         rb_set_parent(tmp, successor);
0280
0281         pc = node->__rb_parent_color;
0282         tmp = __rb_parent(pc);
0283         __rb_change_child(node, successor, tmp, root);
0284
0285         if (child2) {
0286             rb_set_parent_color(child2, parent, RB_BLACK);
0287             rebalance = NULL;
0288         } else {
0289             rebalance = rb_is_black(successor) ? parent : NULL;
0290         }
0291         successor->__rb_parent_color = pc;
0292         tmp = successor;
0293     }
0294
0295     augment->propagate(tmp, NULL);
0296     return rebalance;
0297 }
0298
0299 static __always_inline void
0300 rb_erase_augmented(struct rb_node *node, struct rb_root *root,
0301            const struct rb_augment_callbacks *augment)
0302 {
0303     struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
0304     if (rebalance)
0305         __rb_erase_color(rebalance, root, augment->rotate);
0306 }
0307
0308 static __always_inline void
0309 rb_erase_augmented_cached(struct rb_node *node, struct rb_root_cached *root,
0310               const struct rb_augment_callbacks *augment)
0311 {
0312     if (root->rb_leftmost == node)
0313         root->rb_leftmost = rb_next(node);
0314     rb_erase_augmented(node, &root->rb_root, augment);
0315 }
0316
0317 #endif  /* _LINUX_RBTREE_AUGMENTED_H */