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