rbtree.c 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597
  1. // SPDX-License-Identifier: GPL-2.0-or-later
  2. /*
  3. Red Black Trees
  4. (C) 1999 Andrea Arcangeli <andrea@suse.de>
  5. (C) 2002 David Woodhouse <dwmw2@infradead.org>
  6. (C) 2012 Michel Lespinasse <walken@google.com>
  7. linux/lib/rbtree.c
  8. */
  9. #include <linux/rbtree_augmented.h>
  10. #include <linux/export.h>
  11. /*
  12. * red-black trees properties: https://en.wikipedia.org/wiki/Rbtree
  13. *
  14. * 1) A node is either red or black
  15. * 2) The root is black
  16. * 3) All leaves (NULL) are black
  17. * 4) Both children of every red node are black
  18. * 5) Every simple path from root to leaves contains the same number
  19. * of black nodes.
  20. *
  21. * 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
  22. * consecutive red nodes in a path and every red node is therefore followed by
  23. * a black. So if B is the number of black nodes on every simple path (as per
  24. * 5), then the longest possible path due to 4 is 2B.
  25. *
  26. * We shall indicate color with case, where black nodes are uppercase and red
  27. * nodes will be lowercase. Unknown color nodes shall be drawn as red within
  28. * parentheses and have some accompanying text comment.
  29. */
  30. /*
  31. * Notes on lockless lookups:
  32. *
  33. * All stores to the tree structure (rb_left and rb_right) must be done using
  34. * WRITE_ONCE(). And we must not inadvertently cause (temporary) loops in the
  35. * tree structure as seen in program order.
  36. *
  37. * These two requirements will allow lockless iteration of the tree -- not
  38. * correct iteration mind you, tree rotations are not atomic so a lookup might
  39. * miss entire subtrees.
  40. *
  41. * But they do guarantee that any such traversal will only see valid elements
  42. * and that it will indeed complete -- does not get stuck in a loop.
  43. *
  44. * It also guarantees that if the lookup returns an element it is the 'correct'
  45. * one. But not returning an element does _NOT_ mean it's not present.
  46. *
  47. * NOTE:
  48. *
  49. * Stores to __rb_parent_color are not important for simple lookups so those
  50. * are left undone as of now. Nor did I check for loops involving parent
  51. * pointers.
  52. */
  53. static inline void rb_set_black(struct rb_node *rb)
  54. {
  55. rb->__rb_parent_color |= RB_BLACK;
  56. }
  57. static inline struct rb_node *rb_red_parent(struct rb_node *red)
  58. {
  59. return (struct rb_node *)red->__rb_parent_color;
  60. }
  61. /*
  62. * Helper function for rotations:
  63. * - old's parent and color get assigned to new
  64. * - old gets assigned new as a parent and 'color' as a color.
  65. */
  66. static inline void
  67. __rb_rotate_set_parents(struct rb_node *old, struct rb_node *new,
  68. struct rb_root *root, int color)
  69. {
  70. struct rb_node *parent = rb_parent(old);
  71. new->__rb_parent_color = old->__rb_parent_color;
  72. rb_set_parent_color(old, new, color);
  73. __rb_change_child(old, new, parent, root);
  74. }
  75. static __always_inline void
  76. __rb_insert(struct rb_node *node, struct rb_root *root,
  77. void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
  78. {
  79. struct rb_node *parent = rb_red_parent(node), *gparent, *tmp;
  80. while (true) {
  81. /*
  82. * Loop invariant: node is red.
  83. */
  84. if (unlikely(!parent)) {
  85. /*
  86. * The inserted node is root. Either this is the
  87. * first node, or we recursed at Case 1 below and
  88. * are no longer violating 4).
  89. */
  90. rb_set_parent_color(node, NULL, RB_BLACK);
  91. break;
  92. }
  93. /*
  94. * If there is a black parent, we are done.
  95. * Otherwise, take some corrective action as,
  96. * per 4), we don't want a red root or two
  97. * consecutive red nodes.
  98. */
  99. if(rb_is_black(parent))
  100. break;
  101. gparent = rb_red_parent(parent);
  102. tmp = gparent->rb_right;
  103. if (parent != tmp) { /* parent == gparent->rb_left */
  104. if (tmp && rb_is_red(tmp)) {
  105. /*
  106. * Case 1 - node's uncle is red (color flips).
  107. *
  108. * G g
  109. * / \ / \
  110. * p u --> P U
  111. * / /
  112. * n n
  113. *
  114. * However, since g's parent might be red, and
  115. * 4) does not allow this, we need to recurse
  116. * at g.
  117. */
  118. rb_set_parent_color(tmp, gparent, RB_BLACK);
  119. rb_set_parent_color(parent, gparent, RB_BLACK);
  120. node = gparent;
  121. parent = rb_parent(node);
  122. rb_set_parent_color(node, parent, RB_RED);
  123. continue;
  124. }
  125. tmp = parent->rb_right;
  126. if (node == tmp) {
  127. /*
  128. * Case 2 - node's uncle is black and node is
  129. * the parent's right child (left rotate at parent).
  130. *
  131. * G G
  132. * / \ / \
  133. * p U --> n U
  134. * \ /
  135. * n p
  136. *
  137. * This still leaves us in violation of 4), the
  138. * continuation into Case 3 will fix that.
  139. */
  140. tmp = node->rb_left;
  141. WRITE_ONCE(parent->rb_right, tmp);
  142. WRITE_ONCE(node->rb_left, parent);
  143. if (tmp)
  144. rb_set_parent_color(tmp, parent,
  145. RB_BLACK);
  146. rb_set_parent_color(parent, node, RB_RED);
  147. augment_rotate(parent, node);
  148. parent = node;
  149. tmp = node->rb_right;
  150. }
  151. /*
  152. * Case 3 - node's uncle is black and node is
  153. * the parent's left child (right rotate at gparent).
  154. *
  155. * G P
  156. * / \ / \
  157. * p U --> n g
  158. * / \
  159. * n U
  160. */
  161. WRITE_ONCE(gparent->rb_left, tmp); /* == parent->rb_right */
  162. WRITE_ONCE(parent->rb_right, gparent);
  163. if (tmp)
  164. rb_set_parent_color(tmp, gparent, RB_BLACK);
  165. __rb_rotate_set_parents(gparent, parent, root, RB_RED);
  166. augment_rotate(gparent, parent);
  167. break;
  168. } else {
  169. tmp = gparent->rb_left;
  170. if (tmp && rb_is_red(tmp)) {
  171. /* Case 1 - color flips */
  172. rb_set_parent_color(tmp, gparent, RB_BLACK);
  173. rb_set_parent_color(parent, gparent, RB_BLACK);
  174. node = gparent;
  175. parent = rb_parent(node);
  176. rb_set_parent_color(node, parent, RB_RED);
  177. continue;
  178. }
  179. tmp = parent->rb_left;
  180. if (node == tmp) {
  181. /* Case 2 - right rotate at parent */
  182. tmp = node->rb_right;
  183. WRITE_ONCE(parent->rb_left, tmp);
  184. WRITE_ONCE(node->rb_right, parent);
  185. if (tmp)
  186. rb_set_parent_color(tmp, parent,
  187. RB_BLACK);
  188. rb_set_parent_color(parent, node, RB_RED);
  189. augment_rotate(parent, node);
  190. parent = node;
  191. tmp = node->rb_left;
  192. }
  193. /* Case 3 - left rotate at gparent */
  194. WRITE_ONCE(gparent->rb_right, tmp); /* == parent->rb_left */
  195. WRITE_ONCE(parent->rb_left, gparent);
  196. if (tmp)
  197. rb_set_parent_color(tmp, gparent, RB_BLACK);
  198. __rb_rotate_set_parents(gparent, parent, root, RB_RED);
  199. augment_rotate(gparent, parent);
  200. break;
  201. }
  202. }
  203. }
  204. /*
  205. * Inline version for rb_erase() use - we want to be able to inline
  206. * and eliminate the dummy_rotate callback there
  207. */
  208. static __always_inline void
  209. ____rb_erase_color(struct rb_node *parent, struct rb_root *root,
  210. void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
  211. {
  212. struct rb_node *node = NULL, *sibling, *tmp1, *tmp2;
  213. while (true) {
  214. /*
  215. * Loop invariants:
  216. * - node is black (or NULL on first iteration)
  217. * - node is not the root (parent is not NULL)
  218. * - All leaf paths going through parent and node have a
  219. * black node count that is 1 lower than other leaf paths.
  220. */
  221. sibling = parent->rb_right;
  222. if (node != sibling) { /* node == parent->rb_left */
  223. if (rb_is_red(sibling)) {
  224. /*
  225. * Case 1 - left rotate at parent
  226. *
  227. * P S
  228. * / \ / \
  229. * N s --> p Sr
  230. * / \ / \
  231. * Sl Sr N Sl
  232. */
  233. tmp1 = sibling->rb_left;
  234. WRITE_ONCE(parent->rb_right, tmp1);
  235. WRITE_ONCE(sibling->rb_left, parent);
  236. rb_set_parent_color(tmp1, parent, RB_BLACK);
  237. __rb_rotate_set_parents(parent, sibling, root,
  238. RB_RED);
  239. augment_rotate(parent, sibling);
  240. sibling = tmp1;
  241. }
  242. tmp1 = sibling->rb_right;
  243. if (!tmp1 || rb_is_black(tmp1)) {
  244. tmp2 = sibling->rb_left;
  245. if (!tmp2 || rb_is_black(tmp2)) {
  246. /*
  247. * Case 2 - sibling color flip
  248. * (p could be either color here)
  249. *
  250. * (p) (p)
  251. * / \ / \
  252. * N S --> N s
  253. * / \ / \
  254. * Sl Sr Sl Sr
  255. *
  256. * This leaves us violating 5) which
  257. * can be fixed by flipping p to black
  258. * if it was red, or by recursing at p.
  259. * p is red when coming from Case 1.
  260. */
  261. rb_set_parent_color(sibling, parent,
  262. RB_RED);
  263. if (rb_is_red(parent))
  264. rb_set_black(parent);
  265. else {
  266. node = parent;
  267. parent = rb_parent(node);
  268. if (parent)
  269. continue;
  270. }
  271. break;
  272. }
  273. /*
  274. * Case 3 - right rotate at sibling
  275. * (p could be either color here)
  276. *
  277. * (p) (p)
  278. * / \ / \
  279. * N S --> N sl
  280. * / \ \
  281. * sl Sr S
  282. * \
  283. * Sr
  284. *
  285. * Note: p might be red, and then both
  286. * p and sl are red after rotation(which
  287. * breaks property 4). This is fixed in
  288. * Case 4 (in __rb_rotate_set_parents()
  289. * which set sl the color of p
  290. * and set p RB_BLACK)
  291. *
  292. * (p) (sl)
  293. * / \ / \
  294. * N sl --> P S
  295. * \ / \
  296. * S N Sr
  297. * \
  298. * Sr
  299. */
  300. tmp1 = tmp2->rb_right;
  301. WRITE_ONCE(sibling->rb_left, tmp1);
  302. WRITE_ONCE(tmp2->rb_right, sibling);
  303. WRITE_ONCE(parent->rb_right, tmp2);
  304. if (tmp1)
  305. rb_set_parent_color(tmp1, sibling,
  306. RB_BLACK);
  307. augment_rotate(sibling, tmp2);
  308. tmp1 = sibling;
  309. sibling = tmp2;
  310. }
  311. /*
  312. * Case 4 - left rotate at parent + color flips
  313. * (p and sl could be either color here.
  314. * After rotation, p becomes black, s acquires
  315. * p's color, and sl keeps its color)
  316. *
  317. * (p) (s)
  318. * / \ / \
  319. * N S --> P Sr
  320. * / \ / \
  321. * (sl) sr N (sl)
  322. */
  323. tmp2 = sibling->rb_left;
  324. WRITE_ONCE(parent->rb_right, tmp2);
  325. WRITE_ONCE(sibling->rb_left, parent);
  326. rb_set_parent_color(tmp1, sibling, RB_BLACK);
  327. if (tmp2)
  328. rb_set_parent(tmp2, parent);
  329. __rb_rotate_set_parents(parent, sibling, root,
  330. RB_BLACK);
  331. augment_rotate(parent, sibling);
  332. break;
  333. } else {
  334. sibling = parent->rb_left;
  335. if (rb_is_red(sibling)) {
  336. /* Case 1 - right rotate at parent */
  337. tmp1 = sibling->rb_right;
  338. WRITE_ONCE(parent->rb_left, tmp1);
  339. WRITE_ONCE(sibling->rb_right, parent);
  340. rb_set_parent_color(tmp1, parent, RB_BLACK);
  341. __rb_rotate_set_parents(parent, sibling, root,
  342. RB_RED);
  343. augment_rotate(parent, sibling);
  344. sibling = tmp1;
  345. }
  346. tmp1 = sibling->rb_left;
  347. if (!tmp1 || rb_is_black(tmp1)) {
  348. tmp2 = sibling->rb_right;
  349. if (!tmp2 || rb_is_black(tmp2)) {
  350. /* Case 2 - sibling color flip */
  351. rb_set_parent_color(sibling, parent,
  352. RB_RED);
  353. if (rb_is_red(parent))
  354. rb_set_black(parent);
  355. else {
  356. node = parent;
  357. parent = rb_parent(node);
  358. if (parent)
  359. continue;
  360. }
  361. break;
  362. }
  363. /* Case 3 - left rotate at sibling */
  364. tmp1 = tmp2->rb_left;
  365. WRITE_ONCE(sibling->rb_right, tmp1);
  366. WRITE_ONCE(tmp2->rb_left, sibling);
  367. WRITE_ONCE(parent->rb_left, tmp2);
  368. if (tmp1)
  369. rb_set_parent_color(tmp1, sibling,
  370. RB_BLACK);
  371. augment_rotate(sibling, tmp2);
  372. tmp1 = sibling;
  373. sibling = tmp2;
  374. }
  375. /* Case 4 - right rotate at parent + color flips */
  376. tmp2 = sibling->rb_right;
  377. WRITE_ONCE(parent->rb_left, tmp2);
  378. WRITE_ONCE(sibling->rb_right, parent);
  379. rb_set_parent_color(tmp1, sibling, RB_BLACK);
  380. if (tmp2)
  381. rb_set_parent(tmp2, parent);
  382. __rb_rotate_set_parents(parent, sibling, root,
  383. RB_BLACK);
  384. augment_rotate(parent, sibling);
  385. break;
  386. }
  387. }
  388. }
  389. /* Non-inline version for rb_erase_augmented() use */
  390. void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
  391. void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
  392. {
  393. ____rb_erase_color(parent, root, augment_rotate);
  394. }
  395. /*
  396. * Non-augmented rbtree manipulation functions.
  397. *
  398. * We use dummy augmented callbacks here, and have the compiler optimize them
  399. * out of the rb_insert_color() and rb_erase() function definitions.
  400. */
  401. static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
  402. static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
  403. static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
  404. static const struct rb_augment_callbacks dummy_callbacks = {
  405. .propagate = dummy_propagate,
  406. .copy = dummy_copy,
  407. .rotate = dummy_rotate
  408. };
  409. void rb_insert_color(struct rb_node *node, struct rb_root *root)
  410. {
  411. __rb_insert(node, root, dummy_rotate);
  412. }
  413. void rb_erase(struct rb_node *node, struct rb_root *root)
  414. {
  415. struct rb_node *rebalance;
  416. rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
  417. if (rebalance)
  418. ____rb_erase_color(rebalance, root, dummy_rotate);
  419. }
  420. /*
  421. * Augmented rbtree manipulation functions.
  422. *
  423. * This instantiates the same __always_inline functions as in the non-augmented
  424. * case, but this time with user-defined callbacks.
  425. */
  426. void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
  427. void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
  428. {
  429. __rb_insert(node, root, augment_rotate);
  430. }
  431. /*
  432. * This function returns the first node (in sort order) of the tree.
  433. */
  434. struct rb_node *rb_first(const struct rb_root *root)
  435. {
  436. struct rb_node *n;
  437. n = root->rb_node;
  438. if (!n)
  439. return NULL;
  440. while (n->rb_left)
  441. n = n->rb_left;
  442. return n;
  443. }
  444. struct rb_node *rb_last(const struct rb_root *root)
  445. {
  446. struct rb_node *n;
  447. n = root->rb_node;
  448. if (!n)
  449. return NULL;
  450. while (n->rb_right)
  451. n = n->rb_right;
  452. return n;
  453. }
  454. struct rb_node *rb_next(const struct rb_node *node)
  455. {
  456. struct rb_node *parent;
  457. if (RB_EMPTY_NODE(node))
  458. return NULL;
  459. /*
  460. * If we have a right-hand child, go down and then left as far
  461. * as we can.
  462. */
  463. if (node->rb_right) {
  464. node = node->rb_right;
  465. while (node->rb_left)
  466. node = node->rb_left;
  467. return (struct rb_node *)node;
  468. }
  469. /*
  470. * No right-hand children. Everything down and left is smaller than us,
  471. * so any 'next' node must be in the general direction of our parent.
  472. * Go up the tree; any time the ancestor is a right-hand child of its
  473. * parent, keep going up. First time it's a left-hand child of its
  474. * parent, said parent is our 'next' node.
  475. */
  476. while ((parent = rb_parent(node)) && node == parent->rb_right)
  477. node = parent;
  478. return parent;
  479. }
  480. struct rb_node *rb_prev(const struct rb_node *node)
  481. {
  482. struct rb_node *parent;
  483. if (RB_EMPTY_NODE(node))
  484. return NULL;
  485. /*
  486. * If we have a left-hand child, go down and then right as far
  487. * as we can.
  488. */
  489. if (node->rb_left) {
  490. node = node->rb_left;
  491. while (node->rb_right)
  492. node = node->rb_right;
  493. return (struct rb_node *)node;
  494. }
  495. /*
  496. * No left-hand children. Go up till we find an ancestor which
  497. * is a right-hand child of its parent.
  498. */
  499. while ((parent = rb_parent(node)) && node == parent->rb_left)
  500. node = parent;
  501. return parent;
  502. }
  503. void rb_replace_node(struct rb_node *victim, struct rb_node *new,
  504. struct rb_root *root)
  505. {
  506. struct rb_node *parent = rb_parent(victim);
  507. /* Copy the pointers/colour from the victim to the replacement */
  508. *new = *victim;
  509. /* Set the surrounding nodes to point to the replacement */
  510. if (victim->rb_left)
  511. rb_set_parent(victim->rb_left, new);
  512. if (victim->rb_right)
  513. rb_set_parent(victim->rb_right, new);
  514. __rb_change_child(victim, new, parent, root);
  515. }
  516. static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
  517. {
  518. for (;;) {
  519. if (node->rb_left)
  520. node = node->rb_left;
  521. else if (node->rb_right)
  522. node = node->rb_right;
  523. else
  524. return (struct rb_node *)node;
  525. }
  526. }
  527. struct rb_node *rb_next_postorder(const struct rb_node *node)
  528. {
  529. const struct rb_node *parent;
  530. if (!node)
  531. return NULL;
  532. parent = rb_parent(node);
  533. /* If we're sitting on node, we've already seen our children */
  534. if (parent && node == parent->rb_left && parent->rb_right) {
  535. /* If we are the parent's left node, go to the parent's right
  536. * node then all the way down to the left */
  537. return rb_left_deepest_node(parent->rb_right);
  538. } else
  539. /* Otherwise we are the parent's right node, and the parent
  540. * should be next */
  541. return (struct rb_node *)parent;
  542. }
  543. struct rb_node *rb_first_postorder(const struct rb_root *root)
  544. {
  545. if (!root->rb_node)
  546. return NULL;
  547. return rb_left_deepest_node(root->rb_node);
  548. }