From 548086b34e3dd125edabf5dc1e46b891fad3ea9c Mon Sep 17 00:00:00 2001
From: Jean Boussier <byroot@ruby-lang.org>
Date: Tue, 14 Mar 2023 09:12:55 +0100
Subject: [PATCH] ObjectSpace::WeakMap: fix compaction support

[Bug #19529]

`rb_gc_update_tbl_refs` can't be used on `w->obj2wmap` because it's
not a `VALUE -> VALUE` table, but a `VALUE -> VALUE *` table, so
we need some dedicated iterator.
---
 test/ruby/test_weakmap.rb |  8 ++++++++
 weakmap.c                 | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/test/ruby/test_weakmap.rb b/test/ruby/test_weakmap.rb
index 86550cc221..d84eb8241e 100644
--- a/test/ruby/test_weakmap.rb
+++ b/test/ruby/test_weakmap.rb
@@ -176,4 +176,12 @@ class TestWeakMap < Test::Unit::TestCase
       end
     end;
   end
+
+  def test_compaction_bug_19529
+    obj = Object.new
+    100.times do |i|
+      GC.compact
+      @wm[i] = obj
+    end
+  end
 end
diff --git a/weakmap.c b/weakmap.c
index e6dbb325d7..2145e4d2f8 100644
--- a/weakmap.c
+++ b/weakmap.c
@@ -12,12 +12,47 @@ struct weakmap {
     VALUE final;
 };
 
+static int
+wmap_replace_ref(st_data_t *key, st_data_t *value, st_data_t _argp, int existing)
+{
+    *key = rb_gc_location((VALUE)*key);
+
+    VALUE *values = (VALUE *)value;
+    VALUE size = values[0];
+
+    for (VALUE index = 1; index <= size; index++) {
+        values[index] = rb_gc_location(values[index]);
+    }
+
+    return ST_CONTINUE;
+}
+
+static int
+wmap_foreach_replace(st_data_t key, st_data_t value, st_data_t _argp, int error)
+{
+    if (rb_gc_location((VALUE)key) != (VALUE)key) {
+        return ST_REPLACE;
+    }
+
+    VALUE *values = (VALUE *)value;
+    VALUE size = values[0];
+
+    for (VALUE index = 1; index <= size; index++) {
+        VALUE val = values[index];
+        if (rb_gc_location(val) != val) {
+            return ST_REPLACE;
+        }
+    }
+
+    return ST_CONTINUE;
+}
+
 static void
 wmap_compact(void *ptr)
 {
     struct weakmap *w = ptr;
     if (w->wmap2obj) rb_gc_update_tbl_refs(w->wmap2obj);
-    if (w->obj2wmap) rb_gc_update_tbl_refs(w->obj2wmap);
+    if (w->obj2wmap) st_foreach_with_replace(w->obj2wmap, wmap_foreach_replace, wmap_replace_ref, (st_data_t)NULL);
     w->final = rb_gc_location(w->final);
 }