rbtree.c 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630
  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. EXPORT_SYMBOL(__rb_erase_color);
  396. /*
  397. * Non-augmented rbtree manipulation functions.
  398. *
  399. * We use dummy augmented callbacks here, and have the compiler optimize them
  400. * out of the rb_insert_color() and rb_erase() function definitions.
  401. */
  402. static inline void dummy_propagate(struct rb_node *node, struct rb_node *stop) {}
  403. static inline void dummy_copy(struct rb_node *old, struct rb_node *new) {}
  404. static inline void dummy_rotate(struct rb_node *old, struct rb_node *new) {}
  405. static const struct rb_augment_callbacks dummy_callbacks = {
  406. .propagate = dummy_propagate,
  407. .copy = dummy_copy,
  408. .rotate = dummy_rotate
  409. };
  410. void rb_insert_color(struct rb_node *node, struct rb_root *root)
  411. {
  412. __rb_insert(node, root, dummy_rotate);
  413. }
  414. EXPORT_SYMBOL(rb_insert_color);
  415. void rb_erase(struct rb_node *node, struct rb_root *root)
  416. {
  417. struct rb_node *rebalance;
  418. rebalance = __rb_erase_augmented(node, root, &dummy_callbacks);
  419. if (rebalance)
  420. ____rb_erase_color(rebalance, root, dummy_rotate);
  421. }
  422. EXPORT_SYMBOL(rb_erase);
  423. /*
  424. * Augmented rbtree manipulation functions.
  425. *
  426. * This instantiates the same __always_inline functions as in the non-augmented
  427. * case, but this time with user-defined callbacks.
  428. */
  429. void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
  430. void (*augment_rotate)(struct rb_node *old, struct rb_node *new))
  431. {
  432. __rb_insert(node, root, augment_rotate);
  433. }
  434. EXPORT_SYMBOL(__rb_insert_augmented);
  435. /*
  436. * This function returns the first node (in sort order) of the tree.
  437. */
  438. struct rb_node *rb_first(const struct rb_root *root)
  439. {
  440. struct rb_node *n;
  441. n = root->rb_node;
  442. if (!n)
  443. return NULL;
  444. while (n->rb_left)
  445. n = n->rb_left;
  446. return n;
  447. }
  448. EXPORT_SYMBOL(rb_first);
  449. struct rb_node *rb_last(const struct rb_root *root)
  450. {
  451. struct rb_node *n;
  452. n = root->rb_node;
  453. if (!n)
  454. return NULL;
  455. while (n->rb_right)
  456. n = n->rb_right;
  457. return n;
  458. }
  459. EXPORT_SYMBOL(rb_last);
  460. struct rb_node *rb_next(const struct rb_node *node)
  461. {
  462. struct rb_node *parent;
  463. if (RB_EMPTY_NODE(node))
  464. return NULL;
  465. /*
  466. * If we have a right-hand child, go down and then left as far
  467. * as we can.
  468. */
  469. if (node->rb_right) {
  470. node = node->rb_right;
  471. while (node->rb_left)
  472. node = node->rb_left;
  473. return (struct rb_node *)node;
  474. }
  475. /*
  476. * No right-hand children. Everything down and left is smaller than us,
  477. * so any 'next' node must be in the general direction of our parent.
  478. * Go up the tree; any time the ancestor is a right-hand child of its
  479. * parent, keep going up. First time it's a left-hand child of its
  480. * parent, said parent is our 'next' node.
  481. */
  482. while ((parent = rb_parent(node)) && node == parent->rb_right)
  483. node = parent;
  484. return parent;
  485. }
  486. EXPORT_SYMBOL(rb_next);
  487. struct rb_node *rb_prev(const struct rb_node *node)
  488. {
  489. struct rb_node *parent;
  490. if (RB_EMPTY_NODE(node))
  491. return NULL;
  492. /*
  493. * If we have a left-hand child, go down and then right as far
  494. * as we can.
  495. */
  496. if (node->rb_left) {
  497. node = node->rb_left;
  498. while (node->rb_right)
  499. node = node->rb_right;
  500. return (struct rb_node *)node;
  501. }
  502. /*
  503. * No left-hand children. Go up till we find an ancestor which
  504. * is a right-hand child of its parent.
  505. */
  506. while ((parent = rb_parent(node)) && node == parent->rb_left)
  507. node = parent;
  508. return parent;
  509. }
  510. EXPORT_SYMBOL(rb_prev);
  511. void rb_replace_node(struct rb_node *victim, struct rb_node *new,
  512. struct rb_root *root)
  513. {
  514. struct rb_node *parent = rb_parent(victim);
  515. /* Copy the pointers/colour from the victim to the replacement */
  516. *new = *victim;
  517. /* Set the surrounding nodes to point to the replacement */
  518. if (victim->rb_left)
  519. rb_set_parent(victim->rb_left, new);
  520. if (victim->rb_right)
  521. rb_set_parent(victim->rb_right, new);
  522. __rb_change_child(victim, new, parent, root);
  523. }
  524. EXPORT_SYMBOL(rb_replace_node);
  525. void rb_replace_node_rcu(struct rb_node *victim, struct rb_node *new,
  526. struct rb_root *root)
  527. {
  528. struct rb_node *parent = rb_parent(victim);
  529. /* Copy the pointers/colour from the victim to the replacement */
  530. *new = *victim;
  531. /* Set the surrounding nodes to point to the replacement */
  532. if (victim->rb_left)
  533. rb_set_parent(victim->rb_left, new);
  534. if (victim->rb_right)
  535. rb_set_parent(victim->rb_right, new);
  536. /* Set the parent's pointer to the new node last after an RCU barrier
  537. * so that the pointers onwards are seen to be set correctly when doing
  538. * an RCU walk over the tree.
  539. */
  540. __rb_change_child_rcu(victim, new, parent, root);
  541. }
  542. EXPORT_SYMBOL(rb_replace_node_rcu);
  543. static struct rb_node *rb_left_deepest_node(const struct rb_node *node)
  544. {
  545. for (;;) {
  546. if (node->rb_left)
  547. node = node->rb_left;
  548. else if (node->rb_right)
  549. node = node->rb_right;
  550. else
  551. return (struct rb_node *)node;
  552. }
  553. }
  554. struct rb_node *rb_next_postorder(const struct rb_node *node)
  555. {
  556. const struct rb_node *parent;
  557. if (!node)
  558. return NULL;
  559. parent = rb_parent(node);
  560. /* If we're sitting on node, we've already seen our children */
  561. if (parent && node == parent->rb_left && parent->rb_right) {
  562. /* If we are the parent's left node, go to the parent's right
  563. * node then all the way down to the left */
  564. return rb_left_deepest_node(parent->rb_right);
  565. } else
  566. /* Otherwise we are the parent's right node, and the parent
  567. * should be next */
  568. return (struct rb_node *)parent;
  569. }
  570. EXPORT_SYMBOL(rb_next_postorder);
  571. struct rb_node *rb_first_postorder(const struct rb_root *root)
  572. {
  573. if (!root->rb_node)
  574. return NULL;
  575. return rb_left_deepest_node(root->rb_node);
  576. }
  577. EXPORT_SYMBOL(rb_first_postorder);