rbtree.c

   1 /*
   2   SPDX-License-Identifier: GPL-2.0-or-later
   3
   4   Red Black Trees
   5   (C) 1999  Andrea Arcangeli <andrea@suse.de>
   6   (C) 2002  David Woodhouse <dwmw2@infradead.org>
   7
   8   linux/lib/rbtree.c
   9 */
  10
  11 #include "rbtree.h"
  12
  13 static void __rb_rotate_left(struct rb_node *node, struct rb_root *root)
  14 {
  15         struct rb_node *right = node->rb_right;
  16         struct rb_node *parent = rb_parent(node);
  17
  18         if ((node->rb_right = right->rb_left))
  19                 rb_set_parent(right->rb_left, node);
  20         right->rb_left = node;
  21
  22         rb_set_parent(right, parent);
  23
  24         if (parent)
  25         {
  26                 if (node == parent->rb_left)
  27                         parent->rb_left = right;
  28                 else
  29                         parent->rb_right = right;
  30         }
  31         else
  32                 root->rb_node = right;
  33         rb_set_parent(node, right);
  34 }
  35
  36 static void __rb_rotate_right(struct rb_node *node, struct rb_root *root)
  37 {
  38         struct rb_node *left = node->rb_left;
  39         struct rb_node *parent = rb_parent(node);
  40
  41         if ((node->rb_left = left->rb_right))
  42                 rb_set_parent(left->rb_right, node);
  43         left->rb_right = node;
  44
  45         rb_set_parent(left, parent);
  46
  47         if (parent)
  48         {
  49                 if (node == parent->rb_right)
  50                         parent->rb_right = left;
  51                 else
  52                         parent->rb_left = left;
  53         }
  54         else
  55                 root->rb_node = left;
  56         rb_set_parent(node, left);
  57 }
  58
  59 void rb_insert_color(struct rb_node *node, struct rb_root *root)
  60 {
  61         struct rb_node *parent, *gparent;
  62
  63         while ((parent = rb_parent(node)) && rb_is_red(parent))
  64         {
  65                 gparent = rb_parent(parent);
  66
  67                 if (parent == gparent->rb_left)
  68                 {
  69                         {
  70                                 register struct rb_node *uncle = gparent->rb_right;
  71                                 if (uncle && rb_is_red(uncle))
  72                                 {
  73                                         rb_set_black(uncle);
  74                                         rb_set_black(parent);
  75                                         rb_set_red(gparent);
  76                                         node = gparent;
  77                                         continue;
  78                                 }
  79                         }
  80
  81                         if (parent->rb_right == node)
  82                         {
  83                                 register struct rb_node *tmp;
  84                                 __rb_rotate_left(parent, root);
  85                                 tmp = parent;
  86                                 parent = node;
  87                                 node = tmp;
  88                         }
  89
  90                         rb_set_black(parent);
  91                         rb_set_red(gparent);
  92                         __rb_rotate_right(gparent, root);
  93                 } else {
  94                         {
  95                                 register struct rb_node *uncle = gparent->rb_left;
  96                                 if (uncle && rb_is_red(uncle))
  97                                 {
  98                                         rb_set_black(uncle);
  99                                         rb_set_black(parent);
 100                                         rb_set_red(gparent);
 101                                         node = gparent;
 102                                         continue;
 103                                 }
 104                         }
 105
 106                         if (parent->rb_left == node)
 107                         {
 108                                 register struct rb_node *tmp;
 109                                 __rb_rotate_right(parent, root);
 110                                 tmp = parent;
 111                                 parent = node;
 112                                 node = tmp;
 113                         }
 114
 115                         rb_set_black(parent);
 116                         rb_set_red(gparent);
 117                         __rb_rotate_left(gparent, root);
 118                 }
 119         }
 120
 121         rb_set_black(root->rb_node);
 122 }
 123
 124 static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
 125                              struct rb_root *root)
 126 {
 127         struct rb_node *other;
 128
 129         while ((!node || rb_is_black(node)) && node != root->rb_node)
 130         {
 131                 if (parent->rb_left == node)
 132                 {
 133                         other = parent->rb_right;
 134                         if (rb_is_red(other))
 135                         {
 136                                 rb_set_black(other);
 137                                 rb_set_red(parent);
 138                                 __rb_rotate_left(parent, root);
 139                                 other = parent->rb_right;
 140                         }
 141                         if ((!other->rb_left || rb_is_black(other->rb_left)) &&
 142                             (!other->rb_right || rb_is_black(other->rb_right)))
 143                         {
 144                                 rb_set_red(other);
 145                                 node = parent;
 146                                 parent = rb_parent(node);
 147                         }
 148                         else
 149                         {
 150                                 if (!other->rb_right || rb_is_black(other->rb_right))
 151                                 {
 152                                         rb_set_black(other->rb_left);
 153                                         rb_set_red(other);
 154                                         __rb_rotate_right(other, root);
 155                                         other = parent->rb_right;
 156                                 }
 157                                 rb_set_color(other, rb_color(parent));
 158                                 rb_set_black(parent);
 159                                 rb_set_black(other->rb_right);
 160                                 __rb_rotate_left(parent, root);
 161                                 node = root->rb_node;
 162                                 break;
 163                         }
 164                 }
 165                 else
 166                 {
 167                         other = parent->rb_left;
 168                         if (rb_is_red(other))
 169                         {
 170                                 rb_set_black(other);
 171                                 rb_set_red(parent);
 172                                 __rb_rotate_right(parent, root);
 173                                 other = parent->rb_left;
 174                         }
 175                         if ((!other->rb_left || rb_is_black(other->rb_left)) &&
 176                             (!other->rb_right || rb_is_black(other->rb_right)))
 177                         {
 178                                 rb_set_red(other);
 179                                 node = parent;
 180                                 parent = rb_parent(node);
 181                         }
 182                         else
 183                         {
 184                                 if (!other->rb_left || rb_is_black(other->rb_left))
 185                                 {
 186                                         rb_set_black(other->rb_right);
 187                                         rb_set_red(other);
 188                                         __rb_rotate_left(other, root);
 189                                         other = parent->rb_left;
 190                                 }
 191                                 rb_set_color(other, rb_color(parent));
 192                                 rb_set_black(parent);
 193                                 rb_set_black(other->rb_left);
 194                                 __rb_rotate_right(parent, root);
 195                                 node = root->rb_node;
 196                                 break;
 197                         }
 198                 }
 199         }
 200         if (node)
 201                 rb_set_black(node);
 202 }
 203
 204 void rb_erase(struct rb_node *node, struct rb_root *root)
 205 {
 206         struct rb_node *child, *parent;
 207         int color;
 208
 209         if (!node->rb_left)
 210                 child = node->rb_right;
 211         else if (!node->rb_right)
 212                 child = node->rb_left;
 213         else
 214         {
 215                 struct rb_node *old = node, *left;
 216
 217                 node = node->rb_right;
 218                 while ((left = node->rb_left) != NULL)
 219                         node = left;
 220                 child = node->rb_right;
 221                 parent = rb_parent(node);
 222                 color = rb_color(node);
 223
 224                 if (child)
 225                         rb_set_parent(child, parent);
 226                 if (parent == old) {
 227                         parent->rb_right = child;
 228                         parent = node;
 229                 } else
 230                         parent->rb_left = child;
 231
 232                 node->rb_parent_color = old->rb_parent_color;
 233                 node->rb_right = old->rb_right;
 234                 node->rb_left = old->rb_left;
 235
 236                 if (rb_parent(old))
 237                 {
 238                         if (rb_parent(old)->rb_left == old)
 239                                 rb_parent(old)->rb_left = node;
 240                         else
 241                                 rb_parent(old)->rb_right = node;
 242                 } else
 243                         root->rb_node = node;
 244
 245                 rb_set_parent(old->rb_left, node);
 246                 if (old->rb_right)
 247                         rb_set_parent(old->rb_right, node);
 248                 goto color;
 249         }
 250
 251         parent = rb_parent(node);
 252         color = rb_color(node);
 253
 254         if (child)
 255                 rb_set_parent(child, parent);
 256         if (parent)
 257         {
 258                 if (parent->rb_left == node)
 259                         parent->rb_left = child;
 260                 else
 261                         parent->rb_right = child;
 262         }
 263         else
 264                 root->rb_node = child;
 265
 266  color:
 267         if (color == RB_BLACK)
 268                 __rb_erase_color(child, parent, root);
 269 }
 270
 271 /*
 272  * This function returns the first node (in sort order) of the tree.
 273  */
 274 struct rb_node *rb_first(const struct rb_root *root)
 275 {
 276         struct rb_node  *n;
 277
 278         n = root->rb_node;
 279         if (!n)
 280                 return NULL;
 281         while (n->rb_left)
 282                 n = n->rb_left;
 283         return n;
 284 }
 285
 286 struct rb_node *rb_last(const struct rb_root *root)
 287 {
 288         struct rb_node  *n;
 289
 290         n = root->rb_node;
 291         if (!n)
 292                 return NULL;
 293         while (n->rb_right)
 294                 n = n->rb_right;
 295         return n;
 296 }
 297
 298 struct rb_node *rb_next(const struct rb_node *node)
 299 {
 300         struct rb_node *parent;
 301
 302         if (rb_parent(node) == node)
 303                 return NULL;
 304
 305         /* If we have a right-hand child, go down and then left as far
 306            as we can. */
 307         if (node->rb_right) {
 308                 node = node->rb_right;
 309                 while (node->rb_left)
 310                         node=node->rb_left;
 311                 return (struct rb_node *)node;
 312         }
 313
 314         /* No right-hand children.  Everything down and left is
 315            smaller than us, so any 'next' node must be in the general
 316            direction of our parent. Go up the tree; any time the
 317            ancestor is a right-hand child of its parent, keep going
 318            up. First time it's a left-hand child of its parent, said
 319            parent is our 'next' node. */
 320         while ((parent = rb_parent(node)) && node == parent->rb_right)
 321                 node = parent;
 322
 323         return parent;
 324 }
 325
 326 struct rb_node *rb_prev(const struct rb_node *node)
 327 {
 328         struct rb_node *parent;
 329
 330         if (rb_parent(node) == node)
 331                 return NULL;
 332
 333         /* If we have a left-hand child, go down and then right as far
 334            as we can. */
 335         if (node->rb_left) {
 336                 node = node->rb_left;
 337                 while (node->rb_right)
 338                         node=node->rb_right;
 339                 return (struct rb_node *)node;
 340         }
 341
 342         /* No left-hand children. Go up till we find an ancestor which
 343            is a right-hand child of its parent */
 344         while ((parent = rb_parent(node)) && node == parent->rb_left)
 345                 node = parent;
 346
 347         return parent;
 348 }
 349
 350 void rb_replace_node(struct rb_node *victim, struct rb_node *new,
 351                      struct rb_root *root)
 352 {
 353         struct rb_node *parent = rb_parent(victim);
 354
 355         /* Set the surrounding nodes to point to the replacement */
 356         if (parent) {
 357                 if (victim == parent->rb_left)
 358                         parent->rb_left = new;
 359                 else
 360                         parent->rb_right = new;
 361         } else {
 362                 root->rb_node = new;
 363         }
 364         if (victim->rb_left)
 365                 rb_set_parent(victim->rb_left, new);
 366         if (victim->rb_right)
 367                 rb_set_parent(victim->rb_right, new);
 368
 369         /* Copy the pointers/colour from the victim to the replacement */
 370         *new = *victim;
 371 }