weakmap.c

   1 #include "internal.h"
   2 #include "internal/gc.h"
   3 #include "internal/hash.h"
   4 #include "internal/proc.h"
   5 #include "internal/sanitizers.h"
   6 #include "ruby/st.h"
   7
   8 /* ===== WeakMap =====
   9  *
  10  * WeakMap contains one ST table which contains a pointer to the object as the
  11  * key and a pointer to the object as the value. This means that the key and
  12  * value of the table are both of the type `VALUE *`.
  13  *
  14  * The objects are not directly stored as keys and values in the table because
  15  * `rb_gc_mark_weak` requires a pointer to the memory location to overwrite
  16  * when the object is reclaimed. Using a pointer into the ST table entry is not
  17  * safe because the pointer can change when the ST table is resized.
  18  *
  19  * WeakMap hashes and compares using the pointer address of the object.
  20  *
  21  * For performance and memory efficiency reasons, the key and value
  22  * are allocated at the same time and adjacent to each other.
  23  *
  24  * During GC and while iterating, reclaimed entries (i.e. either the key or
  25  * value points to `Qundef`) are removed from the ST table.
  26  */
  27
  28 struct weakmap {
  29     st_table *table;
  30 };
  31
  32 struct weakmap_entry {
  33     VALUE key;
  34     VALUE val;
  35 };
  36
  37 static bool
  38 wmap_live_p(VALUE obj)
  39 {
  40     return !UNDEF_P(obj);
  41 }
  42
  43 struct wmap_foreach_data {
  44     int (*func)(struct weakmap_entry *, st_data_t);
  45     st_data_t arg;
  46
  47     struct weakmap_entry *dead_entry;
  48 };
  49
  50 static int
  51 wmap_foreach_i(st_data_t key, st_data_t val, st_data_t arg)
  52 {
  53     struct wmap_foreach_data *data = (struct wmap_foreach_data *)arg;
  54
  55     if (data->dead_entry != NULL) {
  56         ruby_sized_xfree(data->dead_entry, sizeof(struct weakmap_entry));
  57         data->dead_entry = NULL;
  58     }
  59
  60     struct weakmap_entry *entry = (struct weakmap_entry *)key;
  61     RUBY_ASSERT(&entry->val == (VALUE *)val);
  62
  63     if (wmap_live_p(entry->key) && wmap_live_p(entry->val)) {
  64         VALUE k = entry->key;
  65         VALUE v = entry->val;
  66
  67         int ret = data->func(entry, data->arg);
  68
  69         RB_GC_GUARD(k);
  70         RB_GC_GUARD(v);
  71
  72         return ret;
  73     }
  74     else {
  75         /* We cannot free the weakmap_entry here because the ST_DELETE could
  76          * hash the key which would read the weakmap_entry and would cause a
  77          * use-after-free. Instead, we store this entry and free it on the next
  78          * iteration. */
  79         data->dead_entry = entry;
  80
  81         return ST_DELETE;
  82     }
  83 }
  84
  85 static void
  86 wmap_foreach(struct weakmap *w, int (*func)(struct weakmap_entry *, st_data_t), st_data_t arg)
  87 {
  88     struct wmap_foreach_data foreach_data = {
  89         .func = func,
  90         .arg = arg,
  91         .dead_entry = NULL,
  92     };
  93
  94     st_foreach(w->table, wmap_foreach_i, (st_data_t)&foreach_data);
  95
  96     ruby_sized_xfree(foreach_data.dead_entry, sizeof(struct weakmap_entry));
  97 }
  98
  99 static int
 100 wmap_mark_weak_table_i(struct weakmap_entry *entry, st_data_t _)
 101 {
 102     rb_gc_mark_weak(&entry->key);
 103     rb_gc_mark_weak(&entry->val);
 104
 105     return ST_CONTINUE;
 106 }
 107
 108 static void
 109 wmap_mark(void *ptr)
 110 {
 111     struct weakmap *w = ptr;
 112     if (w->table) {
 113         wmap_foreach(w, wmap_mark_weak_table_i, (st_data_t)0);
 114     }
 115 }
 116
 117 static int
 118 wmap_free_table_i(st_data_t key, st_data_t val, st_data_t arg)
 119 {
 120     struct weakmap_entry *entry = (struct weakmap_entry *)key;
 121     RUBY_ASSERT(&entry->val == (VALUE *)val);
 122     ruby_sized_xfree(entry, sizeof(struct weakmap_entry));
 123
 124     return ST_CONTINUE;
 125 }
 126
 127 static void
 128 wmap_free(void *ptr)
 129 {
 130     struct weakmap *w = ptr;
 131
 132     st_foreach(w->table, wmap_free_table_i, 0);
 133     st_free_table(w->table);
 134 }
 135
 136 static size_t
 137 wmap_memsize(const void *ptr)
 138 {
 139     const struct weakmap *w = ptr;
 140
 141     size_t size = 0;
 142     size += st_memsize(w->table);
 143     /* The key and value of the table each take sizeof(VALUE) in size. */
 144     size += st_table_size(w->table) * (2 * sizeof(VALUE));
 145
 146     return size;
 147 }
 148
 149 struct wmap_compact_table_data {
 150     st_table *table;
 151     struct weakmap_entry *dead_entry;
 152 };
 153
 154 static int
 155 wmap_compact_table_i(st_data_t key, st_data_t val, st_data_t d)
 156 {
 157     struct wmap_compact_table_data *data = (struct wmap_compact_table_data *)d;
 158     if (data->dead_entry != NULL) {
 159         ruby_sized_xfree(data->dead_entry, sizeof(struct weakmap_entry));
 160         data->dead_entry = NULL;
 161     }
 162
 163     struct weakmap_entry *entry = (struct weakmap_entry *)key;
 164
 165     entry->val = rb_gc_location(entry->val);
 166
 167     VALUE new_key = rb_gc_location(entry->key);
 168
 169     /* If the key object moves, then we must reinsert because the hash is
 170      * based on the pointer rather than the object itself. */
 171     if (entry->key != new_key) {
 172         DURING_GC_COULD_MALLOC_REGION_START();
 173         {
 174             struct weakmap_entry *new_entry = xmalloc(sizeof(struct weakmap_entry));
 175             new_entry->key = new_key;
 176             new_entry->val = entry->val;
 177             st_insert(data->table, (st_data_t)&new_entry->key, (st_data_t)&new_entry->val);
 178         }
 179         DURING_GC_COULD_MALLOC_REGION_END();
 180
 181         /* We cannot free the weakmap_entry here because the ST_DELETE could
 182          * hash the key which would read the weakmap_entry and would cause a
 183          * use-after-free. Instead, we store this entry and free it on the next
 184          * iteration. */
 185         data->dead_entry = entry;
 186
 187         return ST_DELETE;
 188     }
 189
 190     return ST_CONTINUE;
 191 }
 192
 193 static void
 194 wmap_compact(void *ptr)
 195 {
 196     struct weakmap *w = ptr;
 197
 198     if (w->table) {
 199         struct wmap_compact_table_data compact_data = {
 200             .table = w->table,
 201             .dead_entry = NULL,
 202         };
 203
 204         st_foreach(w->table, wmap_compact_table_i, (st_data_t)&compact_data);
 205
 206         ruby_sized_xfree(compact_data.dead_entry, sizeof(struct weakmap_entry));
 207     }
 208 }
 209
 210 static const rb_data_type_t weakmap_type = {
 211     "weakmap",
 212     {
 213         wmap_mark,
 214         wmap_free,
 215         wmap_memsize,
 216         wmap_compact,
 217     },
 218     0, 0, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE
 219 };
 220
 221 static int
 222 wmap_cmp(st_data_t x, st_data_t y)
 223 {
 224     VALUE x_obj = *(VALUE *)x;
 225     VALUE y_obj = *(VALUE *)y;
 226
 227     if (!wmap_live_p(x_obj) && !wmap_live_p(y_obj)) {
 228         return x != y;
 229     }
 230     else {
 231         return x_obj != y_obj;
 232     }
 233 }
 234
 235 static st_index_t
 236 wmap_hash(st_data_t n)
 237 {
 238     return st_numhash(*(VALUE *)n);
 239 }
 240
 241 static const struct st_hash_type wmap_hash_type = {
 242     wmap_cmp,
 243     wmap_hash,
 244 };
 245
 246 static VALUE
 247 wmap_allocate(VALUE klass)
 248 {
 249     struct weakmap *w;
 250     VALUE obj = TypedData_Make_Struct(klass, struct weakmap, &weakmap_type, w);
 251     w->table = st_init_table(&wmap_hash_type);
 252     return obj;
 253 }
 254
 255 static VALUE
 256 wmap_inspect_append(VALUE str, VALUE obj)
 257 {
 258     if (SPECIAL_CONST_P(obj)) {
 259         return rb_str_append(str, rb_inspect(obj));
 260     }
 261     else {
 262         return rb_str_append(str, rb_any_to_s(obj));
 263     }
 264 }
 265
 266 static int
 267 wmap_inspect_i(struct weakmap_entry *entry, st_data_t data)
 268 {
 269     VALUE str = (VALUE)data;
 270
 271     if (RSTRING_PTR(str)[0] == '#') {
 272         rb_str_cat2(str, ", ");
 273     }
 274     else {
 275         rb_str_cat2(str, ": ");
 276         RSTRING_PTR(str)[0] = '#';
 277     }
 278
 279     wmap_inspect_append(str, entry->key);
 280     rb_str_cat2(str, " => ");
 281     wmap_inspect_append(str, entry->val);
 282
 283     return ST_CONTINUE;
 284 }
 285
 286 static VALUE
 287 wmap_inspect(VALUE self)
 288 {
 289     VALUE c = rb_class_name(CLASS_OF(self));
 290     struct weakmap *w;
 291     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 292
 293     VALUE str = rb_sprintf("-<%"PRIsVALUE":%p", c, (void *)self);
 294
 295     wmap_foreach(w, wmap_inspect_i, (st_data_t)str);
 296
 297     RSTRING_PTR(str)[0] = '#';
 298     rb_str_cat2(str, ">");
 299
 300     return str;
 301 }
 302
 303 static int
 304 wmap_each_i(struct weakmap_entry *entry, st_data_t _)
 305 {
 306     rb_yield_values(2, entry->key, entry->val);
 307
 308     return ST_CONTINUE;
 309 }
 310
 311 /*
 312  * call-seq:
 313  *   map.each {|key, val| ... } -> self
 314  *
 315  * Iterates over keys and values. Note that unlike other collections,
 316  * +each+ without block isn't supported.
 317  *
 318  */
 319 static VALUE
 320 wmap_each(VALUE self)
 321 {
 322     struct weakmap *w;
 323     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 324
 325     wmap_foreach(w, wmap_each_i, (st_data_t)0);
 326
 327     return self;
 328 }
 329
 330 static int
 331 wmap_each_key_i(struct weakmap_entry *entry, st_data_t _data)
 332 {
 333     rb_yield(entry->key);
 334
 335     return ST_CONTINUE;
 336 }
 337
 338 /*
 339  * call-seq:
 340  *   map.each_key {|key| ... } -> self
 341  *
 342  * Iterates over keys. Note that unlike other collections,
 343  * +each_key+ without block isn't supported.
 344  *
 345  */
 346 static VALUE
 347 wmap_each_key(VALUE self)
 348 {
 349     struct weakmap *w;
 350     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 351
 352     wmap_foreach(w, wmap_each_key_i, (st_data_t)0);
 353
 354     return self;
 355 }
 356
 357 static int
 358 wmap_each_value_i(struct weakmap_entry *entry, st_data_t _data)
 359 {
 360     rb_yield(entry->val);
 361
 362     return ST_CONTINUE;
 363 }
 364
 365 /*
 366  * call-seq:
 367  *   map.each_value {|val| ... } -> self
 368  *
 369  * Iterates over values. Note that unlike other collections,
 370  * +each_value+ without block isn't supported.
 371  *
 372  */
 373 static VALUE
 374 wmap_each_value(VALUE self)
 375 {
 376     struct weakmap *w;
 377     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 378
 379     wmap_foreach(w, wmap_each_value_i, (st_data_t)0);
 380
 381     return self;
 382 }
 383
 384 static int
 385 wmap_keys_i(struct weakmap_entry *entry, st_data_t arg)
 386 {
 387     VALUE ary = (VALUE)arg;
 388
 389     rb_ary_push(ary, entry->key);
 390
 391     return ST_CONTINUE;
 392 }
 393
 394 /*
 395  * call-seq:
 396  *   map.keys -> new_array
 397  *
 398  * Returns a new Array containing all keys in the map.
 399  *
 400  */
 401 static VALUE
 402 wmap_keys(VALUE self)
 403 {
 404     struct weakmap *w;
 405     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 406
 407     VALUE ary = rb_ary_new();
 408     wmap_foreach(w, wmap_keys_i, (st_data_t)ary);
 409
 410     return ary;
 411 }
 412
 413 static int
 414 wmap_values_i(struct weakmap_entry *entry, st_data_t arg)
 415 {
 416     VALUE ary = (VALUE)arg;
 417
 418     rb_ary_push(ary, entry->val);
 419
 420     return ST_CONTINUE;
 421 }
 422
 423 /*
 424  * call-seq:
 425  *   map.values -> new_array
 426  *
 427  * Returns a new Array containing all values in the map.
 428  *
 429  */
 430 static VALUE
 431 wmap_values(VALUE self)
 432 {
 433     struct weakmap *w;
 434     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 435
 436     VALUE ary = rb_ary_new();
 437     wmap_foreach(w, wmap_values_i, (st_data_t)ary);
 438
 439     return ary;
 440 }
 441
 442 static int
 443 wmap_aset_replace(st_data_t *key, st_data_t *val, st_data_t new_key_ptr, int existing)
 444 {
 445     VALUE new_key = *(VALUE *)new_key_ptr;
 446     VALUE new_val = *(((VALUE *)new_key_ptr) + 1);
 447
 448     if (existing) {
 449         RUBY_ASSERT(*(VALUE *)*key == new_key);
 450     }
 451     else {
 452         struct weakmap_entry *entry = xmalloc(sizeof(struct weakmap_entry));
 453
 454         *key = (st_data_t)&entry->key;
 455         *val = (st_data_t)&entry->val;
 456     }
 457
 458     *(VALUE *)*key = new_key;
 459     *(VALUE *)*val = new_val;
 460
 461     return ST_CONTINUE;
 462 }
 463
 464 /*
 465  *  call-seq:
 466  *    map[key] = value -> value
 467  *
 468  *  Associates the given +value+ with the given +key+.
 469  *
 470  *  If the given +key+ exists, replaces its value with the given +value+;
 471  *  the ordering is not affected.
 472  */
 473 static VALUE
 474 wmap_aset(VALUE self, VALUE key, VALUE val)
 475 {
 476     struct weakmap *w;
 477     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 478
 479     VALUE pair[2] = { key, val };
 480
 481     st_update(w->table, (st_data_t)pair, wmap_aset_replace, (st_data_t)pair);
 482
 483     RB_OBJ_WRITTEN(self, Qundef, key);
 484     RB_OBJ_WRITTEN(self, Qundef, val);
 485
 486     return Qnil;
 487 }
 488
 489 /* Retrieves a weakly referenced object with the given key */
 490 static VALUE
 491 wmap_lookup(VALUE self, VALUE key)
 492 {
 493     RUBY_ASSERT(wmap_live_p(key));
 494
 495     struct weakmap *w;
 496     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 497
 498     st_data_t data;
 499     if (!st_lookup(w->table, (st_data_t)&key, &data)) return Qundef;
 500
 501     if (!wmap_live_p(*(VALUE *)data)) return Qundef;
 502
 503     return *(VALUE *)data;
 504 }
 505
 506 /*
 507  *  call-seq:
 508  *    map[key] -> value
 509  *
 510  *  Returns the value associated with the given +key+ if found.
 511  *
 512  *  If +key+ is not found, returns +nil+.
 513  */
 514 static VALUE
 515 wmap_aref(VALUE self, VALUE key)
 516 {
 517     VALUE obj = wmap_lookup(self, key);
 518     return !UNDEF_P(obj) ? obj : Qnil;
 519 }
 520
 521 /*
 522  *  call-seq:
 523  *    map.delete(key) -> value or nil
 524  *    map.delete(key) {|key| ... } -> object
 525  *
 526  *  Deletes the entry for the given +key+ and returns its associated value.
 527  *
 528  *  If no block is given and +key+ is found, deletes the entry and returns the associated value:
 529  *    m = ObjectSpace::WeakMap.new
 530  *    key = "foo"
 531  *    m[key] = 1
 532  *    m.delete(key) # => 1
 533  *    m[key] # => nil
 534  *
 535  *  If no block is given and +key+ is not found, returns +nil+.
 536  *
 537  *  If a block is given and +key+ is found, ignores the block,
 538  *  deletes the entry, and returns the associated value:
 539  *    m = ObjectSpace::WeakMap.new
 540  *    key = "foo"
 541  *    m[key] = 2
 542  *    m.delete(key) { |key| raise 'Will never happen'} # => 2
 543  *
 544  *  If a block is given and +key+ is not found,
 545  *  yields the +key+ to the block and returns the block's return value:
 546  *    m = ObjectSpace::WeakMap.new
 547  *    m.delete("nosuch") { |key| "Key #{key} not found" } # => "Key nosuch not found"
 548  */
 549 static VALUE
 550 wmap_delete(VALUE self, VALUE key)
 551 {
 552     struct weakmap *w;
 553     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 554
 555     VALUE orig_key = key;
 556     st_data_t orig_key_data = (st_data_t)&orig_key;
 557     st_data_t orig_val_data;
 558     if (st_delete(w->table, &orig_key_data, &orig_val_data)) {
 559         VALUE orig_val = *(VALUE *)orig_val_data;
 560
 561         rb_gc_remove_weak(self, (VALUE *)orig_key_data);
 562         rb_gc_remove_weak(self, (VALUE *)orig_val_data);
 563
 564         struct weakmap_entry *entry = (struct weakmap_entry *)orig_key_data;
 565         ruby_sized_xfree(entry, sizeof(struct weakmap_entry));
 566
 567         if (wmap_live_p(orig_val)) {
 568             return orig_val;
 569         }
 570     }
 571
 572     if (rb_block_given_p()) {
 573         return rb_yield(key);
 574     }
 575     else {
 576         return Qnil;
 577     }
 578 }
 579
 580 /*
 581  *  call-seq:
 582  *    map.key?(key) -> true or false
 583  *
 584  *  Returns +true+ if +key+ is a key in +self+, otherwise +false+.
 585  */
 586 static VALUE
 587 wmap_has_key(VALUE self, VALUE key)
 588 {
 589     return RBOOL(!UNDEF_P(wmap_lookup(self, key)));
 590 }
 591
 592 /*
 593  * call-seq:
 594  *   map.size -> number
 595  *
 596  * Returns the number of referenced objects
 597  */
 598 static VALUE
 599 wmap_size(VALUE self)
 600 {
 601     struct weakmap *w;
 602     TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w);
 603
 604     st_index_t n = st_table_size(w->table);
 605
 606 #if SIZEOF_ST_INDEX_T <= SIZEOF_LONG
 607     return ULONG2NUM(n);
 608 #else
 609     return ULL2NUM(n);
 610 #endif
 611 }
 612
 613 /* ===== WeakKeyMap =====
 614  *
 615  * WeakKeyMap contains one ST table which contains a pointer to the object as
 616  * the key and the object as the value. This means that the key is of the type
 617  * `VALUE *` while the value is of the type `VALUE`.
 618  *
 619  * The object is not directly stored as keys in the table because
 620  * `rb_gc_mark_weak` requires a pointer to the memory location to overwrite
 621  * when the object is reclaimed. Using a pointer into the ST table entry is not
 622  * safe because the pointer can change when the ST table is resized.
 623  *
 624  * WeakKeyMap hashes and compares using the `#hash` and `#==` methods of the
 625  * object, respectively.
 626  *
 627  * During GC and while iterating, reclaimed entries (i.e. the key points to
 628  * `Qundef`) are removed from the ST table.
 629  */
 630
 631 struct weakkeymap {
 632     st_table *table;
 633 };
 634
 635 static int
 636 wkmap_mark_table_i(st_data_t key, st_data_t val_obj, st_data_t data)
 637 {
 638     VALUE **dead_entry = (VALUE **)data;
 639     if (dead_entry != NULL) {
 640         ruby_sized_xfree(*dead_entry, sizeof(VALUE));
 641         *dead_entry = NULL;
 642     }
 643
 644     VALUE *key_ptr = (VALUE *)key;
 645
 646     if (wmap_live_p(*key_ptr)) {
 647         rb_gc_mark_weak(key_ptr);
 648         rb_gc_mark_movable((VALUE)val_obj);
 649
 650         return ST_CONTINUE;
 651     }
 652     else {
 653         *dead_entry = key_ptr;
 654
 655         return ST_DELETE;
 656     }
 657 }
 658
 659 static void
 660 wkmap_mark(void *ptr)
 661 {
 662     struct weakkeymap *w = ptr;
 663     if (w->table) {
 664         VALUE *dead_entry = NULL;
 665         st_foreach(w->table, wkmap_mark_table_i, (st_data_t)&dead_entry);
 666         if (dead_entry != NULL) {
 667             ruby_sized_xfree(dead_entry, sizeof(VALUE));
 668         }
 669     }
 670 }
 671
 672 static int
 673 wkmap_free_table_i(st_data_t key, st_data_t _val, st_data_t _arg)
 674 {
 675     ruby_sized_xfree((VALUE *)key, sizeof(VALUE));
 676     return ST_CONTINUE;
 677 }
 678
 679 static void
 680 wkmap_free(void *ptr)
 681 {
 682     struct weakkeymap *w = ptr;
 683
 684     st_foreach(w->table, wkmap_free_table_i, 0);
 685     st_free_table(w->table);
 686 }
 687
 688 static size_t
 689 wkmap_memsize(const void *ptr)
 690 {
 691     const struct weakkeymap *w = ptr;
 692
 693     size_t size = 0;
 694     size += st_memsize(w->table);
 695     /* Each key of the table takes sizeof(VALUE) in size. */
 696     size += st_table_size(w->table) * sizeof(VALUE);
 697
 698     return size;
 699 }
 700
 701 static int
 702 wkmap_compact_table_i(st_data_t key, st_data_t val_obj, st_data_t data, int _error)
 703 {
 704     VALUE **dead_entry = (VALUE **)data;
 705     if (dead_entry != NULL) {
 706         ruby_sized_xfree(*dead_entry, sizeof(VALUE));
 707         *dead_entry = NULL;
 708     }
 709
 710     VALUE *key_ptr = (VALUE *)key;
 711
 712     if (wmap_live_p(*key_ptr)) {
 713         if (*key_ptr != rb_gc_location(*key_ptr) || val_obj != rb_gc_location(val_obj)) {
 714             return ST_REPLACE;
 715         }
 716
 717         return ST_CONTINUE;
 718     }
 719     else {
 720         *dead_entry = key_ptr;
 721
 722         return ST_DELETE;
 723     }
 724 }
 725
 726 static int
 727 wkmap_compact_table_replace(st_data_t *key_ptr, st_data_t *val_ptr, st_data_t _data, int existing)
 728 {
 729     RUBY_ASSERT(existing);
 730
 731     *(VALUE *)*key_ptr = rb_gc_location(*(VALUE *)*key_ptr);
 732     *val_ptr = (st_data_t)rb_gc_location((VALUE)*val_ptr);
 733
 734     return ST_CONTINUE;
 735 }
 736
 737 static void
 738 wkmap_compact(void *ptr)
 739 {
 740     struct weakkeymap *w = ptr;
 741
 742     if (w->table) {
 743         VALUE *dead_entry = NULL;
 744         st_foreach_with_replace(w->table, wkmap_compact_table_i, wkmap_compact_table_replace, (st_data_t)&dead_entry);
 745         if (dead_entry != NULL) {
 746             ruby_sized_xfree(dead_entry, sizeof(VALUE));
 747         }
 748     }
 749 }
 750
 751 static const rb_data_type_t weakkeymap_type = {
 752     "weakkeymap",
 753     {
 754         wkmap_mark,
 755         wkmap_free,
 756         wkmap_memsize,
 757         wkmap_compact,
 758     },
 759     0, 0, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE
 760 };
 761
 762 static int
 763 wkmap_cmp(st_data_t x, st_data_t y)
 764 {
 765     VALUE x_obj = *(VALUE *)x;
 766     VALUE y_obj = *(VALUE *)y;
 767
 768     if (wmap_live_p(x_obj) && wmap_live_p(y_obj)) {
 769         return rb_any_cmp(x_obj, y_obj);
 770     }
 771     else {
 772         /* If one of the objects is dead, then they cannot be the same. */
 773         return 1;
 774     }
 775 }
 776
 777 static st_index_t
 778 wkmap_hash(st_data_t n)
 779 {
 780     VALUE obj = *(VALUE *)n;
 781     RUBY_ASSERT(wmap_live_p(obj));
 782
 783     return rb_any_hash(obj);
 784 }
 785
 786 static const struct st_hash_type wkmap_hash_type = {
 787     wkmap_cmp,
 788     wkmap_hash,
 789 };
 790
 791 static VALUE
 792 wkmap_allocate(VALUE klass)
 793 {
 794     struct weakkeymap *w;
 795     VALUE obj = TypedData_Make_Struct(klass, struct weakkeymap, &weakkeymap_type, w);
 796     w->table = st_init_table(&wkmap_hash_type);
 797     return obj;
 798 }
 799
 800 static VALUE
 801 wkmap_lookup(VALUE self, VALUE key)
 802 {
 803     struct weakkeymap *w;
 804     TypedData_Get_Struct(self, struct weakkeymap, &weakkeymap_type, w);
 805
 806     st_data_t data;
 807     if (!st_lookup(w->table, (st_data_t)&key, &data)) return Qundef;
 808
 809     return (VALUE)data;
 810 }
 811
 812 /*
 813  *  call-seq:
 814  *    map[key] -> value
 815  *
 816  *  Returns the value associated with the given +key+ if found.
 817  *
 818  *  If +key+ is not found, returns +nil+.
 819  */
 820 static VALUE
 821 wkmap_aref(VALUE self, VALUE key)
 822 {
 823     VALUE obj = wkmap_lookup(self, key);
 824     return !UNDEF_P(obj) ? obj : Qnil;
 825 }
 826
 827 struct wkmap_aset_args {
 828     VALUE new_key;
 829     VALUE new_val;
 830 };
 831
 832 static int
 833 wkmap_aset_replace(st_data_t *key, st_data_t *val, st_data_t data_args, int existing)
 834 {
 835     struct wkmap_aset_args *args = (struct wkmap_aset_args *)data_args;
 836
 837     if (!existing) {
 838         *key = (st_data_t)xmalloc(sizeof(VALUE));
 839     }
 840
 841     *(VALUE *)*key = args->new_key;
 842     *val = (st_data_t)args->new_val;
 843
 844     return ST_CONTINUE;
 845 }
 846
 847 /*
 848  *  call-seq:
 849  *    map[key] = value -> value
 850  *
 851  *  Associates the given +value+ with the given +key+
 852  *
 853  *  The reference to +key+ is weak, so when there is no other reference
 854  *  to +key+ it may be garbage collected.
 855  *
 856  *  If the given +key+ exists, replaces its value with the given +value+;
 857  *  the ordering is not affected
 858  */
 859 static VALUE
 860 wkmap_aset(VALUE self, VALUE key, VALUE val)
 861 {
 862     struct weakkeymap *w;
 863     TypedData_Get_Struct(self, struct weakkeymap, &weakkeymap_type, w);
 864
 865     if (!FL_ABLE(key) || SYMBOL_P(key) || RB_BIGNUM_TYPE_P(key) || RB_TYPE_P(key, T_FLOAT)) {
 866         rb_raise(rb_eArgError, "WeakKeyMap must be garbage collectable");
 867         UNREACHABLE_RETURN(Qnil);
 868     }
 869
 870     struct wkmap_aset_args args = {
 871         .new_key = key,
 872         .new_val = val,
 873     };
 874
 875     st_update(w->table, (st_data_t)&key, wkmap_aset_replace, (st_data_t)&args);
 876
 877     RB_OBJ_WRITTEN(self, Qundef, key);
 878     RB_OBJ_WRITTEN(self, Qundef, val);
 879
 880     return val;
 881 }
 882
 883 /*
 884  *  call-seq:
 885  *    map.delete(key) -> value or nil
 886  *    map.delete(key) {|key| ... } -> object
 887  *
 888  *  Deletes the entry for the given +key+ and returns its associated value.
 889  *
 890  *  If no block is given and +key+ is found, deletes the entry and returns the associated value:
 891  *    m = ObjectSpace::WeakKeyMap.new
 892  *    key = "foo" # to hold reference to the key
 893  *    m[key] = 1
 894  *    m.delete("foo") # => 1
 895  *    m["foo"] # => nil
 896  *
 897  *  If no block given and +key+ is not found, returns +nil+.
 898  *
 899  *  If a block is given and +key+ is found, ignores the block,
 900  *  deletes the entry, and returns the associated value:
 901  *    m = ObjectSpace::WeakKeyMap.new
 902  *    key = "foo" # to hold reference to the key
 903  *    m[key] = 2
 904  *    m.delete("foo") { |key| raise 'Will never happen'} # => 2
 905  *
 906  *  If a block is given and +key+ is not found,
 907  *  yields the +key+ to the block and returns the block's return value:
 908  *    m = ObjectSpace::WeakKeyMap.new
 909  *    m.delete("nosuch") { |key| "Key #{key} not found" } # => "Key nosuch not found"
 910  */
 911
 912 static VALUE
 913 wkmap_delete(VALUE self, VALUE key)
 914 {
 915     struct weakkeymap *w;
 916     TypedData_Get_Struct(self, struct weakkeymap, &weakkeymap_type, w);
 917
 918     VALUE orig_key = key;
 919     st_data_t orig_key_data = (st_data_t)&orig_key;
 920     st_data_t orig_val_data;
 921     if (st_delete(w->table, &orig_key_data, &orig_val_data)) {
 922         VALUE orig_val = (VALUE)orig_val_data;
 923
 924         rb_gc_remove_weak(self, (VALUE *)orig_key_data);
 925
 926         ruby_sized_xfree((VALUE *)orig_key_data, sizeof(VALUE));
 927
 928         return orig_val;
 929     }
 930
 931     if (rb_block_given_p()) {
 932         return rb_yield(key);
 933     }
 934     else {
 935         return Qnil;
 936     }
 937 }
 938
 939 /*
 940  *  call-seq:
 941  *    map.getkey(key) -> existing_key or nil
 942  *
 943  *  Returns the existing equal key if it exists, otherwise returns +nil+.
 944  *
 945  *  This might be useful for implementing caches, so that only one copy of
 946  *  some object would be used everywhere in the program:
 947  *
 948  *    value = {amount: 1, currency: 'USD'}
 949  *
 950  *    # Now if we put this object in a cache:
 951  *    cache = ObjectSpace::WeakKeyMap.new
 952  *    cache[value] = true
 953  *
 954  *    # ...we can always extract from there and use the same object:
 955  *    copy = cache.getkey({amount: 1, currency: 'USD'})
 956  *    copy.object_id == value.object_id #=> true
 957  */
 958 static VALUE
 959 wkmap_getkey(VALUE self, VALUE key)
 960 {
 961     struct weakkeymap *w;
 962     TypedData_Get_Struct(self, struct weakkeymap, &weakkeymap_type, w);
 963
 964     st_data_t orig_key;
 965     if (!st_get_key(w->table, (st_data_t)&key, &orig_key)) return Qnil;
 966
 967     return *(VALUE *)orig_key;
 968 }
 969
 970 /*
 971  *  call-seq:
 972  *    map.key?(key) -> true or false
 973  *
 974  *  Returns +true+ if +key+ is a key in +self+, otherwise +false+.
 975  */
 976 static VALUE
 977 wkmap_has_key(VALUE self, VALUE key)
 978 {
 979     return RBOOL(!UNDEF_P(wkmap_lookup(self, key)));
 980 }
 981
 982 static int
 983 wkmap_clear_i(st_data_t key, st_data_t val, st_data_t data)
 984 {
 985     VALUE self = (VALUE)data;
 986
 987     /* This WeakKeyMap may have already been marked, so we need to remove the
 988      * keys to prevent a use-after-free. */
 989     rb_gc_remove_weak(self, (VALUE *)key);
 990     return wkmap_free_table_i(key, val, 0);
 991 }
 992
 993 /*
 994  *  call-seq:
 995  *    map.clear -> self
 996  *
 997  *  Removes all map entries; returns +self+.
 998  */
 999 static VALUE
1000 wkmap_clear(VALUE self)
1001 {
1002     struct weakkeymap *w;
1003     TypedData_Get_Struct(self, struct weakkeymap, &weakkeymap_type, w);
1004
1005     st_foreach(w->table, wkmap_clear_i, (st_data_t)self);
1006     st_clear(w->table);
1007
1008     return self;
1009 }
1010
1011 /*
1012  *  call-seq:
1013  *    map.inspect -> new_string
1014  *
1015  *  Returns a new String containing informations about the map:
1016  *
1017  *    m = ObjectSpace::WeakKeyMap.new
1018  *    m[key] = value
1019  *    m.inspect # => "#<ObjectSpace::WeakKeyMap:0x00000001028dcba8 size=1>"
1020  *
1021  */
1022 static VALUE
1023 wkmap_inspect(VALUE self)
1024 {
1025     struct weakkeymap *w;
1026     TypedData_Get_Struct(self, struct weakkeymap, &weakkeymap_type, w);
1027
1028     st_index_t n = st_table_size(w->table);
1029
1030 #if SIZEOF_ST_INDEX_T <= SIZEOF_LONG
1031     const char * format = "#<%"PRIsVALUE":%p size=%lu>";
1032 #else
1033     const char * format = "#<%"PRIsVALUE":%p size=%llu>";
1034 #endif
1035
1036     VALUE str = rb_sprintf(format, rb_class_name(CLASS_OF(self)), (void *)self, n);
1037     return str;
1038 }
1039
1040 /*
1041  *  Document-class: ObjectSpace::WeakMap
1042  *
1043  *  An ObjectSpace::WeakMap is a key-value map that holds weak references
1044  *  to its keys and values, so they can be garbage-collected when there are
1045  *  no more references left.
1046  *
1047  *  Keys in the map are compared by identity.
1048  *
1049  *     m = ObjectSpace::WeakMap.new
1050  *     key1 = "foo"
1051  *     val1 = Object.new
1052  *     m[key1] = val1
1053  *
1054  *     key2 = "bar"
1055  *     val2 = Object.new
1056  *     m[key2] = val2
1057  *
1058  *     m[key1] #=> #<Object:0x0...>
1059  *     m[key2] #=> #<Object:0x0...>
1060  *
1061  *     val1 = nil # remove the other reference to value
1062  *     GC.start
1063  *
1064  *     m[key1] #=> nil
1065  *     m.keys #=> ["bar"]
1066  *
1067  *     key2 = nil # remove the other reference to key
1068  *     GC.start
1069  *
1070  *     m[key2] #=> nil
1071  *     m.keys #=> []
1072  *
1073  *  (Note that GC.start is used here only for demonstrational purposes and might
1074  *  not always lead to demonstrated results.)
1075  *
1076  *
1077  *  See also ObjectSpace::WeakKeyMap map class, which compares keys by value,
1078  *  and holds weak references only to the keys.
1079  */
1080
1081 /*
1082  *  Document-class: ObjectSpace::WeakKeyMap
1083  *
1084  *  An ObjectSpace::WeakKeyMap is a key-value map that holds weak references
1085  *  to its keys, so they can be garbage collected when there is no more references.
1086  *
1087  *  Unlike ObjectSpace::WeakMap:
1088  *
1089  *  * references to values are _strong_, so they aren't garbage collected while
1090  *    they are in the map;
1091  *  * keys are compared by value (using Object#eql?), not by identity;
1092  *  * only garbage-collectable objects can be used as keys.
1093  *
1094  *       map = ObjectSpace::WeakKeyMap.new
1095  *       val = Time.new(2023, 12, 7)
1096  *       key = "name"
1097  *       map[key] = val
1098  *
1099  *       # Value is fetched by equality: the instance of string "name" is
1100  *       # different here, but it is equal to the key
1101  *       map["name"] #=> 2023-12-07 00:00:00 +0200
1102  *
1103  *       val = nil
1104  *       GC.start
1105  *       # There are no more references to `val`, yet the pair isn't
1106  *       # garbage-collected.
1107  *       map["name"] #=> 2023-12-07 00:00:00 +0200
1108  *
1109  *       key = nil
1110  *       GC.start
1111  *       # There are no more references to `key`, key and value are
1112  *       # garbage-collected.
1113  *       map["name"] #=> nil
1114  *
1115  *  (Note that GC.start is used here only for demonstrational purposes and might
1116  *  not always lead to demonstrated results.)
1117  *
1118  *  The collection is especially useful for implementing caches of lightweight value
1119  *  objects, so that only one copy of each value representation would be stored in
1120  *  memory, but the copies that aren't used would be garbage-collected.
1121  *
1122  *    CACHE = ObjectSpace::WeakKeyMap
1123  *
1124  *    def make_value(**)
1125  *       val = ValueObject.new(**)
1126  *       if (existing = @cache.getkey(val))
1127  *          # if the object with this value exists, we return it
1128  *          existing
1129  *       else
1130  *          # otherwise, put it in the cache
1131  *          @cache[val] = true
1132  *          val
1133  *       end
1134  *    end
1135  *
1136  *  This will result in +make_value+ returning the same object for same set of attributes
1137  *  always, but the values that aren't needed anymore wouldn't be sitting in the cache forever.
1138  */
1139
1140 void
1141 Init_WeakMap(void)
1142 {
1143     VALUE rb_mObjectSpace = rb_define_module("ObjectSpace");
1144
1145     VALUE rb_cWeakMap = rb_define_class_under(rb_mObjectSpace, "WeakMap", rb_cObject);
1146     rb_define_alloc_func(rb_cWeakMap, wmap_allocate);
1147     rb_define_method(rb_cWeakMap, "[]=", wmap_aset, 2);
1148     rb_define_method(rb_cWeakMap, "[]", wmap_aref, 1);
1149     rb_define_method(rb_cWeakMap, "delete", wmap_delete, 1);
1150     rb_define_method(rb_cWeakMap, "include?", wmap_has_key, 1);
1151     rb_define_method(rb_cWeakMap, "member?", wmap_has_key, 1);
1152     rb_define_method(rb_cWeakMap, "key?", wmap_has_key, 1);
1153     rb_define_method(rb_cWeakMap, "inspect", wmap_inspect, 0);
1154     rb_define_method(rb_cWeakMap, "each", wmap_each, 0);
1155     rb_define_method(rb_cWeakMap, "each_pair", wmap_each, 0);
1156     rb_define_method(rb_cWeakMap, "each_key", wmap_each_key, 0);
1157     rb_define_method(rb_cWeakMap, "each_value", wmap_each_value, 0);
1158     rb_define_method(rb_cWeakMap, "keys", wmap_keys, 0);
1159     rb_define_method(rb_cWeakMap, "values", wmap_values, 0);
1160     rb_define_method(rb_cWeakMap, "size", wmap_size, 0);
1161     rb_define_method(rb_cWeakMap, "length", wmap_size, 0);
1162     rb_include_module(rb_cWeakMap, rb_mEnumerable);
1163
1164     VALUE rb_cWeakKeyMap = rb_define_class_under(rb_mObjectSpace, "WeakKeyMap", rb_cObject);
1165     rb_define_alloc_func(rb_cWeakKeyMap, wkmap_allocate);
1166     rb_define_method(rb_cWeakKeyMap, "[]=", wkmap_aset, 2);
1167     rb_define_method(rb_cWeakKeyMap, "[]", wkmap_aref, 1);
1168     rb_define_method(rb_cWeakKeyMap, "delete", wkmap_delete, 1);
1169     rb_define_method(rb_cWeakKeyMap, "getkey", wkmap_getkey, 1);
1170     rb_define_method(rb_cWeakKeyMap, "key?", wkmap_has_key, 1);
1171     rb_define_method(rb_cWeakKeyMap, "clear", wkmap_clear, 0);
1172     rb_define_method(rb_cWeakKeyMap, "inspect", wkmap_inspect, 0);
1173 }