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