From 52449b5b75068872ce568ed00c4c7fb8e6c28072 Mon Sep 17 00:00:00 2001 From: Jean Boussier Date: Thu, 30 Mar 2023 16:23:14 +0200 Subject: [PATCH] Implement ObjectSpace::WeakMap#delete and ObjectSpace::WeakKeyMap#delete [Feature #19561] It's useful to be able to remove references from weak maps. --- .../objectspace/weakkeymap/delete_spec.rb | 40 +++++++++ .../core/objectspace/weakmap/delete_spec.rb | 30 +++++++ .../{weakkeymap.rb => test_weakkeymap.rb} | 21 ++++- test/ruby/test_weakmap.rb | 16 ++++ weakmap.c | 82 +++++++++++++++++++ 5 files changed, 185 insertions(+), 4 deletions(-) create mode 100644 spec/ruby/core/objectspace/weakkeymap/delete_spec.rb create mode 100644 spec/ruby/core/objectspace/weakmap/delete_spec.rb rename test/ruby/{weakkeymap.rb => test_weakkeymap.rb} (85%) diff --git a/spec/ruby/core/objectspace/weakkeymap/delete_spec.rb b/spec/ruby/core/objectspace/weakkeymap/delete_spec.rb new file mode 100644 index 00000000000000..6e534b8ea8b6a6 --- /dev/null +++ b/spec/ruby/core/objectspace/weakkeymap/delete_spec.rb @@ -0,0 +1,40 @@ +require_relative '../../../spec_helper' + +ruby_version_is '3.3' do + describe "ObjectSpace::WeakKeyMap#delete" do + it "removes the entry and returns the deleted value" do + m = ObjectSpace::WeakKeyMap.new + key = Object.new + value = Object.new + m[key] = value + + m.delete(key).should == value + m.key?(key).should == false + end + + it "uses equality semantic" do + m = ObjectSpace::WeakKeyMap.new + key = "foo".upcase + value = Object.new + m[key] = value + + m.delete("foo".upcase).should == value + m.key?(key).should == false + end + + it "calls supplied block if the key is not found" do + key = Object.new + m = ObjectSpace::WeakKeyMap.new + return_value = m.delete(key) do |yielded_key| + yielded_key.should == key + 5 + end + return_value.should == 5 + end + + it "returns nil if the key is not found when no block is given" do + m = ObjectSpace::WeakMap.new + m.delete(Object.new).should == nil + end + end +end diff --git a/spec/ruby/core/objectspace/weakmap/delete_spec.rb b/spec/ruby/core/objectspace/weakmap/delete_spec.rb new file mode 100644 index 00000000000000..302de264fb2998 --- /dev/null +++ b/spec/ruby/core/objectspace/weakmap/delete_spec.rb @@ -0,0 +1,30 @@ +require_relative '../../../spec_helper' + +ruby_version_is '3.3' do + describe "ObjectSpace::WeakMap#delete" do + it "removes the entry and returns the deleted value" do + m = ObjectSpace::WeakMap.new + key = Object.new + value = Object.new + m[key] = value + + m.delete(key).should == value + m.key?(key).should == false + end + + it "calls supplied block if the key is not found" do + key = Object.new + m = ObjectSpace::WeakMap.new + return_value = m.delete(key) do |yielded_key| + yielded_key.should == key + 5 + end + return_value.should == 5 + end + + it "returns nil if the key is not found when no block is given" do + m = ObjectSpace::WeakMap.new + m.delete(Object.new).should == nil + end + end +end diff --git a/test/ruby/weakkeymap.rb b/test/ruby/test_weakkeymap.rb similarity index 85% rename from test/ruby/weakkeymap.rb rename to test/ruby/test_weakkeymap.rb index 20bfe5ec2c9bf2..5df53749ca04ae 100644 --- a/test/ruby/weakkeymap.rb +++ b/test/ruby/test_weakkeymap.rb @@ -33,11 +33,24 @@ def test_getkey end def test_key? - 1.times do - assert_weak_include(:key?, "foo") + assert_weak_include(:key?, "foo") + assert_not_send([@wm, :key?, "bar"]) + end + + def test_delete + k1 = "foo" + x1 = Object.new + @wm[k1] = x1 + assert_equal x1, @wm[k1] + assert_equal x1, @wm.delete(k1) + assert_nil @wm[k1] + assert_nil @wm.delete(k1) + + fallback = @wm.delete(k1) do |key| + assert_equal k1, key + 42 end - GC.start - assert_not_send([@wm, :key?, "FOO".downcase]) + assert_equal 42, fallback end def test_clear diff --git a/test/ruby/test_weakmap.rb b/test/ruby/test_weakmap.rb index 7fc956dfae4747..f9358e11f30afe 100644 --- a/test/ruby/test_weakmap.rb +++ b/test/ruby/test_weakmap.rb @@ -82,6 +82,22 @@ def test_inspect_garbage @wm.inspect) end + def test_delete + k1 = "foo" + x1 = Object.new + @wm[k1] = x1 + assert_equal x1, @wm[k1] + assert_equal x1, @wm.delete(k1) + assert_nil @wm[k1] + assert_nil @wm.delete(k1) + + fallback = @wm.delete(k1) do |key| + assert_equal k1, key + 42 + end + assert_equal 42, fallback + end + def test_each m = __callee__[/test_(.*)/, 1] x1 = Object.new diff --git a/weakmap.c b/weakmap.c index f5b9a5fbf496ef..17883f05a44f9c 100644 --- a/weakmap.c +++ b/weakmap.c @@ -515,6 +515,31 @@ wmap_aref(VALUE self, VALUE key) return !UNDEF_P(obj) ? obj : Qnil; } +/* Delete the given key from the map */ +static VALUE +wmap_delete(VALUE self, VALUE key) +{ + assert(wmap_live_p(key)); + + struct weakmap *w; + TypedData_Get_Struct(self, struct weakmap, &weakmap_type, w); + + VALUE old_value = Qnil; + if (st_delete(w->wmap2obj, (st_data_t *)&key, (st_data_t *)&old_value)) { + if (FL_ABLE(old_value)) { + // That key existed and had an inverse reference, we need to clear the outdated inverse reference. + st_update(w->obj2wmap, (st_data_t)old_value, wmap_remove_inverse_ref, key); + } + return old_value; + } + else if (rb_block_given_p()) { + return rb_yield(key); + } + else { + return Qnil; + } +} + /* Returns +true+ if +key+ is registered */ static VALUE wmap_has_key(VALUE self, VALUE key) @@ -751,6 +776,61 @@ wkmap_aset(VALUE self, VALUE key, VALUE value) return value; } +/* + * call-seq: + * map.delete(key) -> value or nil + * map.delete(key) {|key| ... } -> object + * + * Deletes the entry for the given +key+ and returns its associated value. + * + * If no block is given and +key+ is found, deletes the entry and returns the associated value: + * m = ObjectSpace::WeakKeyMap.new + * m["foo"] = 1 + * m.delete("foo") # => 1 + * m["foo"] # => nil + * + * If no block given and +key+ is not found, returns +nil+. + * + * If a block is given and +key+ is found, ignores the block, + * deletes the entry, and returns the associated value: + * m = ObjectSpace::WeakKeyMap.new + * m["foo"] = 2 + * h.delete("foo") { |key| raise 'Will never happen'} # => 2 + * + * If a block is given and +key+ is not found, + * calls the block and returns the block's return value: + * m = ObjectSpace::WeakKeyMap.new + * h.delete("nosuch") { |key| "Key #{key} not found" } # => "Key nosuch not found" + */ + +static VALUE +wkmap_delete(VALUE self, VALUE key) +{ + struct weakkeymap *w; + TypedData_Get_Struct(self, struct weakkeymap, &weakkeymap_type, w); + + st_index_t hash = rb_any_hash(key); + weakkeymap_entry_t lookup_entry = {key, hash}; + weakkeymap_entry_t *deleted_entry = NULL; + if (st_get_key(w->map, (st_data_t)&lookup_entry, (st_data_t *)&deleted_entry)) { + st_data_t deleted_value; + if (st_delete(w->map, (st_data_t *)&deleted_entry, &deleted_value)) { + xfree(deleted_entry); + st_delete(w->obj2hash, (st_data_t *)key, &hash); + return (VALUE)deleted_value; + } + else { + rb_bug("WeakKeyMap: miss on delete, corrupted memory?"); + } + } + else if (rb_block_given_p()) { + return rb_yield(key); + } + else { + return Qnil; + } +} + /* * call-seq: * map.getkey(key) -> existing_key or nil @@ -870,6 +950,7 @@ Init_WeakMap(void) rb_define_alloc_func(rb_cWeakMap, wmap_allocate); rb_define_method(rb_cWeakMap, "[]=", wmap_aset, 2); rb_define_method(rb_cWeakMap, "[]", wmap_aref, 1); + rb_define_method(rb_cWeakMap, "delete", wmap_delete, 1); rb_define_method(rb_cWeakMap, "include?", wmap_has_key, 1); rb_define_method(rb_cWeakMap, "member?", wmap_has_key, 1); rb_define_method(rb_cWeakMap, "key?", wmap_has_key, 1); @@ -888,6 +969,7 @@ Init_WeakMap(void) rb_define_alloc_func(rb_cWeakKeyMap, wkmap_allocate); rb_define_method(rb_cWeakKeyMap, "[]=", wkmap_aset, 2); rb_define_method(rb_cWeakKeyMap, "[]", wkmap_aref, 1); + rb_define_method(rb_cWeakKeyMap, "delete", wkmap_delete, 1); rb_define_method(rb_cWeakKeyMap, "getkey", wkmap_getkey, 1); rb_define_method(rb_cWeakKeyMap, "key?", wkmap_has_key, 1); rb_define_method(rb_cWeakKeyMap, "clear", wkmap_clear, 0);