tree.c

   1 /*
   2  * Copyright (C) 2024 Mikulas Patocka
   3  *
   4  * This file is part of Ajla.
   5  *
   6  * Ajla is free software: you can redistribute it and/or modify it under the
   7  * terms of the GNU General Public License as published by the Free Software
   8  * Foundation, either version 3 of the License, or (at your option) any later
   9  * version.
  10  *
  11  * Ajla is distributed in the hope that it will be useful, but WITHOUT ANY
  12  * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  13  * A PARTICULAR PURPOSE. See the GNU General Public License for more details.
  14  *
  15  * You should have received a copy of the GNU General Public License along with
  16  * Ajla. If not, see <https://www.gnu.org/licenses/>.
  17  */
  18
  19 #include "ajla.h"
  20
  21 #include "tree.h"
  22
  23 #define NEG(idx)        ((idx) ^ 1)
  24
  25 static struct tree_entry *rb_parent(struct tree_entry *n)
  26 {
  27 #ifdef DEBUG_RBTREE
  28         if (n->idx != RB_IDX_LEFT && n->idx != RB_IDX_RIGHT)
  29                 internal(file_line, "rb_parent: invalid index %d", n->idx);
  30 #endif
  31         return get_struct(n->parent - n->idx, struct tree_entry, children[0]);
  32 }
  33
  34 #ifdef DEBUG_RBTREE
  35 static void rb_verify_node_no_color(struct tree_entry *n)
  36 {
  37         if (unlikely(n->color != RB_BLACK && n->color != RB_RED))
  38                 internal(file_line, "rb_verify_node_no_color(%p): invalid color %d", n, n->color);
  39         if (unlikely(!n->parent))
  40                 internal(file_line, "rb_verify_node_no_color(%p): no parent link", n);
  41         if (unlikely(*n->parent != n))
  42                 internal(file_line, "rb_verify_node_no_color(%p): bad parent link: %p, %p", n, *n->parent, n);
  43         if (unlikely(n->idx != RB_IDX_ROOT)) {
  44                 struct tree_entry *p = rb_parent(n);
  45                 if (unlikely(p->idx != RB_IDX_LEFT && p->idx != RB_IDX_RIGHT && p->idx != RB_IDX_ROOT))
  46                         internal(file_line, "rb_verify_node_no_color(%p): invalid parent: %p, %d", n, p, p->idx);
  47         }
  48         if (n->children[0] && (n->children[0]->parent != &n->children[0] || n->children[0]->idx != RB_IDX_LEFT))
  49                 internal(file_line, "rb_verify_node_no_color(%p): bad left child: %p, %p, %d", n, n->children[0]->parent, &n->children[0], n->children[0]->idx);
  50         if (n->children[1] && (n->children[1]->parent != &n->children[1] || n->children[1]->idx != RB_IDX_RIGHT))
  51                 internal(file_line, "rb_verify_node_no_color(%p): bad right child: %p, %p, %d", n, n->children[1]->parent, &n->children[1], n->children[1]->idx);
  52 }
  53 void tree_verify_node(struct tree_entry *n)
  54 {
  55         rb_verify_node_no_color(n);
  56         if (n->idx == RB_IDX_ROOT && n->color == RB_RED)
  57                 internal(file_line, "tree_verify_node(%p): root is red", n);
  58         if (n->color == RB_RED) {
  59                 if ((n->children[0] && n->children[0]->color != RB_BLACK) ||
  60                     (n->children[1] && n->children[1]->color != RB_BLACK))
  61                         internal(file_line, "tree_verify_node(%p): red node has a red child", n);
  62                 if (n->idx != RB_IDX_ROOT) {
  63                         struct tree_entry *p = rb_parent(n);
  64                         if (p->color != RB_BLACK)
  65                                 internal(file_line, "tree_verify_node(%p): red parent %p has a red child", n, p);
  66                 }
  67         }
  68 }
  69 #else
  70 #define rb_verify_node_no_color(n)      ((void)0)
  71 #endif
  72
  73 static attr_noinline void rb_rotate(struct tree_entry *p, struct tree_entry *n)
  74 {
  75         uchar_efficient_t neg_n_idx;
  76         struct tree_entry *ch;
  77
  78         rb_verify_node_no_color(p);
  79         rb_verify_node_no_color(n);
  80
  81         *p->parent = n;
  82         n->parent = p->parent;
  83         neg_n_idx = NEG(n->idx);
  84         ch = p->children[n->idx] = n->children[neg_n_idx];
  85         if (ch) {
  86                 ch->parent = &p->children[n->idx];
  87                 ch->idx = n->idx;
  88         }
  89         p->parent = &n->children[neg_n_idx];
  90         n->children[neg_n_idx] = p;
  91         n->idx = p->idx;
  92         p->idx = neg_n_idx;
  93
  94         rb_verify_node_no_color(p);
  95         rb_verify_node_no_color(n);
  96 }
  97
  98 void attr_fastcall tree_insert_after_find_impl(struct tree_entry *n, uchar_efficient_t idx, struct tree_entry **link)
  99 {
 100         *link = n;
 101         n->color = RB_RED;
 102         n->idx = idx;
 103         n->parent = link;
 104         n->children[0] = n->children[1] = NULL;
 105
 106         rb_verify_node_no_color(n);
 107
 108         while (1) {
 109                 struct tree_entry *p, *gp, *un;
 110                 if (n->idx == RB_IDX_ROOT) {
 111                         n->color = RB_BLACK;
 112                         break;
 113                 }
 114                 p = rb_parent(n);
 115                 rb_verify_node_no_color(p);
 116                 if (p->color == RB_BLACK)
 117                         break;
 118                 gp = rb_parent(p);
 119                 rb_verify_node_no_color(gp);
 120                 un = gp->children[NEG(p->idx)];
 121                 if (un && un->color == RB_RED) {
 122                         rb_verify_node_no_color(un);
 123                         gp->color = RB_RED;
 124                         p->color = RB_BLACK;
 125                         un->color = RB_BLACK;
 126                         n = gp;
 127                 } else {
 128                         if (n->idx != p->idx) {
 129                                 struct tree_entry *tmp;
 130                                 rb_rotate(p, n);
 131                                 tmp = p;
 132                                 p = n;
 133                                 n = tmp;
 134                         }
 135                         rb_rotate(gp, p);
 136                         p->color = RB_BLACK;
 137                         gp->color = RB_RED;
 138                         break;
 139                 }
 140         }
 141 }
 142
 143 static void rb_fix_ptrs(struct tree_entry *n)
 144 {
 145         *n->parent = n;
 146         if (n->children[0]) n->children[0]->parent = &n->children[0];
 147         if (n->children[1]) n->children[1]->parent = &n->children[1];
 148 }
 149
 150 void attr_fastcall tree_delete(struct tree_entry *n)
 151 {
 152         uchar_efficient_t idx;
 153         struct tree_entry *ch, *p;
 154
 155         tree_verify_node(n);
 156
 157         if (!n->children[0]) {
 158                 idx = RB_IDX_RIGHT;
 159         } else if (!n->children[1]) {
 160                 idx = RB_IDX_LEFT;
 161         } else {
 162                 struct tree_entry *c, tmp;
 163                 c = n->children[1];
 164                 tree_verify_node(c);
 165                 while (c->children[0]) {
 166                         c = c->children[0];
 167                         tree_verify_node(c);
 168                 }
 169                 tmp = *c;
 170                 if (tmp.parent == &n->children[1]) tmp.parent = &c->children[1];
 171                 *c = *n;
 172                 if (c->children[1] == c) c->children[1] = n;
 173                 *n = tmp;
 174                 rb_fix_ptrs(c);
 175                 rb_fix_ptrs(n);
 176                 tree_verify_node(c);
 177                 tree_verify_node(n);
 178                 idx = RB_IDX_RIGHT;
 179         }
 180         ch = n->children[idx];
 181         *n->parent = ch;
 182         if (ch) {
 183                 ch->parent = n->parent;
 184                 ch->idx = n->idx;
 185         }
 186         if (n->color == RB_RED)
 187                 return;
 188 again:
 189         idx = n->idx;
 190         if (idx == RB_IDX_ROOT)
 191                 goto set_ch_black;
 192         p = rb_parent(n);
 193         rb_verify_node_no_color(p);
 194         if (!ch || ch->color == RB_BLACK) {
 195                 struct tree_entry *s, *z;
 196 again2:
 197                 s = p->children[NEG(idx)];
 198                 rb_verify_node_no_color(s);
 199                 if (s->color == RB_BLACK) {
 200                         z = s->children[NEG(idx)];
 201                         if (z && z->color == RB_RED) {
 202                                 rb_verify_node_no_color(z);
 203                                 z->color = RB_BLACK;
 204 rotate_p_s_return:
 205                                 s->color = p->color;
 206                                 p->color = RB_BLACK;
 207                                 rb_rotate(p, s);
 208                                 return;
 209                         }
 210                         z = s->children[idx];
 211                         if (z && z->color == RB_RED) {
 212                                 rb_verify_node_no_color(z);
 213                                 z->color = RB_BLACK;
 214                                 rb_rotate(s, z);
 215                                 s = z;
 216                                 goto rotate_p_s_return;
 217                         }
 218                         s->color = RB_RED;
 219
 220                         ch = n = p;
 221                         goto again;
 222                 } else {
 223                         rb_rotate(p, s);
 224                         p->color = RB_RED;
 225                         s->color = RB_BLACK;
 226                         goto again2;
 227                 }
 228         }
 229 set_ch_black:
 230         if (ch)
 231                 ch->color = RB_BLACK;
 232 }
 233
 234 struct tree_entry * attr_fastcall tree_first(struct tree *root)
 235 {
 236         struct tree_entry *n = root->root;
 237         if (unlikely(!n))
 238                 return NULL;
 239         tree_verify_node(n);
 240         while (n->children[0]) {
 241                 n = n->children[0];
 242                 tree_verify_node(n);
 243         }
 244         return n;
 245 }
 246
 247 struct tree_entry * attr_fastcall tree_next(struct tree_entry *e)
 248 {
 249         tree_verify_node(e);
 250         if (e->children[1]) {
 251                 e = e->children[1];
 252                 tree_verify_node(e);
 253                 while (e->children[0]) {
 254                         e = e->children[0];
 255                         tree_verify_node(e);
 256                 }
 257                 return e;
 258         }
 259         while (e->idx == RB_IDX_RIGHT) {
 260                 e = rb_parent(e);
 261                 tree_verify_node(e);
 262         }
 263         if (e->idx == RB_IDX_LEFT) {
 264                 e = rb_parent(e);
 265                 tree_verify_node(e);
 266                 return e;
 267         }
 268         return NULL;
 269 }
 270
 271 struct tree_entry * attr_fastcall tree_last(struct tree *root)
 272 {
 273         struct tree_entry *n = root->root;
 274         if (unlikely(!n))
 275                 return NULL;
 276         tree_verify_node(n);
 277         while (n->children[1]) {
 278                 n = n->children[1];
 279                 tree_verify_node(n);
 280         }
 281         return n;
 282 }