From 15765eac0ae156afe69b53ab317c4096f2c2c0ec Mon Sep 17 00:00:00 2001
From: Peter Zhu <peter@peterzhu.ca>
Date: Fri, 13 Dec 2024 12:20:30 -0500
Subject: [PATCH] Fix ObjectSpace.trace_object_allocations for compaction

We need to reinsert into the ST table when an object moves because it is
a numtable that hashes on the object address, so when an object moves we
need to reinsert it rather than just updating the key.
---
 ext/objspace/object_tracing.c  | 28 +++++++++++++++-------------
 test/objspace/test_objspace.rb | 18 ++++++++++++++++++
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/ext/objspace/object_tracing.c b/ext/objspace/object_tracing.c
index b92fda6826..8ea35ef7b7 100644
--- a/ext/objspace/object_tracing.c
+++ b/ext/objspace/object_tracing.c
@@ -190,31 +190,33 @@ allocation_info_tracer_memsize(const void *ptr)
 }
 
 static int
-hash_foreach_should_replace_key(st_data_t key, st_data_t value, st_data_t argp, int error)
+allocation_info_tracer_compact_update_object_table_i(st_data_t key, st_data_t value, st_data_t data)
 {
-    VALUE allocated_object = (VALUE)key;
-    if (allocated_object != rb_gc_location(allocated_object)) {
-        return ST_REPLACE;
+    st_table *table = (st_table *)data;
+
+    if (key != rb_gc_location(key)) {
+        DURING_GC_COULD_MALLOC_REGION_START();
+        {
+            st_insert(table, rb_gc_location(key), value);
+        }
+        DURING_GC_COULD_MALLOC_REGION_END();
+
+        return ST_DELETE;
     }
 
     return ST_CONTINUE;
 }
 
-static int
-hash_replace_key(st_data_t *key, st_data_t *value, st_data_t argp, int existing)
-{
-    *key = rb_gc_location((VALUE)*key);
-
-    return ST_CONTINUE;
-}
-
 static void
 allocation_info_tracer_compact(void *ptr)
 {
     struct traceobj_arg *trace_arg = (struct traceobj_arg *)ptr;
 
     if (trace_arg->object_table &&
-            st_foreach_with_replace(trace_arg->object_table, hash_foreach_should_replace_key, hash_replace_key, 0)) {
+            st_foreach(
+                trace_arg->object_table,
+                allocation_info_tracer_compact_update_object_table_i,
+                (st_data_t)trace_arg->object_table)) {
         rb_raise(rb_eRuntimeError, "hash modified during iteration");
     }
 }
diff --git a/test/objspace/test_objspace.rb b/test/objspace/test_objspace.rb
index 3f08413378..394b2e49ad 100644
--- a/test/objspace/test_objspace.rb
+++ b/test/objspace/test_objspace.rb
@@ -287,6 +287,24 @@ class TestObjSpace < Test::Unit::TestCase
     assert true # success
   end
 
+  def test_trace_object_allocations_compaction
+    omit "compaction is not supported on this platform" unless GC.respond_to?(:compact)
+
+    assert_separately(%w(-robjspace), <<~RUBY)
+      ObjectSpace.trace_object_allocations do
+        objs = 100.times.map do
+          Object.new
+        end
+
+        assert_equal(__FILE__, ObjectSpace.allocation_sourcefile(objs[0]))
+
+        GC.verify_compaction_references(expand_heap: true, toward: :empty)
+
+        assert_equal(__FILE__, ObjectSpace.allocation_sourcefile(objs[0]))
+      end
+    RUBY
+  end
+
   def test_dump_flags
     # Ensure that the fstring is promoted to old generation
     4.times { GC.start }