From 0fc8422a05e21a3d1247794e8ec34a43010d51c9 Mon Sep 17 00:00:00 2001
From: Peter Zhu <peter@peterzhu.ca>
Date: Thu, 12 Sep 2024 10:27:55 -0400
Subject: [PATCH] Move checks for heap traversal to rb_gc_mark_weak

If we are during heap traversal, we don't want to call rb_gc_impl_mark_weak.
This commit moves that check from rb_gc_impl_mark_weak to rb_gc_mark_weak.
---
 gc.c         | 11 ++++++++++-
 gc/default.c |  2 --
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/gc.c b/gc.c
index aed788d39b..ee8b7bc376 100644
--- a/gc.c
+++ b/gc.c
@@ -2104,7 +2104,16 @@ rb_gc_mark_weak(VALUE *ptr)
 {
     if (RB_SPECIAL_CONST_P(*ptr)) return;
 
-    rb_gc_impl_mark_weak(rb_gc_get_objspace(), ptr);
+    rb_vm_t *vm = GET_VM();
+    void *objspace = vm->gc.objspace;
+    if (LIKELY(vm->gc.mark_func_data == NULL)) {
+        GC_ASSERT(rb_gc_impl_during_gc_p(objspace));
+
+        rb_gc_impl_mark_weak(objspace, ptr);
+    }
+    else {
+        GC_ASSERT(!rb_gc_impl_during_gc_p(objspace));
+    }
 }
 
 void
diff --git a/gc/default.c b/gc/default.c
index b5fa1a6b4b..94e5acf3d7 100644
--- a/gc/default.c
+++ b/gc/default.c
@@ -4607,8 +4607,6 @@ rb_gc_impl_mark_weak(void *objspace_ptr, VALUE *ptr)
 {
     rb_objspace_t *objspace = objspace_ptr;
 
-    if (RB_UNLIKELY(!during_gc)) return;
-
     GC_ASSERT(objspace->rgengc.parent_object == 0 || FL_TEST(objspace->rgengc.parent_object, FL_WB_PROTECTED));
 
     VALUE obj = *ptr;