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