newlib/treemap.ajla

   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 unit treemap;
  20
  21 private type xtreemap(key : type, cls : class_ord(key), value : type);
  22 type treemap(key : type, value : type, cls : class_ord(key)) := xtreemap(key, cls, value);
  23
  24 record treemap_key_value(key : type, value : type) [
  25         k : key;
  26         v : value;
  27 ]
  28
  29 fn treemap_init(key : type, value : type, const cls : class_ord(key)) : treemap(key, value, cls);
  30 fn treemap_is_nonempty(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : bool;
  31 fn treemap_test(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls), k : key) : bool;
  32 fn treemap_search(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls), k : key) : maybe(value);
  33 fn treemap_search_default(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls), k : key, def : value) : value;
  34 fn treemap_first(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : maybe(treemap_key_value(key, value));
  35 fn treemap_last(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : maybe(treemap_key_value(key, value));
  36 fn treemap_next(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key) : maybe(treemap_key_value(key, value));
  37 fn treemap_prev(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key) : maybe(treemap_key_value(key, value));
  38 fn treemap_size(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : int;
  39 fn treemap_insert(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key, v : value) : treemap(key, value, cls);
  40 fn treemap_delete(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key) : treemap(key, value, cls);
  41
  42 implicit fn instance_functor_treemap~inline(key : type, const cls : class_ord(key)) : class_functor(xtreemap(key, cls,));
  43
  44 conversion fn treemap_iterator~type(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : class_iterator :=
  45         class_iterator.[
  46                 state : maybe(treemap_key_value(key, value)),
  47                 element : treemap_key_value(key, value),
  48                 init : treemap_first(tm),
  49                 test : lambda(st : maybe(treemap_key_value(key, value))) [ return st is j; ],
  50                 get_element : lambda(st : maybe(treemap_key_value(key, value))) [ return st.j; ],
  51                 increment : lambda(st : maybe(treemap_key_value(key, value))) [ return treemap_next(tm, st.j.k); ],
  52         ];
  53
  54 fn treemap_iterator_reverse~type(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : class_iterator :=
  55         class_iterator.[
  56                 state : maybe(treemap_key_value(key, value)),
  57                 element : treemap_key_value(key, value),
  58                 init : treemap_last(tm),
  59                 test : lambda(st : maybe(treemap_key_value(key, value))) [ return st is j; ],
  60                 get_element : lambda(st : maybe(treemap_key_value(key, value))) [ return st.j; ],
  61                 increment : lambda(st : maybe(treemap_key_value(key, value))) [ return treemap_prev(tm, st.j.k); ],
  62         ];
  63
  64 type treeset(key : type, cls : class_ord(key));
  65
  66 fn treeset_init(key : type, const cls : class_ord(key)) : treeset(key, cls);
  67 fn treeset_is_nonempty(key : type, const cls : class_ord(key), ts : treeset(key, cls)) : bool;
  68 fn treeset_test(key : type, const cls : class_ord(key), ts : treeset(key, cls), k : key) : bool;
  69 fn treeset_first(key : type, const cls : class_ord(key), ts : treeset(key, cls)) : maybe(key);
  70 fn treeset_last(key : type, const cls : class_ord(key), ts : treeset(key, cls)) : maybe(key);
  71 fn treeset_next(key : type, const implicit cls : class_ord(key), ts : treeset(key, cls), k : key) : maybe(key);
  72 fn treeset_prev(key : type, const implicit cls : class_ord(key), ts : treeset(key, cls), k : key) : maybe(key);
  73 fn treeset_size(key : type, const cls : class_ord(key), ts : treeset(key, cls)) : int;
  74 fn treeset_set(key : type, const cls : class_ord(key), ts : treeset(key, cls), k : key) : treeset(key, cls);
  75 fn treeset_clear(key : type, const cls : class_ord(key), ts : treeset(key, cls), k : key) : treeset(key, cls);
  76 fn treeset_from_list(key : type, const cls : class_ord(key), l : list(key)) : treeset(key, cls);
  77
  78 conversion fn treeset_iterator~type(key : type, const cls : class_ord(key), ts : treeset(key, cls)) : class_iterator :=
  79         class_iterator.[
  80                 state : maybe(key),
  81                 element : key,
  82                 init : treeset_first(ts),
  83                 test : lambda(st : maybe(key)) [ return st is j; ],
  84                 get_element : lambda(st : maybe(key)) [ return st.j; ],
  85                 increment : lambda(st : maybe(key)) [ return treeset_next(ts, st.j); ],
  86         ];
  87
  88 fn treeset_iterator_reverse~type(key : type, const cls : class_ord(key), ts : treeset(key, cls)) : class_iterator :=
  89         class_iterator.[
  90                 state : maybe(key),
  91                 element : key,
  92                 init : treeset_last(ts),
  93                 test : lambda(st : maybe(key)) [ return st is j; ],
  94                 get_element : lambda(st : maybe(key)) [ return st.j; ],
  95                 increment : lambda(st : maybe(key)) [ return treeset_prev(ts, st.j); ],
  96         ];
  97
  98 implementation
  99
 100 record treemap_entry(key : type, value : type, cls : class_ord(key)) [
 101         k : key;
 102         v : value;
 103         left : treemap(key, value, cls);
 104         right : treemap(key, value, cls);
 105         balance : sint8;
 106 ]
 107
 108 type xtreemap(key : type, cls : class_ord(key), value : type) := maybe(treemap_entry(key, value, cls));
 109
 110
 111 {
 112 fn treemap_verify(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls)) : int
 113 [
 114         if tm is n then
 115                 return 0;
 116         var depth1 := treemap_verify(tm.j.left);
 117         var depth2 := treemap_verify(tm.j.right);
 118         var diff := tm.j.balance;
 119         if abs(diff) > 1 or depth2 - depth1 <> diff then
 120                 abort internal("treemap_verify: diff " + ntos(diff) + ", depth1 " + ntos(depth1) + ", depth2 " + ntos(depth2));
 121         return max(depth1, depth2) + 1;
 122 ]
 123 }
 124
 125 fn treemap_init(key : type, value : type, const cls : class_ord(key)) : treemap(key, value, cls)
 126 [
 127         return treemap(key, value, cls).n;
 128 ]
 129
 130 fn treemap_is_nonempty(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : bool
 131 [
 132         return tm is j;
 133 ]
 134
 135 fn treemap_test(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key) : bool
 136 [
 137 again:
 138         if tm is n then
 139                 return false;
 140         if k = tm.j.k then
 141                 return true;
 142         if k < tm.j.k then
 143                 tm := tm.j.left;
 144         else
 145                 tm := tm.j.right;
 146         goto again;
 147 ]
 148
 149 fn treemap_search(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key) : maybe(value)
 150 [
 151 again:
 152         if tm is n then
 153                 return maybe(value).n;
 154         if k = tm.j.k then
 155                 return maybe(value).j.(tm.j.v);
 156         if k < tm.j.k then
 157                 tm := tm.j.left;
 158         else
 159                 tm := tm.j.right;
 160         goto again;
 161 ]
 162
 163 fn treemap_search_default(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls), k : key, def : value) : value
 164 [
 165         var s := treemap_search(tm, k);
 166         if s is j then
 167                 return s.j;
 168         return def;
 169 ]
 170
 171 fn treemap_first(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : maybe(treemap_key_value(key, value))
 172 [
 173         if tm is n then
 174                 return maybe(treemap_key_value(key, value)).n;
 175         while tm.j.left is j do
 176                 tm := tm.j.left;
 177         return maybe(treemap_key_value(key, value)).j.(treemap_key_value(key, value).[ k : tm.j.k, v : tm.j.v ]);
 178 ]
 179
 180 fn treemap_last(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : maybe(treemap_key_value(key, value))
 181 [
 182         if tm is n then
 183                 return maybe(treemap_key_value(key, value)).n;
 184         while tm.j.right is j do
 185                 tm := tm.j.right;
 186         return maybe(treemap_key_value(key, value)).j.(treemap_key_value(key, value).[ k : tm.j.k, v : tm.j.v ]);
 187 ]
 188
 189 fn treemap_next(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key) : maybe(treemap_key_value(key, value))
 190 [
 191         var result := maybe(treemap_key_value(key, value)).n;
 192 again:
 193         if tm is n then
 194                 return result;
 195         if k < tm.j.k then [
 196                 result := maybe(treemap_key_value(key, value)).j.(treemap_key_value(key, value).[ k : tm.j.k, v : tm.j.v ]);
 197                 tm := tm.j.left;
 198         ] else [
 199                 tm := tm.j.right;
 200         ]
 201         goto again;
 202 ]
 203
 204 fn treemap_prev(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key) : maybe(treemap_key_value(key, value))
 205 [
 206         var result := maybe(treemap_key_value(key, value)).n;
 207 again:
 208         if tm is n then
 209                 return result;
 210         if k > tm.j.k then [
 211                 result := maybe(treemap_key_value(key, value)).j.(treemap_key_value(key, value).[ k : tm.j.k, v : tm.j.v ]);
 212                 tm := tm.j.right;
 213         ] else [
 214                 tm := tm.j.left;
 215         ]
 216         goto again;
 217 ]
 218
 219 fn treemap_size(key : type, value : type, const cls : class_ord(key), tm : treemap(key, value, cls)) : int
 220 [
 221         if tm is n then
 222                 return 0;
 223         return 1 + treemap_size(tm.j.left) + treemap_size(tm.j.right);
 224 ]
 225
 226 fn treemap_insert_internal(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key, v : value) : (treemap(key, value, cls), bool)
 227 [
 228         const te := treemap_entry(key, value);
 229         if tm is n then
 230                 return maybe(te).j.(te.[ k : k, v : v, left : maybe(te).n, right : maybe(te).n, balance : 0 ]), true;
 231         if k = tm.j.k then [
 232                 tm.j.v := v;
 233                 return tm, false;
 234         ]
 235         var height_changed : bool;
 236         if k < tm.j.k then [
 237                 var s := maybe(te).n;
 238                 s, tm.j.left := tm.j.left, s;
 239                 s, height_changed := treemap_insert_internal(s, k, v);
 240                 if height_changed then [
 241                         if tm.j.balance = 1 then [
 242                                 tm.j.balance := 0;
 243                                 height_changed := false;
 244                         ] else if tm.j.balance = 0 then [
 245                                 tm.j.balance := -1;
 246                         ] else [
 247                                 if s.j.balance = 0 then
 248                                         abort internal("treemap_insert_internal: s.j.balance = 0");
 249                                 else if s.j.balance = -1 then [
 250                                         //eval debug("insert rotate 1");
 251                                         tm.j.left := s.j.right;
 252                                         tm.j.balance := 0;
 253                                         s.j.right := tm;
 254                                         s.j.balance := 0;
 255                                         return s, false;
 256                                 ] else [
 257                                         //eval debug("insert rotate 2");
 258                                         var x := s.j.right;
 259                                         var balance := x.j.balance;
 260                                         tm.j.left := x.j.right;
 261                                         s.j.right := x.j.left;
 262                                         tm.j.balance := select(balance = -1, 0, 1);
 263                                         s.j.balance := select(balance = 1, 0, -1);
 264                                         x.j.left := s;
 265                                         x.j.right := tm;
 266                                         x.j.balance := 0;
 267                                         return x, false;
 268                                 ]
 269                         ]
 270                 ]
 271                 s, tm.j.left := tm.j.left, s;
 272         ] else [
 273                 var s := maybe(te).n;
 274                 s, tm.j.right := tm.j.right, s;
 275                 s, height_changed := treemap_insert_internal(s, k, v);
 276                 if height_changed then [
 277                         if tm.j.balance = -1 then [
 278                                 tm.j.balance := 0;
 279                                 height_changed := false;
 280                         ] else if tm.j.balance = 0 then [
 281                                 tm.j.balance := 1;
 282                         ] else [
 283                                 if s.j.balance = 0 then
 284                                         abort internal("treemap_insert_internal: s.j.balance = 0");
 285                                 else if s.j.balance = 1 then [
 286                                         //eval debug("insert rotate 3");
 287                                         tm.j.right := s.j.left;
 288                                         tm.j.balance := 0;
 289                                         s.j.left := tm;
 290                                         s.j.balance := 0;
 291                                         return s, false;
 292                                 ] else [
 293                                         //eval debug("insert rotate 4");
 294                                         var x := s.j.left;
 295                                         var balance := x.j.balance;
 296                                         tm.j.right := x.j.left;
 297                                         s.j.left := x.j.right;
 298                                         tm.j.balance := select(balance = 1, 0, -1);
 299                                         s.j.balance := select(balance = -1, 0, 1);
 300                                         x.j.right := s;
 301                                         x.j.left := tm;
 302                                         x.j.balance := 0;
 303                                         return x, false;
 304                                 ]
 305                         ]
 306                 ]
 307                 s, tm.j.right := tm.j.right, s;
 308         ]
 309         return tm, height_changed;
 310 ]
 311
 312 fn treemap_insert(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key, v : value) : treemap(key, value, cls)
 313 [
 314         var tm, height_changed := treemap_insert_internal(tm, k, v);
 315         //xeval treemap_verify(tm);
 316         return tm;
 317 ]
 318
 319 fn treemap_delete_internal(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key) : (treemap(key, value, cls), bool)
 320 [
 321         const te := treemap_entry(key, value);
 322         if tm is n then
 323                 return tm, false;
 324         var height_changed : bool;
 325         if k = tm.j.k then [
 326                 if tm.j.left is n then
 327                         return tm.j.right, true;
 328                 if tm.j.right is n then
 329                         return tm.j.left, true;
 330                 var last := tm.j.left;
 331                 while last.j.right is j do
 332                         last := last.j.right;
 333                 tm.j.k := last.j.k;
 334                 tm.j.v := last.j.v;
 335                 k := last.j.k;
 336                 goto left;
 337         ]
 338         if k < tm.j.k then [
 339 left:
 340                 var s := maybe(te).n;
 341                 s, tm.j.left := tm.j.left, s;
 342                 s, height_changed := treemap_delete_internal(s, k);
 343                 s, tm.j.left := tm.j.left, s;
 344                 if height_changed then [
 345                         if tm.j.balance = -1 then [
 346                                 tm.j.balance := 0;
 347                         ] else if tm.j.balance = 0 then [
 348                                 tm.j.balance := 1;
 349                                 height_changed := false;
 350                         ] else [
 351                                 s := tm.j.right;
 352                                 var balance := s.j.balance;
 353                                 if balance >= 0 then [
 354                                         //eval debug("delete rotate 1");
 355                                         tm.j.right := s.j.left;
 356                                         tm.j.balance := 1 - balance;
 357                                         s.j.left := tm;
 358                                         s.j.balance := balance - 1;
 359                                         return s, balance <> 0;
 360                                 ] else [
 361                                         //eval debug("delete rotate 2");
 362                                         var x := s.j.left;
 363                                         balance := x.j.balance;
 364                                         tm.j.right := x.j.left;
 365                                         s.j.left := x.j.right;
 366                                         tm.j.balance := select(balance = 1, 0, -1);
 367                                         s.j.balance := select(balance = -1, 0, 1);
 368                                         x.j.right := s;
 369                                         x.j.left := tm;
 370                                         x.j.balance := 0;
 371                                         return x, true;
 372                                 ]
 373                         ]
 374                 ]
 375         ] else [
 376                 var s := maybe(te).n;
 377                 s, tm.j.right := tm.j.right, s;
 378                 s, height_changed := treemap_delete_internal(s, k);
 379                 s, tm.j.right := tm.j.right, s;
 380                 if height_changed then [
 381                         if tm.j.balance = 1 then [
 382                                 tm.j.balance := 0;
 383                         ] else if tm.j.balance = 0 then [
 384                                 tm.j.balance := -1;
 385                                 height_changed := false;
 386                         ] else [
 387                                 s := tm.j.left;
 388                                 var balance := s.j.balance;
 389                                 if balance <= 0 then [
 390                                         //eval debug("delete rotate 3");
 391                                         tm.j.left := s.j.right;
 392                                         tm.j.balance := -1 - balance;
 393                                         s.j.right := tm;
 394                                         s.j.balance := balance + 1;
 395                                         return s, balance <> 0;
 396                                 ] else [
 397                                         //eval debug("delete rotate 4");
 398                                         var x := s.j.right;
 399                                         balance := x.j.balance;
 400                                         tm.j.left := x.j.right;
 401                                         s.j.right := x.j.left;
 402                                         tm.j.balance := select(balance = -1, 0, 1);
 403                                         s.j.balance := select(balance = 1, 0, -1);
 404                                         x.j.left := s;
 405                                         x.j.right := tm;
 406                                         x.j.balance := 0;
 407                                         return x, true;
 408                                 ]
 409                         ]
 410                 ]
 411         ]
 412         return tm, height_changed;
 413 ]
 414
 415 fn treemap_delete(key : type, value : type, const implicit cls : class_ord(key), tm : treemap(key, value, cls), k : key) : treemap(key, value, cls)
 416 [
 417         var tm, height_changed := treemap_delete_internal(tm, k);
 418         //xeval treemap_verify(tm);
 419         return tm;
 420 ]
 421
 422 fn treemap_map(key : type, const implicit cls : class_ord(key), value : type, new_value : type, tm : treemap(key, value, cls), m : fn(value) : new_value) : treemap(key, new_value, cls)
 423 [
 424         if tm is n then
 425                 return treemap_init(key, new_value, cls);
 426         const te := treemap_entry(key, new_value);
 427         return maybe(te).j.(te.[
 428                 k : tm.j.k,
 429                 v : m(tm.j.v),
 430                 balance : tm.j.balance,
 431                 left : treemap_map(key, cls, value, new_value, tm.j.left, m),
 432                 right : treemap_map(key, cls, value, new_value, tm.j.right, m),
 433         ]);
 434 ]
 435
 436 implicit fn instance_functor_treemap~inline(key : type, const cls : class_ord(key)) : class_functor(xtreemap(key, cls,)) :=
 437         class_functor(xtreemap(key, cls,)).[
 438                 map : treemap_map(key, cls,,,,),
 439         ];
 440
 441
 442 type treeset(key : type, cls : class_ord(key)) := treemap(key, unit_type, cls);
 443
 444 fn treeset_init(key : type, const cls : class_ord(key)) : treeset(key, cls) := treemap_init(key, unit_type, cls);
 445 fn treeset_is_nonempty(key : type, const cls : class_ord(key), ts : treeset(key, cls)) : bool := treemap_is_nonempty(ts);
 446 fn treeset_test(key : type, const cls : class_ord(key), ts : treeset(key, cls), k : key) : bool := treemap_test(ts, k);
 447 fn treeset_first(key : type, const cls : class_ord(key), ts : treeset(key, cls)) : maybe(key)
 448 [
 449         var f := treemap_first(ts);
 450         if f is n then
 451                 return maybe(key).n;
 452         else
 453                 return maybe(key).j.(f.j.k);
 454 ]
 455 fn treeset_last(key : type, const cls : class_ord(key), ts : treeset(key, cls)) : maybe(key)
 456 [
 457         var f := treemap_last(ts);
 458         if f is n then
 459                 return maybe(key).n;
 460         else
 461                 return maybe(key).j.(f.j.k);
 462 ]
 463 fn treeset_next(key : type, const implicit cls : class_ord(key), ts : treeset(key, cls), k : key) : maybe(key)
 464 [
 465         var f := treemap_next(ts, k);
 466         if f is n then
 467                 return maybe(key).n;
 468         else
 469                 return maybe(key).j.(f.j.k);
 470 ]
 471 fn treeset_prev(key : type, const implicit cls : class_ord(key), ts : treeset(key, cls), k : key) : maybe(key)
 472 [
 473         var f := treemap_prev(ts, k);
 474         if f is n then
 475                 return maybe(key).n;
 476         else
 477                 return maybe(key).j.(f.j.k);
 478 ]
 479 fn treeset_size(key : type, const cls : class_ord(key), tm : treeset(key, cls)) : int := treemap_size(tm);
 480 fn treeset_set(key : type, const cls : class_ord(key), ts : treeset(key, cls), k : key) : treeset(key, cls) := treemap_insert(ts, k, unit_value);
 481 fn treeset_clear(key : type, const cls : class_ord(key), ts : treeset(key, cls), k : key) : treeset(key, cls) := treemap_delete(ts, k);
 482 fn treeset_from_list(key : type, const cls : class_ord(key), l : list(key)) : treeset(key, cls)
 483 [
 484         var ts := treeset_init(key, cls);
 485         for f in l do
 486                 ts := treeset_set(ts, f);
 487         return ts;
 488 ]