From 3592b24cdc07ed89eecb39161f21fe721a89a5de Mon Sep 17 00:00:00 2001
From: Jean Boussier <byroot@ruby-lang.org>
Date: Thu, 16 Mar 2023 08:52:22 +0000
Subject: ObjectSpace::WeakMap: clean inverse reference when an entry is
 re-assigned

[Bug #19531]

```ruby
wmap[1] = "A"
wmap[1] = "B"
```

In the example above, we need to remove the `"A" => 1` inverse reference
so that when `"A"` is GCed the `1` key isn't deleted.
---
 test/ruby/test_weakmap.rb | 17 +++++++++
 weakmap.c                 | 91 ++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 91 insertions(+), 17 deletions(-)

diff --git a/test/ruby/test_weakmap.rb b/test/ruby/test_weakmap.rb
index 9bbe2d6b81..c72e7310db 100644
--- a/test/ruby/test_weakmap.rb
+++ b/test/ruby/test_weakmap.rb
@@ -194,4 +194,21 @@ class TestWeakMap < Test::Unit::TestCase
       GC.compact
     end;
   end
+
+  def test_replaced_values_bug_19531
+    a = "A".dup
+    b = "B".dup
+
+    @wm[1] = a
+    @wm[1] = a
+    @wm[1] = a
+
+    @wm[1] = b
+    assert_equal b, @wm[1]
+
+    a = nil
+    GC.start
+
+    assert_equal b, @wm[1]
+  end
 end
diff --git a/weakmap.c b/weakmap.c
index 3a615d6868..8629d77ffb 100644
--- a/weakmap.c
+++ b/weakmap.c
@@ -148,25 +148,40 @@ wmap_live_p(VALUE obj)
 }
 
 static int
-wmap_final_func(st_data_t *key, st_data_t *value, st_data_t arg, int existing)
+wmap_remove_inverse_ref(st_data_t *key, st_data_t *val, st_data_t arg, int existing)
 {
-    VALUE wmap, *ptr, size, i, j;
     if (!existing) return ST_STOP;
-    wmap = (VALUE)arg, ptr = (VALUE *)*value;
-    for (i = j = 1, size = ptr[0]; i <= size; ++i) {
-        if (ptr[i] != wmap) {
-            ptr[j++] = ptr[i];
-        }
-    }
-    if (j == 1) {
-        ruby_sized_xfree(ptr, i * sizeof(VALUE));
+
+    VALUE old_ref = (VALUE)arg;
+
+    VALUE *values = (VALUE *)*val;
+    VALUE size = values[0];
+
+    if (size == 1) {
+        // fast path, we only had one backref
+        RUBY_ASSERT(values[1] == old_ref);
+        ruby_sized_xfree(values, 2 * sizeof(VALUE));
         return ST_DELETE;
     }
-    if (j < i) {
-        SIZED_REALLOC_N(ptr, VALUE, j, i);
-        ptr[0] = j - 1;
-        *value = (st_data_t)ptr;
+
+    bool found = false;
+    VALUE index = 1;
+    for (; index <= size; index++) {
+        if (values[index] == old_ref) {
+            found = true;
+            break;
+        }
     }
+    if (!found) return ST_STOP;
+
+    if (size > index) {
+        MEMMOVE(&values[index], &values[index + 1], VALUE, size - index);
+    }
+
+    size -= 1;
+    values[0] = size;
+    SIZED_REALLOC_N(values, VALUE, size + 1, size + 2);
+    *val = (st_data_t)values;
     return ST_CONTINUE;
 }
 
@@ -199,7 +214,7 @@ wmap_finalize(RB_BLOCK_CALL_FUNC_ARGLIST(objid, self))
     wmap = (st_data_t)obj;
     if (st_delete(w->wmap2obj, &wmap, &orig)) {
         wmap = (st_data_t)obj;
-        st_update(w->obj2wmap, orig, wmap_final_func, wmap);
+        st_update(w->obj2wmap, orig, wmap_remove_inverse_ref, wmap);
     }
     return self;
 }
@@ -387,6 +402,14 @@ wmap_aset_update(st_data_t *key, st_data_t *val, st_data_t arg, int existing)
     VALUE size, *ptr, *optr;
     if (existing) {
         size = (ptr = optr = (VALUE *)*val)[0];
+
+        for (VALUE index = 1; index <= size; index++) {
+            if (ptr[index] == (VALUE)arg) {
+                // The reference was already registered.
+                return ST_STOP;
+            }
+        }
+
         ++size;
         SIZED_REALLOC_N(ptr, VALUE, size + 1, size);
     }
@@ -414,6 +437,23 @@ nonspecial_obj_id(VALUE obj)
 #endif
 }
 
+struct wmap_aset_replace_args {
+    VALUE new_value;
+    VALUE old_value;
+};
+
+static int
+wmap_aset_replace_value(st_data_t *key, st_data_t *val, st_data_t _args, int existing)
+{
+    struct wmap_aset_replace_args *args = (struct wmap_aset_replace_args *)_args;
+
+    if (existing) {
+        args->old_value = *val;
+    }
+    *val = (st_data_t)args->new_value;
+    return ST_CONTINUE;
+}
+
 /* Creates a weak reference from the given key to the given value */
 static VALUE
 wmap_aset(VALUE self, VALUE key, VALUE value)
@@ -428,8 +468,25 @@ wmap_aset(VALUE self, VALUE key, VALUE value)
         rb_define_finalizer_no_check(key, w->final);
     }
 
-    st_update(w->obj2wmap, (st_data_t)value, wmap_aset_update, key);
-    st_insert(w->wmap2obj, (st_data_t)key, (st_data_t)value);
+    struct wmap_aset_replace_args aset_args = {
+        .new_value = value,
+        .old_value = Qundef,
+    };
+    st_update(w->wmap2obj, (st_data_t)key, wmap_aset_replace_value, (st_data_t)&aset_args);
+
+    // If the value is unchanged, we have nothing to do.
+    if (value != aset_args.old_value) {
+        if (!UNDEF_P(aset_args.old_value) && FL_ABLE(aset_args.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)aset_args.old_value, wmap_remove_inverse_ref, key);
+        }
+
+        if (FL_ABLE(value)) {
+            // If the value has no finalizer, we don't need to keep the inverse reference
+            st_update(w->obj2wmap, (st_data_t)value, wmap_aset_update, key);
+        }
+    }
+
     return nonspecial_obj_id(value);
 }
 
-- 
cgit v1.2.3