Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 1 | /* SPDX-License-Identifier: GPL-2.0-or-later */ |
| 2 | |
| 3 | #include "qemu/osdep.h" |
| 4 | #include "qemu/interval-tree.h" |
| 5 | #include "qemu/atomic.h" |
| 6 | |
| 7 | /* |
| 8 | * Red Black Trees. |
| 9 | * |
| 10 | * For now, don't expose Linux Red-Black Trees separately, but retain the |
| 11 | * separate type definitions to keep the implementation sane, and allow |
| 12 | * the possibility of separating them later. |
| 13 | * |
| 14 | * Derived from include/linux/rbtree_augmented.h and its dependencies. |
| 15 | */ |
| 16 | |
| 17 | /* |
| 18 | * red-black trees properties: https://en.wikipedia.org/wiki/Rbtree |
| 19 | * |
| 20 | * 1) A node is either red or black |
| 21 | * 2) The root is black |
| 22 | * 3) All leaves (NULL) are black |
| 23 | * 4) Both children of every red node are black |
| 24 | * 5) Every simple path from root to leaves contains the same number |
| 25 | * of black nodes. |
| 26 | * |
| 27 | * 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two |
| 28 | * consecutive red nodes in a path and every red node is therefore followed by |
| 29 | * a black. So if B is the number of black nodes on every simple path (as per |
| 30 | * 5), then the longest possible path due to 4 is 2B. |
| 31 | * |
| 32 | * We shall indicate color with case, where black nodes are uppercase and red |
| 33 | * nodes will be lowercase. Unknown color nodes shall be drawn as red within |
| 34 | * parentheses and have some accompanying text comment. |
| 35 | * |
| 36 | * Notes on lockless lookups: |
| 37 | * |
| 38 | * All stores to the tree structure (rb_left and rb_right) must be done using |
| 39 | * WRITE_ONCE [qatomic_set for QEMU]. And we must not inadvertently cause |
| 40 | * (temporary) loops in the tree structure as seen in program order. |
| 41 | * |
| 42 | * These two requirements will allow lockless iteration of the tree -- not |
| 43 | * correct iteration mind you, tree rotations are not atomic so a lookup might |
| 44 | * miss entire subtrees. |
| 45 | * |
| 46 | * But they do guarantee that any such traversal will only see valid elements |
| 47 | * and that it will indeed complete -- does not get stuck in a loop. |
| 48 | * |
| 49 | * It also guarantees that if the lookup returns an element it is the 'correct' |
| 50 | * one. But not returning an element does _NOT_ mean it's not present. |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 51 | */ |
| 52 | |
| 53 | typedef enum RBColor |
| 54 | { |
| 55 | RB_RED, |
| 56 | RB_BLACK, |
| 57 | } RBColor; |
| 58 | |
| 59 | typedef struct RBAugmentCallbacks { |
| 60 | void (*propagate)(RBNode *node, RBNode *stop); |
| 61 | void (*copy)(RBNode *old, RBNode *new); |
| 62 | void (*rotate)(RBNode *old, RBNode *new); |
| 63 | } RBAugmentCallbacks; |
| 64 | |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 65 | static inline uintptr_t rb_pc(const RBNode *n) |
| 66 | { |
| 67 | return qatomic_read(&n->rb_parent_color); |
| 68 | } |
| 69 | |
| 70 | static inline void rb_set_pc(RBNode *n, uintptr_t pc) |
| 71 | { |
| 72 | qatomic_set(&n->rb_parent_color, pc); |
| 73 | } |
| 74 | |
Richard Henderson | d37a259 | 2023-07-22 21:20:05 +0100 | [diff] [blame] | 75 | static inline RBNode *pc_parent(uintptr_t pc) |
| 76 | { |
| 77 | return (RBNode *)(pc & ~1); |
| 78 | } |
| 79 | |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 80 | static inline RBNode *rb_parent(const RBNode *n) |
| 81 | { |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 82 | return pc_parent(rb_pc(n)); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 83 | } |
| 84 | |
| 85 | static inline RBNode *rb_red_parent(const RBNode *n) |
| 86 | { |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 87 | return (RBNode *)rb_pc(n); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 88 | } |
| 89 | |
| 90 | static inline RBColor pc_color(uintptr_t pc) |
| 91 | { |
| 92 | return (RBColor)(pc & 1); |
| 93 | } |
| 94 | |
| 95 | static inline bool pc_is_red(uintptr_t pc) |
| 96 | { |
| 97 | return pc_color(pc) == RB_RED; |
| 98 | } |
| 99 | |
| 100 | static inline bool pc_is_black(uintptr_t pc) |
| 101 | { |
| 102 | return !pc_is_red(pc); |
| 103 | } |
| 104 | |
| 105 | static inline RBColor rb_color(const RBNode *n) |
| 106 | { |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 107 | return pc_color(rb_pc(n)); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 108 | } |
| 109 | |
| 110 | static inline bool rb_is_red(const RBNode *n) |
| 111 | { |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 112 | return pc_is_red(rb_pc(n)); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 113 | } |
| 114 | |
| 115 | static inline bool rb_is_black(const RBNode *n) |
| 116 | { |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 117 | return pc_is_black(rb_pc(n)); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 118 | } |
| 119 | |
| 120 | static inline void rb_set_black(RBNode *n) |
| 121 | { |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 122 | rb_set_pc(n, rb_pc(n) | RB_BLACK); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 123 | } |
| 124 | |
| 125 | static inline void rb_set_parent_color(RBNode *n, RBNode *p, RBColor color) |
| 126 | { |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 127 | rb_set_pc(n, (uintptr_t)p | color); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 128 | } |
| 129 | |
| 130 | static inline void rb_set_parent(RBNode *n, RBNode *p) |
| 131 | { |
| 132 | rb_set_parent_color(n, p, rb_color(n)); |
| 133 | } |
| 134 | |
| 135 | static inline void rb_link_node(RBNode *node, RBNode *parent, RBNode **rb_link) |
| 136 | { |
| 137 | node->rb_parent_color = (uintptr_t)parent; |
| 138 | node->rb_left = node->rb_right = NULL; |
| 139 | |
Richard Henderson | 4c8baa0 | 2023-07-22 15:25:30 +0100 | [diff] [blame] | 140 | /* |
| 141 | * Ensure that node is initialized before insertion, |
| 142 | * as viewed by a concurrent search. |
| 143 | */ |
| 144 | qatomic_set_mb(rb_link, node); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 145 | } |
| 146 | |
| 147 | static RBNode *rb_next(RBNode *node) |
| 148 | { |
| 149 | RBNode *parent; |
| 150 | |
| 151 | /* OMIT: if empty node, return null. */ |
| 152 | |
| 153 | /* |
| 154 | * If we have a right-hand child, go down and then left as far as we can. |
| 155 | */ |
| 156 | if (node->rb_right) { |
| 157 | node = node->rb_right; |
| 158 | while (node->rb_left) { |
| 159 | node = node->rb_left; |
| 160 | } |
| 161 | return node; |
| 162 | } |
| 163 | |
| 164 | /* |
| 165 | * No right-hand children. Everything down and left is smaller than us, |
| 166 | * so any 'next' node must be in the general direction of our parent. |
| 167 | * Go up the tree; any time the ancestor is a right-hand child of its |
| 168 | * parent, keep going up. First time it's a left-hand child of its |
| 169 | * parent, said parent is our 'next' node. |
| 170 | */ |
| 171 | while ((parent = rb_parent(node)) && node == parent->rb_right) { |
| 172 | node = parent; |
| 173 | } |
| 174 | |
| 175 | return parent; |
| 176 | } |
| 177 | |
| 178 | static inline void rb_change_child(RBNode *old, RBNode *new, |
| 179 | RBNode *parent, RBRoot *root) |
| 180 | { |
| 181 | if (!parent) { |
| 182 | qatomic_set(&root->rb_node, new); |
| 183 | } else if (parent->rb_left == old) { |
| 184 | qatomic_set(&parent->rb_left, new); |
| 185 | } else { |
| 186 | qatomic_set(&parent->rb_right, new); |
| 187 | } |
| 188 | } |
| 189 | |
| 190 | static inline void rb_rotate_set_parents(RBNode *old, RBNode *new, |
| 191 | RBRoot *root, RBColor color) |
| 192 | { |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 193 | uintptr_t pc = rb_pc(old); |
| 194 | RBNode *parent = pc_parent(pc); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 195 | |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 196 | rb_set_pc(new, pc); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 197 | rb_set_parent_color(old, new, color); |
| 198 | rb_change_child(old, new, parent, root); |
| 199 | } |
| 200 | |
| 201 | static void rb_insert_augmented(RBNode *node, RBRoot *root, |
| 202 | const RBAugmentCallbacks *augment) |
| 203 | { |
| 204 | RBNode *parent = rb_red_parent(node), *gparent, *tmp; |
| 205 | |
| 206 | while (true) { |
| 207 | /* |
| 208 | * Loop invariant: node is red. |
| 209 | */ |
| 210 | if (unlikely(!parent)) { |
| 211 | /* |
| 212 | * The inserted node is root. Either this is the first node, or |
| 213 | * we recursed at Case 1 below and are no longer violating 4). |
| 214 | */ |
| 215 | rb_set_parent_color(node, NULL, RB_BLACK); |
| 216 | break; |
| 217 | } |
| 218 | |
| 219 | /* |
| 220 | * If there is a black parent, we are done. Otherwise, take some |
| 221 | * corrective action as, per 4), we don't want a red root or two |
| 222 | * consecutive red nodes. |
| 223 | */ |
| 224 | if (rb_is_black(parent)) { |
| 225 | break; |
| 226 | } |
| 227 | |
| 228 | gparent = rb_red_parent(parent); |
| 229 | |
| 230 | tmp = gparent->rb_right; |
| 231 | if (parent != tmp) { /* parent == gparent->rb_left */ |
| 232 | if (tmp && rb_is_red(tmp)) { |
| 233 | /* |
| 234 | * Case 1 - node's uncle is red (color flips). |
| 235 | * |
| 236 | * G g |
| 237 | * / \ / \ |
| 238 | * p u --> P U |
| 239 | * / / |
| 240 | * n n |
| 241 | * |
| 242 | * However, since g's parent might be red, and 4) does not |
| 243 | * allow this, we need to recurse at g. |
| 244 | */ |
| 245 | rb_set_parent_color(tmp, gparent, RB_BLACK); |
| 246 | rb_set_parent_color(parent, gparent, RB_BLACK); |
| 247 | node = gparent; |
| 248 | parent = rb_parent(node); |
| 249 | rb_set_parent_color(node, parent, RB_RED); |
| 250 | continue; |
| 251 | } |
| 252 | |
| 253 | tmp = parent->rb_right; |
| 254 | if (node == tmp) { |
| 255 | /* |
| 256 | * Case 2 - node's uncle is black and node is |
| 257 | * the parent's right child (left rotate at parent). |
| 258 | * |
| 259 | * G G |
| 260 | * / \ / \ |
| 261 | * p U --> n U |
| 262 | * \ / |
| 263 | * n p |
| 264 | * |
| 265 | * This still leaves us in violation of 4), the |
| 266 | * continuation into Case 3 will fix that. |
| 267 | */ |
| 268 | tmp = node->rb_left; |
| 269 | qatomic_set(&parent->rb_right, tmp); |
| 270 | qatomic_set(&node->rb_left, parent); |
| 271 | if (tmp) { |
| 272 | rb_set_parent_color(tmp, parent, RB_BLACK); |
| 273 | } |
| 274 | rb_set_parent_color(parent, node, RB_RED); |
| 275 | augment->rotate(parent, node); |
| 276 | parent = node; |
| 277 | tmp = node->rb_right; |
| 278 | } |
| 279 | |
| 280 | /* |
| 281 | * Case 3 - node's uncle is black and node is |
| 282 | * the parent's left child (right rotate at gparent). |
| 283 | * |
| 284 | * G P |
| 285 | * / \ / \ |
| 286 | * p U --> n g |
| 287 | * / \ |
| 288 | * n U |
| 289 | */ |
| 290 | qatomic_set(&gparent->rb_left, tmp); /* == parent->rb_right */ |
| 291 | qatomic_set(&parent->rb_right, gparent); |
| 292 | if (tmp) { |
| 293 | rb_set_parent_color(tmp, gparent, RB_BLACK); |
| 294 | } |
| 295 | rb_rotate_set_parents(gparent, parent, root, RB_RED); |
| 296 | augment->rotate(gparent, parent); |
| 297 | break; |
| 298 | } else { |
| 299 | tmp = gparent->rb_left; |
| 300 | if (tmp && rb_is_red(tmp)) { |
| 301 | /* Case 1 - color flips */ |
| 302 | rb_set_parent_color(tmp, gparent, RB_BLACK); |
| 303 | rb_set_parent_color(parent, gparent, RB_BLACK); |
| 304 | node = gparent; |
| 305 | parent = rb_parent(node); |
| 306 | rb_set_parent_color(node, parent, RB_RED); |
| 307 | continue; |
| 308 | } |
| 309 | |
| 310 | tmp = parent->rb_left; |
| 311 | if (node == tmp) { |
| 312 | /* Case 2 - right rotate at parent */ |
| 313 | tmp = node->rb_right; |
| 314 | qatomic_set(&parent->rb_left, tmp); |
| 315 | qatomic_set(&node->rb_right, parent); |
| 316 | if (tmp) { |
| 317 | rb_set_parent_color(tmp, parent, RB_BLACK); |
| 318 | } |
| 319 | rb_set_parent_color(parent, node, RB_RED); |
| 320 | augment->rotate(parent, node); |
| 321 | parent = node; |
| 322 | tmp = node->rb_left; |
| 323 | } |
| 324 | |
| 325 | /* Case 3 - left rotate at gparent */ |
| 326 | qatomic_set(&gparent->rb_right, tmp); /* == parent->rb_left */ |
| 327 | qatomic_set(&parent->rb_left, gparent); |
| 328 | if (tmp) { |
| 329 | rb_set_parent_color(tmp, gparent, RB_BLACK); |
| 330 | } |
| 331 | rb_rotate_set_parents(gparent, parent, root, RB_RED); |
| 332 | augment->rotate(gparent, parent); |
| 333 | break; |
| 334 | } |
| 335 | } |
| 336 | } |
| 337 | |
| 338 | static void rb_insert_augmented_cached(RBNode *node, |
| 339 | RBRootLeftCached *root, bool newleft, |
| 340 | const RBAugmentCallbacks *augment) |
| 341 | { |
| 342 | if (newleft) { |
| 343 | root->rb_leftmost = node; |
| 344 | } |
| 345 | rb_insert_augmented(node, &root->rb_root, augment); |
| 346 | } |
| 347 | |
| 348 | static void rb_erase_color(RBNode *parent, RBRoot *root, |
| 349 | const RBAugmentCallbacks *augment) |
| 350 | { |
| 351 | RBNode *node = NULL, *sibling, *tmp1, *tmp2; |
| 352 | |
| 353 | while (true) { |
| 354 | /* |
| 355 | * Loop invariants: |
| 356 | * - node is black (or NULL on first iteration) |
| 357 | * - node is not the root (parent is not NULL) |
| 358 | * - All leaf paths going through parent and node have a |
| 359 | * black node count that is 1 lower than other leaf paths. |
| 360 | */ |
| 361 | sibling = parent->rb_right; |
| 362 | if (node != sibling) { /* node == parent->rb_left */ |
| 363 | if (rb_is_red(sibling)) { |
| 364 | /* |
| 365 | * Case 1 - left rotate at parent |
| 366 | * |
| 367 | * P S |
| 368 | * / \ / \ |
| 369 | * N s --> p Sr |
| 370 | * / \ / \ |
| 371 | * Sl Sr N Sl |
| 372 | */ |
| 373 | tmp1 = sibling->rb_left; |
| 374 | qatomic_set(&parent->rb_right, tmp1); |
| 375 | qatomic_set(&sibling->rb_left, parent); |
| 376 | rb_set_parent_color(tmp1, parent, RB_BLACK); |
| 377 | rb_rotate_set_parents(parent, sibling, root, RB_RED); |
| 378 | augment->rotate(parent, sibling); |
| 379 | sibling = tmp1; |
| 380 | } |
| 381 | tmp1 = sibling->rb_right; |
| 382 | if (!tmp1 || rb_is_black(tmp1)) { |
| 383 | tmp2 = sibling->rb_left; |
| 384 | if (!tmp2 || rb_is_black(tmp2)) { |
| 385 | /* |
| 386 | * Case 2 - sibling color flip |
| 387 | * (p could be either color here) |
| 388 | * |
| 389 | * (p) (p) |
| 390 | * / \ / \ |
| 391 | * N S --> N s |
| 392 | * / \ / \ |
| 393 | * Sl Sr Sl Sr |
| 394 | * |
| 395 | * This leaves us violating 5) which |
| 396 | * can be fixed by flipping p to black |
| 397 | * if it was red, or by recursing at p. |
| 398 | * p is red when coming from Case 1. |
| 399 | */ |
| 400 | rb_set_parent_color(sibling, parent, RB_RED); |
| 401 | if (rb_is_red(parent)) { |
| 402 | rb_set_black(parent); |
| 403 | } else { |
| 404 | node = parent; |
| 405 | parent = rb_parent(node); |
| 406 | if (parent) { |
| 407 | continue; |
| 408 | } |
| 409 | } |
| 410 | break; |
| 411 | } |
| 412 | /* |
| 413 | * Case 3 - right rotate at sibling |
| 414 | * (p could be either color here) |
| 415 | * |
| 416 | * (p) (p) |
| 417 | * / \ / \ |
| 418 | * N S --> N sl |
| 419 | * / \ \ |
| 420 | * sl Sr S |
| 421 | * \ |
| 422 | * Sr |
| 423 | * |
| 424 | * Note: p might be red, and then bot |
| 425 | * p and sl are red after rotation (which |
| 426 | * breaks property 4). This is fixed in |
| 427 | * Case 4 (in rb_rotate_set_parents() |
| 428 | * which set sl the color of p |
| 429 | * and set p RB_BLACK) |
| 430 | * |
| 431 | * (p) (sl) |
| 432 | * / \ / \ |
| 433 | * N sl --> P S |
| 434 | * \ / \ |
| 435 | * S N Sr |
| 436 | * \ |
| 437 | * Sr |
| 438 | */ |
| 439 | tmp1 = tmp2->rb_right; |
| 440 | qatomic_set(&sibling->rb_left, tmp1); |
| 441 | qatomic_set(&tmp2->rb_right, sibling); |
| 442 | qatomic_set(&parent->rb_right, tmp2); |
| 443 | if (tmp1) { |
| 444 | rb_set_parent_color(tmp1, sibling, RB_BLACK); |
| 445 | } |
| 446 | augment->rotate(sibling, tmp2); |
| 447 | tmp1 = sibling; |
| 448 | sibling = tmp2; |
| 449 | } |
| 450 | /* |
| 451 | * Case 4 - left rotate at parent + color flips |
| 452 | * (p and sl could be either color here. |
| 453 | * After rotation, p becomes black, s acquires |
| 454 | * p's color, and sl keeps its color) |
| 455 | * |
| 456 | * (p) (s) |
| 457 | * / \ / \ |
| 458 | * N S --> P Sr |
| 459 | * / \ / \ |
| 460 | * (sl) sr N (sl) |
| 461 | */ |
| 462 | tmp2 = sibling->rb_left; |
| 463 | qatomic_set(&parent->rb_right, tmp2); |
| 464 | qatomic_set(&sibling->rb_left, parent); |
| 465 | rb_set_parent_color(tmp1, sibling, RB_BLACK); |
| 466 | if (tmp2) { |
| 467 | rb_set_parent(tmp2, parent); |
| 468 | } |
| 469 | rb_rotate_set_parents(parent, sibling, root, RB_BLACK); |
| 470 | augment->rotate(parent, sibling); |
| 471 | break; |
| 472 | } else { |
| 473 | sibling = parent->rb_left; |
| 474 | if (rb_is_red(sibling)) { |
| 475 | /* Case 1 - right rotate at parent */ |
| 476 | tmp1 = sibling->rb_right; |
| 477 | qatomic_set(&parent->rb_left, tmp1); |
| 478 | qatomic_set(&sibling->rb_right, parent); |
| 479 | rb_set_parent_color(tmp1, parent, RB_BLACK); |
| 480 | rb_rotate_set_parents(parent, sibling, root, RB_RED); |
| 481 | augment->rotate(parent, sibling); |
| 482 | sibling = tmp1; |
| 483 | } |
| 484 | tmp1 = sibling->rb_left; |
| 485 | if (!tmp1 || rb_is_black(tmp1)) { |
| 486 | tmp2 = sibling->rb_right; |
| 487 | if (!tmp2 || rb_is_black(tmp2)) { |
| 488 | /* Case 2 - sibling color flip */ |
| 489 | rb_set_parent_color(sibling, parent, RB_RED); |
| 490 | if (rb_is_red(parent)) { |
| 491 | rb_set_black(parent); |
| 492 | } else { |
| 493 | node = parent; |
| 494 | parent = rb_parent(node); |
| 495 | if (parent) { |
| 496 | continue; |
| 497 | } |
| 498 | } |
| 499 | break; |
| 500 | } |
| 501 | /* Case 3 - left rotate at sibling */ |
| 502 | tmp1 = tmp2->rb_left; |
| 503 | qatomic_set(&sibling->rb_right, tmp1); |
| 504 | qatomic_set(&tmp2->rb_left, sibling); |
| 505 | qatomic_set(&parent->rb_left, tmp2); |
| 506 | if (tmp1) { |
| 507 | rb_set_parent_color(tmp1, sibling, RB_BLACK); |
| 508 | } |
| 509 | augment->rotate(sibling, tmp2); |
| 510 | tmp1 = sibling; |
| 511 | sibling = tmp2; |
| 512 | } |
| 513 | /* Case 4 - right rotate at parent + color flips */ |
| 514 | tmp2 = sibling->rb_right; |
| 515 | qatomic_set(&parent->rb_left, tmp2); |
| 516 | qatomic_set(&sibling->rb_right, parent); |
| 517 | rb_set_parent_color(tmp1, sibling, RB_BLACK); |
| 518 | if (tmp2) { |
| 519 | rb_set_parent(tmp2, parent); |
| 520 | } |
| 521 | rb_rotate_set_parents(parent, sibling, root, RB_BLACK); |
| 522 | augment->rotate(parent, sibling); |
| 523 | break; |
| 524 | } |
| 525 | } |
| 526 | } |
| 527 | |
| 528 | static void rb_erase_augmented(RBNode *node, RBRoot *root, |
| 529 | const RBAugmentCallbacks *augment) |
| 530 | { |
| 531 | RBNode *child = node->rb_right; |
| 532 | RBNode *tmp = node->rb_left; |
| 533 | RBNode *parent, *rebalance; |
| 534 | uintptr_t pc; |
| 535 | |
| 536 | if (!tmp) { |
| 537 | /* |
| 538 | * Case 1: node to erase has no more than 1 child (easy!) |
| 539 | * |
| 540 | * Note that if there is one child it must be red due to 5) |
| 541 | * and node must be black due to 4). We adjust colors locally |
| 542 | * so as to bypass rb_erase_color() later on. |
| 543 | */ |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 544 | pc = rb_pc(node); |
Richard Henderson | d37a259 | 2023-07-22 21:20:05 +0100 | [diff] [blame] | 545 | parent = pc_parent(pc); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 546 | rb_change_child(node, child, parent, root); |
| 547 | if (child) { |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 548 | rb_set_pc(child, pc); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 549 | rebalance = NULL; |
| 550 | } else { |
| 551 | rebalance = pc_is_black(pc) ? parent : NULL; |
| 552 | } |
| 553 | tmp = parent; |
| 554 | } else if (!child) { |
| 555 | /* Still case 1, but this time the child is node->rb_left */ |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 556 | pc = rb_pc(node); |
Richard Henderson | d37a259 | 2023-07-22 21:20:05 +0100 | [diff] [blame] | 557 | parent = pc_parent(pc); |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 558 | rb_set_pc(tmp, pc); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 559 | rb_change_child(node, tmp, parent, root); |
| 560 | rebalance = NULL; |
| 561 | tmp = parent; |
| 562 | } else { |
| 563 | RBNode *successor = child, *child2; |
| 564 | tmp = child->rb_left; |
| 565 | if (!tmp) { |
| 566 | /* |
| 567 | * Case 2: node's successor is its right child |
| 568 | * |
| 569 | * (n) (s) |
| 570 | * / \ / \ |
| 571 | * (x) (s) -> (x) (c) |
| 572 | * \ |
| 573 | * (c) |
| 574 | */ |
| 575 | parent = successor; |
| 576 | child2 = successor->rb_right; |
| 577 | |
| 578 | augment->copy(node, successor); |
| 579 | } else { |
| 580 | /* |
| 581 | * Case 3: node's successor is leftmost under |
| 582 | * node's right child subtree |
| 583 | * |
| 584 | * (n) (s) |
| 585 | * / \ / \ |
| 586 | * (x) (y) -> (x) (y) |
| 587 | * / / |
| 588 | * (p) (p) |
| 589 | * / / |
| 590 | * (s) (c) |
| 591 | * \ |
| 592 | * (c) |
| 593 | */ |
| 594 | do { |
| 595 | parent = successor; |
| 596 | successor = tmp; |
| 597 | tmp = tmp->rb_left; |
| 598 | } while (tmp); |
| 599 | child2 = successor->rb_right; |
| 600 | qatomic_set(&parent->rb_left, child2); |
| 601 | qatomic_set(&successor->rb_right, child); |
| 602 | rb_set_parent(child, successor); |
| 603 | |
| 604 | augment->copy(node, successor); |
| 605 | augment->propagate(parent, successor); |
| 606 | } |
| 607 | |
| 608 | tmp = node->rb_left; |
| 609 | qatomic_set(&successor->rb_left, tmp); |
| 610 | rb_set_parent(tmp, successor); |
| 611 | |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 612 | pc = rb_pc(node); |
Richard Henderson | d37a259 | 2023-07-22 21:20:05 +0100 | [diff] [blame] | 613 | tmp = pc_parent(pc); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 614 | rb_change_child(node, successor, tmp, root); |
| 615 | |
| 616 | if (child2) { |
| 617 | rb_set_parent_color(child2, parent, RB_BLACK); |
| 618 | rebalance = NULL; |
| 619 | } else { |
| 620 | rebalance = rb_is_black(successor) ? parent : NULL; |
| 621 | } |
Richard Henderson | 79e2985 | 2023-07-22 15:17:38 +0100 | [diff] [blame] | 622 | rb_set_pc(successor, pc); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 623 | tmp = successor; |
| 624 | } |
| 625 | |
| 626 | augment->propagate(tmp, NULL); |
| 627 | |
| 628 | if (rebalance) { |
| 629 | rb_erase_color(rebalance, root, augment); |
| 630 | } |
| 631 | } |
| 632 | |
| 633 | static void rb_erase_augmented_cached(RBNode *node, RBRootLeftCached *root, |
| 634 | const RBAugmentCallbacks *augment) |
| 635 | { |
| 636 | if (root->rb_leftmost == node) { |
| 637 | root->rb_leftmost = rb_next(node); |
| 638 | } |
| 639 | rb_erase_augmented(node, &root->rb_root, augment); |
| 640 | } |
| 641 | |
| 642 | |
| 643 | /* |
| 644 | * Interval trees. |
| 645 | * |
| 646 | * Derived from lib/interval_tree.c and its dependencies, |
| 647 | * especially include/linux/interval_tree_generic.h. |
| 648 | */ |
| 649 | |
| 650 | #define rb_to_itree(N) container_of(N, IntervalTreeNode, rb) |
| 651 | |
| 652 | static bool interval_tree_compute_max(IntervalTreeNode *node, bool exit) |
| 653 | { |
| 654 | IntervalTreeNode *child; |
| 655 | uint64_t max = node->last; |
| 656 | |
| 657 | if (node->rb.rb_left) { |
| 658 | child = rb_to_itree(node->rb.rb_left); |
| 659 | if (child->subtree_last > max) { |
| 660 | max = child->subtree_last; |
| 661 | } |
| 662 | } |
| 663 | if (node->rb.rb_right) { |
| 664 | child = rb_to_itree(node->rb.rb_right); |
| 665 | if (child->subtree_last > max) { |
| 666 | max = child->subtree_last; |
| 667 | } |
| 668 | } |
| 669 | if (exit && node->subtree_last == max) { |
| 670 | return true; |
| 671 | } |
| 672 | node->subtree_last = max; |
| 673 | return false; |
| 674 | } |
| 675 | |
| 676 | static void interval_tree_propagate(RBNode *rb, RBNode *stop) |
| 677 | { |
| 678 | while (rb != stop) { |
| 679 | IntervalTreeNode *node = rb_to_itree(rb); |
| 680 | if (interval_tree_compute_max(node, true)) { |
| 681 | break; |
| 682 | } |
| 683 | rb = rb_parent(&node->rb); |
| 684 | } |
| 685 | } |
| 686 | |
| 687 | static void interval_tree_copy(RBNode *rb_old, RBNode *rb_new) |
| 688 | { |
| 689 | IntervalTreeNode *old = rb_to_itree(rb_old); |
| 690 | IntervalTreeNode *new = rb_to_itree(rb_new); |
| 691 | |
| 692 | new->subtree_last = old->subtree_last; |
| 693 | } |
| 694 | |
| 695 | static void interval_tree_rotate(RBNode *rb_old, RBNode *rb_new) |
| 696 | { |
| 697 | IntervalTreeNode *old = rb_to_itree(rb_old); |
| 698 | IntervalTreeNode *new = rb_to_itree(rb_new); |
| 699 | |
| 700 | new->subtree_last = old->subtree_last; |
| 701 | interval_tree_compute_max(old, false); |
| 702 | } |
| 703 | |
| 704 | static const RBAugmentCallbacks interval_tree_augment = { |
| 705 | .propagate = interval_tree_propagate, |
| 706 | .copy = interval_tree_copy, |
| 707 | .rotate = interval_tree_rotate, |
| 708 | }; |
| 709 | |
| 710 | /* Insert / remove interval nodes from the tree */ |
| 711 | void interval_tree_insert(IntervalTreeNode *node, IntervalTreeRoot *root) |
| 712 | { |
| 713 | RBNode **link = &root->rb_root.rb_node, *rb_parent = NULL; |
| 714 | uint64_t start = node->start, last = node->last; |
| 715 | IntervalTreeNode *parent; |
| 716 | bool leftmost = true; |
| 717 | |
| 718 | while (*link) { |
| 719 | rb_parent = *link; |
| 720 | parent = rb_to_itree(rb_parent); |
| 721 | |
| 722 | if (parent->subtree_last < last) { |
| 723 | parent->subtree_last = last; |
| 724 | } |
| 725 | if (start < parent->start) { |
| 726 | link = &parent->rb.rb_left; |
| 727 | } else { |
| 728 | link = &parent->rb.rb_right; |
| 729 | leftmost = false; |
| 730 | } |
| 731 | } |
| 732 | |
| 733 | node->subtree_last = last; |
| 734 | rb_link_node(&node->rb, rb_parent, link); |
| 735 | rb_insert_augmented_cached(&node->rb, root, leftmost, |
| 736 | &interval_tree_augment); |
| 737 | } |
| 738 | |
| 739 | void interval_tree_remove(IntervalTreeNode *node, IntervalTreeRoot *root) |
| 740 | { |
| 741 | rb_erase_augmented_cached(&node->rb, root, &interval_tree_augment); |
| 742 | } |
| 743 | |
| 744 | /* |
| 745 | * Iterate over intervals intersecting [start;last] |
| 746 | * |
| 747 | * Note that a node's interval intersects [start;last] iff: |
| 748 | * Cond1: node->start <= last |
| 749 | * and |
| 750 | * Cond2: start <= node->last |
| 751 | */ |
| 752 | |
| 753 | static IntervalTreeNode *interval_tree_subtree_search(IntervalTreeNode *node, |
| 754 | uint64_t start, |
| 755 | uint64_t last) |
| 756 | { |
| 757 | while (true) { |
| 758 | /* |
| 759 | * Loop invariant: start <= node->subtree_last |
| 760 | * (Cond2 is satisfied by one of the subtree nodes) |
| 761 | */ |
Richard Henderson | 055b86e | 2023-07-22 20:43:45 +0100 | [diff] [blame] | 762 | RBNode *tmp = qatomic_read(&node->rb.rb_left); |
| 763 | if (tmp) { |
| 764 | IntervalTreeNode *left = rb_to_itree(tmp); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 765 | |
| 766 | if (start <= left->subtree_last) { |
| 767 | /* |
| 768 | * Some nodes in left subtree satisfy Cond2. |
| 769 | * Iterate to find the leftmost such node N. |
| 770 | * If it also satisfies Cond1, that's the |
| 771 | * match we are looking for. Otherwise, there |
| 772 | * is no matching interval as nodes to the |
| 773 | * right of N can't satisfy Cond1 either. |
| 774 | */ |
| 775 | node = left; |
| 776 | continue; |
| 777 | } |
| 778 | } |
| 779 | if (node->start <= last) { /* Cond1 */ |
| 780 | if (start <= node->last) { /* Cond2 */ |
| 781 | return node; /* node is leftmost match */ |
| 782 | } |
Richard Henderson | 055b86e | 2023-07-22 20:43:45 +0100 | [diff] [blame] | 783 | tmp = qatomic_read(&node->rb.rb_right); |
| 784 | if (tmp) { |
| 785 | node = rb_to_itree(tmp); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 786 | if (start <= node->subtree_last) { |
| 787 | continue; |
| 788 | } |
| 789 | } |
| 790 | } |
| 791 | return NULL; /* no match */ |
| 792 | } |
| 793 | } |
| 794 | |
| 795 | IntervalTreeNode *interval_tree_iter_first(IntervalTreeRoot *root, |
| 796 | uint64_t start, uint64_t last) |
| 797 | { |
| 798 | IntervalTreeNode *node, *leftmost; |
| 799 | |
Helge Deller | 47d1e98 | 2023-08-09 17:11:39 +0200 | [diff] [blame] | 800 | if (!root || !root->rb_root.rb_node) { |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 801 | return NULL; |
| 802 | } |
| 803 | |
| 804 | /* |
| 805 | * Fastpath range intersection/overlap between A: [a0, a1] and |
| 806 | * B: [b0, b1] is given by: |
| 807 | * |
| 808 | * a0 <= b1 && b0 <= a1 |
| 809 | * |
| 810 | * ... where A holds the lock range and B holds the smallest |
| 811 | * 'start' and largest 'last' in the tree. For the later, we |
| 812 | * rely on the root node, which by augmented interval tree |
| 813 | * property, holds the largest value in its last-in-subtree. |
| 814 | * This allows mitigating some of the tree walk overhead for |
| 815 | * for non-intersecting ranges, maintained and consulted in O(1). |
| 816 | */ |
| 817 | node = rb_to_itree(root->rb_root.rb_node); |
| 818 | if (node->subtree_last < start) { |
| 819 | return NULL; |
| 820 | } |
| 821 | |
| 822 | leftmost = rb_to_itree(root->rb_leftmost); |
| 823 | if (leftmost->start > last) { |
| 824 | return NULL; |
| 825 | } |
| 826 | |
| 827 | return interval_tree_subtree_search(node, start, last); |
| 828 | } |
| 829 | |
| 830 | IntervalTreeNode *interval_tree_iter_next(IntervalTreeNode *node, |
| 831 | uint64_t start, uint64_t last) |
| 832 | { |
Richard Henderson | 055b86e | 2023-07-22 20:43:45 +0100 | [diff] [blame] | 833 | RBNode *rb, *prev; |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 834 | |
Richard Henderson | 055b86e | 2023-07-22 20:43:45 +0100 | [diff] [blame] | 835 | rb = qatomic_read(&node->rb.rb_right); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 836 | while (true) { |
| 837 | /* |
| 838 | * Loop invariants: |
| 839 | * Cond1: node->start <= last |
| 840 | * rb == node->rb.rb_right |
| 841 | * |
| 842 | * First, search right subtree if suitable |
| 843 | */ |
| 844 | if (rb) { |
| 845 | IntervalTreeNode *right = rb_to_itree(rb); |
| 846 | |
| 847 | if (start <= right->subtree_last) { |
| 848 | return interval_tree_subtree_search(right, start, last); |
| 849 | } |
| 850 | } |
| 851 | |
| 852 | /* Move up the tree until we come from a node's left child */ |
| 853 | do { |
| 854 | rb = rb_parent(&node->rb); |
| 855 | if (!rb) { |
| 856 | return NULL; |
| 857 | } |
| 858 | prev = &node->rb; |
| 859 | node = rb_to_itree(rb); |
Richard Henderson | 055b86e | 2023-07-22 20:43:45 +0100 | [diff] [blame] | 860 | rb = qatomic_read(&node->rb.rb_right); |
Richard Henderson | 0d99d37 | 2022-09-17 14:05:54 +0200 | [diff] [blame] | 861 | } while (prev == rb); |
| 862 | |
| 863 | /* Check if the node intersects [start;last] */ |
| 864 | if (last < node->start) { /* !Cond1 */ |
| 865 | return NULL; |
| 866 | } |
| 867 | if (start <= node->last) { /* Cond2 */ |
| 868 | return node; |
| 869 | } |
| 870 | } |
| 871 | } |
| 872 | |
| 873 | /* Occasionally useful for calling from within the debugger. */ |
| 874 | #if 0 |
| 875 | static void debug_interval_tree_int(IntervalTreeNode *node, |
| 876 | const char *dir, int level) |
| 877 | { |
| 878 | printf("%4d %*s %s [%" PRIu64 ",%" PRIu64 "] subtree_last:%" PRIu64 "\n", |
| 879 | level, level + 1, dir, rb_is_red(&node->rb) ? "r" : "b", |
| 880 | node->start, node->last, node->subtree_last); |
| 881 | |
| 882 | if (node->rb.rb_left) { |
| 883 | debug_interval_tree_int(rb_to_itree(node->rb.rb_left), "<", level + 1); |
| 884 | } |
| 885 | if (node->rb.rb_right) { |
| 886 | debug_interval_tree_int(rb_to_itree(node->rb.rb_right), ">", level + 1); |
| 887 | } |
| 888 | } |
| 889 | |
| 890 | void debug_interval_tree(IntervalTreeNode *node); |
| 891 | void debug_interval_tree(IntervalTreeNode *node) |
| 892 | { |
| 893 | if (node) { |
| 894 | debug_interval_tree_int(node, "*", 0); |
| 895 | } else { |
| 896 | printf("null\n"); |
| 897 | } |
| 898 | } |
| 899 | #endif |