rbtree_augmented.h 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308
  1. /* SPDX-License-Identifier: GPL-2.0-or-later */
  2. /*
  3. Red Black Trees
  4. (C) 1999 Andrea Arcangeli <[email protected]>
  5. (C) 2002 David Woodhouse <[email protected]>
  6. (C) 2012 Michel Lespinasse <[email protected]>
  7. tools/linux/include/linux/rbtree_augmented.h
  8. Copied from:
  9. linux/include/linux/rbtree_augmented.h
  10. */
  11. #ifndef _TOOLS_LINUX_RBTREE_AUGMENTED_H
  12. #define _TOOLS_LINUX_RBTREE_AUGMENTED_H
  13. #include <linux/compiler.h>
  14. #include <linux/rbtree.h>
  15. /*
  16. * Please note - only struct rb_augment_callbacks and the prototypes for
  17. * rb_insert_augmented() and rb_erase_augmented() are intended to be public.
  18. * The rest are implementation details you are not expected to depend on.
  19. *
  20. * See Documentation/core-api/rbtree.rst for documentation and samples.
  21. */
  22. struct rb_augment_callbacks {
  23. void (*propagate)(struct rb_node *node, struct rb_node *stop);
  24. void (*copy)(struct rb_node *old, struct rb_node *new);
  25. void (*rotate)(struct rb_node *old, struct rb_node *new);
  26. };
  27. extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
  28. void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
  29. /*
  30. * Fixup the rbtree and update the augmented information when rebalancing.
  31. *
  32. * On insertion, the user must update the augmented information on the path
  33. * leading to the inserted node, then call rb_link_node() as usual and
  34. * rb_insert_augmented() instead of the usual rb_insert_color() call.
  35. * If rb_insert_augmented() rebalances the rbtree, it will callback into
  36. * a user provided function to update the augmented information on the
  37. * affected subtrees.
  38. */
  39. static inline void
  40. rb_insert_augmented(struct rb_node *node, struct rb_root *root,
  41. const struct rb_augment_callbacks *augment)
  42. {
  43. __rb_insert_augmented(node, root, augment->rotate);
  44. }
  45. static inline void
  46. rb_insert_augmented_cached(struct rb_node *node,
  47. struct rb_root_cached *root, bool newleft,
  48. const struct rb_augment_callbacks *augment)
  49. {
  50. if (newleft)
  51. root->rb_leftmost = node;
  52. rb_insert_augmented(node, &root->rb_root, augment);
  53. }
  54. /*
  55. * Template for declaring augmented rbtree callbacks (generic case)
  56. *
  57. * RBSTATIC: 'static' or empty
  58. * RBNAME: name of the rb_augment_callbacks structure
  59. * RBSTRUCT: struct type of the tree nodes
  60. * RBFIELD: name of struct rb_node field within RBSTRUCT
  61. * RBAUGMENTED: name of field within RBSTRUCT holding data for subtree
  62. * RBCOMPUTE: name of function that recomputes the RBAUGMENTED data
  63. */
  64. #define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME, \
  65. RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE) \
  66. static inline void \
  67. RBNAME ## _propagate(struct rb_node *rb, struct rb_node *stop) \
  68. { \
  69. while (rb != stop) { \
  70. RBSTRUCT *node = rb_entry(rb, RBSTRUCT, RBFIELD); \
  71. if (RBCOMPUTE(node, true)) \
  72. break; \
  73. rb = rb_parent(&node->RBFIELD); \
  74. } \
  75. } \
  76. static inline void \
  77. RBNAME ## _copy(struct rb_node *rb_old, struct rb_node *rb_new) \
  78. { \
  79. RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD); \
  80. RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD); \
  81. new->RBAUGMENTED = old->RBAUGMENTED; \
  82. } \
  83. static void \
  84. RBNAME ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new) \
  85. { \
  86. RBSTRUCT *old = rb_entry(rb_old, RBSTRUCT, RBFIELD); \
  87. RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD); \
  88. new->RBAUGMENTED = old->RBAUGMENTED; \
  89. RBCOMPUTE(old, false); \
  90. } \
  91. RBSTATIC const struct rb_augment_callbacks RBNAME = { \
  92. .propagate = RBNAME ## _propagate, \
  93. .copy = RBNAME ## _copy, \
  94. .rotate = RBNAME ## _rotate \
  95. };
  96. /*
  97. * Template for declaring augmented rbtree callbacks,
  98. * computing RBAUGMENTED scalar as max(RBCOMPUTE(node)) for all subtree nodes.
  99. *
  100. * RBSTATIC: 'static' or empty
  101. * RBNAME: name of the rb_augment_callbacks structure
  102. * RBSTRUCT: struct type of the tree nodes
  103. * RBFIELD: name of struct rb_node field within RBSTRUCT
  104. * RBTYPE: type of the RBAUGMENTED field
  105. * RBAUGMENTED: name of RBTYPE field within RBSTRUCT holding data for subtree
  106. * RBCOMPUTE: name of function that returns the per-node RBTYPE scalar
  107. */
  108. #define RB_DECLARE_CALLBACKS_MAX(RBSTATIC, RBNAME, RBSTRUCT, RBFIELD, \
  109. RBTYPE, RBAUGMENTED, RBCOMPUTE) \
  110. static inline bool RBNAME ## _compute_max(RBSTRUCT *node, bool exit) \
  111. { \
  112. RBSTRUCT *child; \
  113. RBTYPE max = RBCOMPUTE(node); \
  114. if (node->RBFIELD.rb_left) { \
  115. child = rb_entry(node->RBFIELD.rb_left, RBSTRUCT, RBFIELD); \
  116. if (child->RBAUGMENTED > max) \
  117. max = child->RBAUGMENTED; \
  118. } \
  119. if (node->RBFIELD.rb_right) { \
  120. child = rb_entry(node->RBFIELD.rb_right, RBSTRUCT, RBFIELD); \
  121. if (child->RBAUGMENTED > max) \
  122. max = child->RBAUGMENTED; \
  123. } \
  124. if (exit && node->RBAUGMENTED == max) \
  125. return true; \
  126. node->RBAUGMENTED = max; \
  127. return false; \
  128. } \
  129. RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME, \
  130. RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max)
  131. #define RB_RED 0
  132. #define RB_BLACK 1
  133. #define __rb_parent(pc) ((struct rb_node *)(pc & ~3))
  134. #define __rb_color(pc) ((pc) & 1)
  135. #define __rb_is_black(pc) __rb_color(pc)
  136. #define __rb_is_red(pc) (!__rb_color(pc))
  137. #define rb_color(rb) __rb_color((rb)->__rb_parent_color)
  138. #define rb_is_red(rb) __rb_is_red((rb)->__rb_parent_color)
  139. #define rb_is_black(rb) __rb_is_black((rb)->__rb_parent_color)
  140. static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
  141. {
  142. rb->__rb_parent_color = rb_color(rb) | (unsigned long)p;
  143. }
  144. static inline void rb_set_parent_color(struct rb_node *rb,
  145. struct rb_node *p, int color)
  146. {
  147. rb->__rb_parent_color = (unsigned long)p | color;
  148. }
  149. static inline void
  150. __rb_change_child(struct rb_node *old, struct rb_node *new,
  151. struct rb_node *parent, struct rb_root *root)
  152. {
  153. if (parent) {
  154. if (parent->rb_left == old)
  155. WRITE_ONCE(parent->rb_left, new);
  156. else
  157. WRITE_ONCE(parent->rb_right, new);
  158. } else
  159. WRITE_ONCE(root->rb_node, new);
  160. }
  161. extern void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
  162. void (*augment_rotate)(struct rb_node *old, struct rb_node *new));
  163. static __always_inline struct rb_node *
  164. __rb_erase_augmented(struct rb_node *node, struct rb_root *root,
  165. const struct rb_augment_callbacks *augment)
  166. {
  167. struct rb_node *child = node->rb_right;
  168. struct rb_node *tmp = node->rb_left;
  169. struct rb_node *parent, *rebalance;
  170. unsigned long pc;
  171. if (!tmp) {
  172. /*
  173. * Case 1: node to erase has no more than 1 child (easy!)
  174. *
  175. * Note that if there is one child it must be red due to 5)
  176. * and node must be black due to 4). We adjust colors locally
  177. * so as to bypass __rb_erase_color() later on.
  178. */
  179. pc = node->__rb_parent_color;
  180. parent = __rb_parent(pc);
  181. __rb_change_child(node, child, parent, root);
  182. if (child) {
  183. child->__rb_parent_color = pc;
  184. rebalance = NULL;
  185. } else
  186. rebalance = __rb_is_black(pc) ? parent : NULL;
  187. tmp = parent;
  188. } else if (!child) {
  189. /* Still case 1, but this time the child is node->rb_left */
  190. tmp->__rb_parent_color = pc = node->__rb_parent_color;
  191. parent = __rb_parent(pc);
  192. __rb_change_child(node, tmp, parent, root);
  193. rebalance = NULL;
  194. tmp = parent;
  195. } else {
  196. struct rb_node *successor = child, *child2;
  197. tmp = child->rb_left;
  198. if (!tmp) {
  199. /*
  200. * Case 2: node's successor is its right child
  201. *
  202. * (n) (s)
  203. * / \ / \
  204. * (x) (s) -> (x) (c)
  205. * \
  206. * (c)
  207. */
  208. parent = successor;
  209. child2 = successor->rb_right;
  210. augment->copy(node, successor);
  211. } else {
  212. /*
  213. * Case 3: node's successor is leftmost under
  214. * node's right child subtree
  215. *
  216. * (n) (s)
  217. * / \ / \
  218. * (x) (y) -> (x) (y)
  219. * / /
  220. * (p) (p)
  221. * / /
  222. * (s) (c)
  223. * \
  224. * (c)
  225. */
  226. do {
  227. parent = successor;
  228. successor = tmp;
  229. tmp = tmp->rb_left;
  230. } while (tmp);
  231. child2 = successor->rb_right;
  232. WRITE_ONCE(parent->rb_left, child2);
  233. WRITE_ONCE(successor->rb_right, child);
  234. rb_set_parent(child, successor);
  235. augment->copy(node, successor);
  236. augment->propagate(parent, successor);
  237. }
  238. tmp = node->rb_left;
  239. WRITE_ONCE(successor->rb_left, tmp);
  240. rb_set_parent(tmp, successor);
  241. pc = node->__rb_parent_color;
  242. tmp = __rb_parent(pc);
  243. __rb_change_child(node, successor, tmp, root);
  244. if (child2) {
  245. successor->__rb_parent_color = pc;
  246. rb_set_parent_color(child2, parent, RB_BLACK);
  247. rebalance = NULL;
  248. } else {
  249. unsigned long pc2 = successor->__rb_parent_color;
  250. successor->__rb_parent_color = pc;
  251. rebalance = __rb_is_black(pc2) ? parent : NULL;
  252. }
  253. tmp = successor;
  254. }
  255. augment->propagate(tmp, NULL);
  256. return rebalance;
  257. }
  258. static __always_inline void
  259. rb_erase_augmented(struct rb_node *node, struct rb_root *root,
  260. const struct rb_augment_callbacks *augment)
  261. {
  262. struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
  263. if (rebalance)
  264. __rb_erase_color(rebalance, root, augment->rotate);
  265. }
  266. static __always_inline void
  267. rb_erase_augmented_cached(struct rb_node *node, struct rb_root_cached *root,
  268. const struct rb_augment_callbacks *augment)
  269. {
  270. if (root->rb_leftmost == node)
  271. root->rb_leftmost = rb_next(node);
  272. rb_erase_augmented(node, &root->rb_root, augment);
  273. }
  274. #endif /* _TOOLS_LINUX_RBTREE_AUGMENTED_H */