From d7b0d49f047f7038cbe461fc0e36cc500e6d6630 Mon Sep 17 00:00:00 2001
From: Kunshan Wang <wks1986@gmail.com>
Date: Tue, 30 Jul 2024 14:01:28 +0800
Subject: [PATCH] Experimental support for heap traversal

---
 internal/mmtk.h         |  4 ++++
 internal/mmtk_support.h |  1 +
 iseq.c                  | 25 +++++++++++++++++++++++++
 mmtk_support.c          | 14 ++++++++++++++
 4 files changed, 44 insertions(+)

diff --git a/internal/mmtk.h b/internal/mmtk.h
index 1938db8ef760c5..f06e6b9faa1115 100644
--- a/internal/mmtk.h
+++ b/internal/mmtk.h
@@ -253,4 +253,8 @@ bool mmtk_is_object_wb_unprotected(MMTk_ObjectReference object);
 
 void mmtk_object_reference_write_post(MMTk_Mutator *mutator, MMTk_ObjectReference object);
 
+void mmtk_print_all_objects(void);
+
+void mmtk_enumerate_objects(void (*callback)(MMTk_ObjectReference, void*), void *data);
+
 #endif /* MMTK_H */
diff --git a/internal/mmtk_support.h b/internal/mmtk_support.h
index 4e2736e02c2367..99e2c6de1ac9ce 100644
--- a/internal/mmtk_support.h
+++ b/internal/mmtk_support.h
@@ -114,6 +114,7 @@ VALUE rb_mmtk_plan_name(VALUE _);
 VALUE rb_mmtk_enabled(VALUE _);
 VALUE rb_mmtk_harness_begin(VALUE _);
 VALUE rb_mmtk_harness_end(VALUE _);
+VALUE rb_mmtk_print_all_objects(VALUE _);
 
 // Debugging
 bool rb_mmtk_is_mmtk_worker(void);
diff --git a/iseq.c b/iseq.c
index d16730154e79da..b24cce1830812c 100644
--- a/iseq.c
+++ b/iseq.c
@@ -45,6 +45,9 @@
 #include "insns.inc"
 #include "insns_info.inc"
 
+// MMTk conditional compilation macros
+#include "internal/mmtk_macros.h"
+
 VALUE rb_cISeq;
 static VALUE iseqw_new(const rb_iseq_t *iseq);
 static const rb_iseq_t *iseqw_check(VALUE iseqw);
@@ -3989,6 +3992,23 @@ rb_clear_bf_ccs(void)
     rb_objspace_each_objects(clear_bf_ccs_i, NULL);
 }
 
+#if USE_MMTK
+static void
+rb_mmtk_trace_set_i(MMTk_ObjectReference objref, void *data)
+{
+    rb_event_flag_t turnon_events = *(rb_event_flag_t *)data;
+    VALUE v = (VALUE) objref;
+
+    if (rb_obj_is_iseq(v)) {
+        rb_iseq_trace_set(rb_iseq_check((rb_iseq_t *)v), turnon_events);
+    }
+    else if (clear_attr_cc(v)) {
+    }
+    else if (clear_bf_cc(v)) {
+    }
+}
+#endif
+
 static int
 trace_set_i(void *vstart, void *vend, size_t stride, void *data)
 {
@@ -4015,6 +4035,11 @@ trace_set_i(void *vstart, void *vend, size_t stride, void *data)
 void
 rb_iseq_trace_set_all(rb_event_flag_t turnon_events)
 {
+    WHEN_USING_MMTK({
+        mmtk_enumerate_objects(rb_mmtk_trace_set_i, &turnon_events);
+        return;
+    });
+
     rb_objspace_each_objects(trace_set_i, &turnon_events);
 }
 
diff --git a/mmtk_support.c b/mmtk_support.c
index 65ed7b19bc9298..0b81c6b04e2d82 100644
--- a/mmtk_support.c
+++ b/mmtk_support.c
@@ -1188,6 +1188,7 @@ rb_mmtk_define_gc_mmtk_module(void)
     rb_define_singleton_method(rb_mMMTk, "enabled?", rb_mmtk_enabled, 0);
     rb_define_singleton_method(rb_mMMTk, "harness_begin", rb_mmtk_harness_begin, 0);
     rb_define_singleton_method(rb_mMMTk, "harness_end", rb_mmtk_harness_end, 0);
+    rb_define_singleton_method(rb_mMMTk, "print_all_objects", rb_mmtk_print_all_objects, 0);
 }
 
 /*
@@ -1303,6 +1304,19 @@ rb_mmtk_harness_end(VALUE _)
     return Qnil;
 }
 
+/*
+ *  call-seq:
+ *      GC::MMTk.print_all_objects
+ *
+ *  Traverse the heap and print the addresses of all objects.
+ */
+VALUE
+rb_mmtk_print_all_objects(VALUE _)
+{
+    mmtk_print_all_objects();
+    return Qnil;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Debugging
 ////////////////////////////////////////////////////////////////////////////////