Implement ObjectSpace::WeakMap#delete and ObjectSpace::WeakKeyMap#delete

[Feature #19561]

It's useful to be able to remove references from weak maps.

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

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

test/ruby/weakkeymap.rb → 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

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

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);