1*83d290c5STom Rini /* SPDX-License-Identifier: GPL-2.0+ */
29dd228b5SHeiko Schocher /*
39dd228b5SHeiko Schocher Red Black Trees
49dd228b5SHeiko Schocher (C) 1999 Andrea Arcangeli <andrea@suse.de>
59dd228b5SHeiko Schocher (C) 2002 David Woodhouse <dwmw2@infradead.org>
69dd228b5SHeiko Schocher (C) 2012 Michel Lespinasse <walken@google.com>
79dd228b5SHeiko Schocher
89dd228b5SHeiko Schocher linux/include/linux/rbtree_augmented.h
99dd228b5SHeiko Schocher */
109dd228b5SHeiko Schocher
119dd228b5SHeiko Schocher #ifndef _LINUX_RBTREE_AUGMENTED_H
129dd228b5SHeiko Schocher #define _LINUX_RBTREE_AUGMENTED_H
139dd228b5SHeiko Schocher
149dd228b5SHeiko Schocher #include <linux/compiler.h>
159dd228b5SHeiko Schocher #include <linux/rbtree.h>
169dd228b5SHeiko Schocher
179dd228b5SHeiko Schocher /*
189dd228b5SHeiko Schocher * Please note - only struct rb_augment_callbacks and the prototypes for
199dd228b5SHeiko Schocher * rb_insert_augmented() and rb_erase_augmented() are intended to be public.
209dd228b5SHeiko Schocher * The rest are implementation details you are not expected to depend on.
219dd228b5SHeiko Schocher *
229dd228b5SHeiko Schocher * See Documentation/rbtree.txt for documentation and samples.
239dd228b5SHeiko Schocher */
249dd228b5SHeiko Schocher
259dd228b5SHeiko Schocher struct rb_augment_callbacks {
269dd228b5SHeiko Schocher void (*propagate)(struct rb_node *node, struct rb_node *stop);
279dd228b5SHeiko Schocher void (*copy)(struct rb_node *old, struct rb_node *new);
289dd228b5SHeiko Schocher void (*rotate)(struct rb_node *old, struct rb_node *new);
299dd228b5SHeiko Schocher };
309dd228b5SHeiko Schocher
319dd228b5SHeiko Schocher extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
329dd228b5SHeiko Schocher void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
339dd228b5SHeiko Schocher static inline void
rb_insert_augmented(struct rb_node * node,struct rb_root * root,const struct rb_augment_callbacks * augment)349dd228b5SHeiko Schocher rb_insert_augmented(struct rb_node *node, struct rb_root *root,
359dd228b5SHeiko Schocher const struct rb_augment_callbacks *augment)
369dd228b5SHeiko Schocher {
379dd228b5SHeiko Schocher __rb_insert_augmented(node, root, augment->rotate);
389dd228b5SHeiko Schocher }
399dd228b5SHeiko Schocher
409dd228b5SHeiko Schocher #define RB_DECLARE_CALLBACKS(rbstatic, rbname, rbstruct, rbfield, \
419dd228b5SHeiko Schocher rbtype, rbaugmented, rbcompute) \
429dd228b5SHeiko Schocher static inline void \
439dd228b5SHeiko Schocher rbname ## _propagate(struct rb_node *rb, struct rb_node *stop) \
449dd228b5SHeiko Schocher { \
459dd228b5SHeiko Schocher while (rb != stop) { \
469dd228b5SHeiko Schocher rbstruct *node = rb_entry(rb, rbstruct, rbfield); \
479dd228b5SHeiko Schocher rbtype augmented = rbcompute(node); \
489dd228b5SHeiko Schocher if (node->rbaugmented == augmented) \
499dd228b5SHeiko Schocher break; \
509dd228b5SHeiko Schocher node->rbaugmented = augmented; \
519dd228b5SHeiko Schocher rb = rb_parent(&node->rbfield); \
529dd228b5SHeiko Schocher } \
539dd228b5SHeiko Schocher } \
549dd228b5SHeiko Schocher static inline void \
559dd228b5SHeiko Schocher rbname ## _copy(struct rb_node *rb_old, struct rb_node *rb_new) \
569dd228b5SHeiko Schocher { \
579dd228b5SHeiko Schocher rbstruct *old = rb_entry(rb_old, rbstruct, rbfield); \
589dd228b5SHeiko Schocher rbstruct *new = rb_entry(rb_new, rbstruct, rbfield); \
599dd228b5SHeiko Schocher new->rbaugmented = old->rbaugmented; \
609dd228b5SHeiko Schocher } \
619dd228b5SHeiko Schocher static void \
629dd228b5SHeiko Schocher rbname ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new) \
639dd228b5SHeiko Schocher { \
649dd228b5SHeiko Schocher rbstruct *old = rb_entry(rb_old, rbstruct, rbfield); \
659dd228b5SHeiko Schocher rbstruct *new = rb_entry(rb_new, rbstruct, rbfield); \
669dd228b5SHeiko Schocher new->rbaugmented = old->rbaugmented; \
679dd228b5SHeiko Schocher old->rbaugmented = rbcompute(old); \
689dd228b5SHeiko Schocher } \
699dd228b5SHeiko Schocher rbstatic const struct rb_augment_callbacks rbname = { \
709dd228b5SHeiko Schocher rbname ## _propagate, rbname ## _copy, rbname ## _rotate \
719dd228b5SHeiko Schocher };
729dd228b5SHeiko Schocher
739dd228b5SHeiko Schocher
749dd228b5SHeiko Schocher #define RB_RED 0
759dd228b5SHeiko Schocher #define RB_BLACK 1
769dd228b5SHeiko Schocher
779dd228b5SHeiko Schocher #define __rb_parent(pc) ((struct rb_node *)(pc & ~3))
789dd228b5SHeiko Schocher
799dd228b5SHeiko Schocher #define __rb_color(pc) ((pc) & 1)
809dd228b5SHeiko Schocher #define __rb_is_black(pc) __rb_color(pc)
819dd228b5SHeiko Schocher #define __rb_is_red(pc) (!__rb_color(pc))
829dd228b5SHeiko Schocher #define rb_color(rb) __rb_color((rb)->__rb_parent_color)
839dd228b5SHeiko Schocher #define rb_is_red(rb) __rb_is_red((rb)->__rb_parent_color)
849dd228b5SHeiko Schocher #define rb_is_black(rb) __rb_is_black((rb)->__rb_parent_color)
859dd228b5SHeiko Schocher
rb_set_parent(struct rb_node * rb,struct rb_node * p)869dd228b5SHeiko Schocher static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
879dd228b5SHeiko Schocher {
889dd228b5SHeiko Schocher rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
899dd228b5SHeiko Schocher }
909dd228b5SHeiko Schocher
rb_set_parent_color(struct rb_node * rb,struct rb_node * p,int color)919dd228b5SHeiko Schocher static inline void rb_set_parent_color(struct rb_node *rb,
929dd228b5SHeiko Schocher struct rb_node *p, int color)
939dd228b5SHeiko Schocher {
949dd228b5SHeiko Schocher rb->__rb_parent_color = (unsigned long)p | color;
959dd228b5SHeiko Schocher }
969dd228b5SHeiko Schocher
979dd228b5SHeiko Schocher static inline void
__rb_change_child(struct rb_node * old,struct rb_node * new,struct rb_node * parent,struct rb_root * root)989dd228b5SHeiko Schocher __rb_change_child(struct rb_node *old, struct rb_node *new,
999dd228b5SHeiko Schocher struct rb_node *parent, struct rb_root *root)
1009dd228b5SHeiko Schocher {
1019dd228b5SHeiko Schocher if (parent) {
1029dd228b5SHeiko Schocher if (parent->rb_left == old)
1039dd228b5SHeiko Schocher parent->rb_left = new;
1049dd228b5SHeiko Schocher else
1059dd228b5SHeiko Schocher parent->rb_right = new;
1069dd228b5SHeiko Schocher } else
1079dd228b5SHeiko Schocher root->rb_node = new;
1089dd228b5SHeiko Schocher }
1099dd228b5SHeiko Schocher
1109dd228b5SHeiko Schocher extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
1119dd228b5SHeiko Schocher void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
1129dd228b5SHeiko Schocher
1139dd228b5SHeiko Schocher static __always_inline struct rb_node *
__rb_erase_augmented(struct rb_node * node,struct rb_root * root,const struct rb_augment_callbacks * augment)1149dd228b5SHeiko Schocher __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
1159dd228b5SHeiko Schocher const struct rb_augment_callbacks *augment)
1169dd228b5SHeiko Schocher {
1179dd228b5SHeiko Schocher struct rb_node *child = node->rb_right, *tmp = node->rb_left;
1189dd228b5SHeiko Schocher struct rb_node *parent, *rebalance;
1199dd228b5SHeiko Schocher unsigned long pc;
1209dd228b5SHeiko Schocher
1219dd228b5SHeiko Schocher if (!tmp) {
1229dd228b5SHeiko Schocher /*
1239dd228b5SHeiko Schocher * Case 1: node to erase has no more than 1 child (easy!)
1249dd228b5SHeiko Schocher *
1259dd228b5SHeiko Schocher * Note that if there is one child it must be red due to 5)
1269dd228b5SHeiko Schocher * and node must be black due to 4). We adjust colors locally
1279dd228b5SHeiko Schocher * so as to bypass __rb_erase_color() later on.
1289dd228b5SHeiko Schocher */
1299dd228b5SHeiko Schocher pc = node->__rb_parent_color;
1309dd228b5SHeiko Schocher parent = __rb_parent(pc);
1319dd228b5SHeiko Schocher __rb_change_child(node, child, parent, root);
1329dd228b5SHeiko Schocher if (child) {
1339dd228b5SHeiko Schocher child->__rb_parent_color = pc;
1349dd228b5SHeiko Schocher rebalance = NULL;
1359dd228b5SHeiko Schocher } else
1369dd228b5SHeiko Schocher rebalance = __rb_is_black(pc) ? parent : NULL;
1379dd228b5SHeiko Schocher tmp = parent;
1389dd228b5SHeiko Schocher } else if (!child) {
1399dd228b5SHeiko Schocher /* Still case 1, but this time the child is node->rb_left */
1409dd228b5SHeiko Schocher tmp->__rb_parent_color = pc = node->__rb_parent_color;
1419dd228b5SHeiko Schocher parent = __rb_parent(pc);
1429dd228b5SHeiko Schocher __rb_change_child(node, tmp, parent, root);
1439dd228b5SHeiko Schocher rebalance = NULL;
1449dd228b5SHeiko Schocher tmp = parent;
1459dd228b5SHeiko Schocher } else {
1469dd228b5SHeiko Schocher struct rb_node *successor = child, *child2;
1479dd228b5SHeiko Schocher tmp = child->rb_left;
1489dd228b5SHeiko Schocher if (!tmp) {
1499dd228b5SHeiko Schocher /*
1509dd228b5SHeiko Schocher * Case 2: node's successor is its right child
1519dd228b5SHeiko Schocher *
1529dd228b5SHeiko Schocher * (n) (s)
1539dd228b5SHeiko Schocher * / \ / \
1549dd228b5SHeiko Schocher * (x) (s) -> (x) (c)
1559dd228b5SHeiko Schocher * \
1569dd228b5SHeiko Schocher * (c)
1579dd228b5SHeiko Schocher */
1589dd228b5SHeiko Schocher parent = successor;
1599dd228b5SHeiko Schocher child2 = successor->rb_right;
1609dd228b5SHeiko Schocher augment->copy(node, successor);
1619dd228b5SHeiko Schocher } else {
1629dd228b5SHeiko Schocher /*
1639dd228b5SHeiko Schocher * Case 3: node's successor is leftmost under
1649dd228b5SHeiko Schocher * node's right child subtree
1659dd228b5SHeiko Schocher *
1669dd228b5SHeiko Schocher * (n) (s)
1679dd228b5SHeiko Schocher * / \ / \
1689dd228b5SHeiko Schocher * (x) (y) -> (x) (y)
1699dd228b5SHeiko Schocher * / /
1709dd228b5SHeiko Schocher * (p) (p)
1719dd228b5SHeiko Schocher * / /
1729dd228b5SHeiko Schocher * (s) (c)
1739dd228b5SHeiko Schocher * \
1749dd228b5SHeiko Schocher * (c)
1759dd228b5SHeiko Schocher */
1769dd228b5SHeiko Schocher do {
1779dd228b5SHeiko Schocher parent = successor;
1789dd228b5SHeiko Schocher successor = tmp;
1799dd228b5SHeiko Schocher tmp = tmp->rb_left;
1809dd228b5SHeiko Schocher } while (tmp);
1819dd228b5SHeiko Schocher parent->rb_left = child2 = successor->rb_right;
1829dd228b5SHeiko Schocher successor->rb_right = child;
1839dd228b5SHeiko Schocher rb_set_parent(child, successor);
1849dd228b5SHeiko Schocher augment->copy(node, successor);
1859dd228b5SHeiko Schocher augment->propagate(parent, successor);
1869dd228b5SHeiko Schocher }
1879dd228b5SHeiko Schocher
1889dd228b5SHeiko Schocher successor->rb_left = tmp = node->rb_left;
1899dd228b5SHeiko Schocher rb_set_parent(tmp, successor);
1909dd228b5SHeiko Schocher
1919dd228b5SHeiko Schocher pc = node->__rb_parent_color;
1929dd228b5SHeiko Schocher tmp = __rb_parent(pc);
1939dd228b5SHeiko Schocher __rb_change_child(node, successor, tmp, root);
1949dd228b5SHeiko Schocher if (child2) {
1959dd228b5SHeiko Schocher successor->__rb_parent_color = pc;
1969dd228b5SHeiko Schocher rb_set_parent_color(child2, parent, RB_BLACK);
1979dd228b5SHeiko Schocher rebalance = NULL;
1989dd228b5SHeiko Schocher } else {
1999dd228b5SHeiko Schocher unsigned long pc2 = successor->__rb_parent_color;
2009dd228b5SHeiko Schocher successor->__rb_parent_color = pc;
2019dd228b5SHeiko Schocher rebalance = __rb_is_black(pc2) ? parent : NULL;
2029dd228b5SHeiko Schocher }
2039dd228b5SHeiko Schocher tmp = successor;
2049dd228b5SHeiko Schocher }
2059dd228b5SHeiko Schocher
2069dd228b5SHeiko Schocher augment->propagate(tmp, NULL);
2079dd228b5SHeiko Schocher return rebalance;
2089dd228b5SHeiko Schocher }
2099dd228b5SHeiko Schocher
2109dd228b5SHeiko Schocher static __always_inline void
rb_erase_augmented(struct rb_node * node,struct rb_root * root,const struct rb_augment_callbacks * augment)2119dd228b5SHeiko Schocher rb_erase_augmented(struct rb_node *node, struct rb_root *root,
2129dd228b5SHeiko Schocher const struct rb_augment_callbacks *augment)
2139dd228b5SHeiko Schocher {
2149dd228b5SHeiko Schocher struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
2159dd228b5SHeiko Schocher if (rebalance)
2169dd228b5SHeiko Schocher __rb_erase_color(rebalance, root, augment->rotate);
2179dd228b5SHeiko Schocher }
2189dd228b5SHeiko Schocher
2199dd228b5SHeiko Schocher #endif /* _LINUX_RBTREE_AUGMENTED_H */
220