From f358d8666fe4d55dfbff3e7211111a1c95a46e38 Mon Sep 17 00:00:00 2001
From: "Jeong, YunWon" <jeong@youknowone.org>
Date: Tue, 3 Feb 2026 15:56:47 +0900
Subject: [PATCH] Improve object traversal for heap types and mro
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Fix heap type instance traversal to include type reference
  (enables correct cycle detection for instance ↔ type references)
- Enable mro traversal in PyType (was previously disabled)
---
 crates/vm/src/builtins/type.rs          |  3 +--
 crates/vm/src/object/traverse_object.rs | 23 ++++++++++++++++-------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs
index 5e984ff6f3d..8944d051a1e 100644
--- a/crates/vm/src/builtins/type.rs
+++ b/crates/vm/src/builtins/type.rs
@@ -49,8 +49,7 @@ unsafe impl crate::object::Traverse for PyType {
     fn traverse(&self, tracer_fn: &mut crate::object::TraverseFn<'_>) {
         self.base.traverse(tracer_fn);
         self.bases.traverse(tracer_fn);
-        // mro contains self as mro[0], so skip traversing to avoid circular reference
-        // self.mro.traverse(tracer_fn);
+        self.mro.traverse(tracer_fn);
         self.subclasses.traverse(tracer_fn);
         self.attributes
             .read_recursive()
diff --git a/crates/vm/src/object/traverse_object.rs b/crates/vm/src/object/traverse_object.rs
index b297864245e..af90e31934b 100644
--- a/crates/vm/src/object/traverse_object.rs
+++ b/crates/vm/src/object/traverse_object.rs
@@ -45,9 +45,16 @@ unsafe impl Traverse for InstanceDict {
 unsafe impl Traverse for PyInner<Erased> {
     /// Because PyObject hold a `PyInner<Erased>`, so we need to trace it
     fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
-        // 1. trace `dict` and `slots` field(`typ` can't trace for it's a AtomicRef while is leaked by design)
-        // 2. call vtable's trace function to trace payload
-        // self.typ.trace(tracer_fn);
+        // For heap type instances, traverse the type reference.
+        // PyAtomicRef holds a strong reference (via PyRef::leak), so GC must
+        // account for it to correctly detect instance ↔ type cycles.
+        // Static types are always alive and don't need this.
+        let typ = &*self.typ;
+        if typ.heaptype_ext.is_some() {
+            // Safety: Py<PyType> and PyObject share the same memory layout
+            let typ_obj: &PyObject = unsafe { &*(typ as *const _ as *const PyObject) };
+            tracer_fn(typ_obj);
+        }
         self.dict.traverse(tracer_fn);
         // weak_list is inline atomic pointers, no heap allocation, no trace
         self.slots.traverse(tracer_fn);
@@ -64,10 +71,12 @@ unsafe impl Traverse for PyInner<Erased> {
 unsafe impl<T: MaybeTraverse> Traverse for PyInner<T> {
     /// Type is known, so we can call `try_trace` directly instead of using erased type vtable
     fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
-        // 1. trace `dict` and `slots` field(`typ` can't trace for it's a AtomicRef while is leaked by design)
-        // 2. call corresponding `try_trace` function to trace payload
-        // (No need to call vtable's trace function because we already know the type)
-        // self.typ.trace(tracer_fn);
+        // For heap type instances, traverse the type reference (same as erased version)
+        let typ = &*self.typ;
+        if typ.heaptype_ext.is_some() {
+            let typ_obj: &PyObject = unsafe { &*(typ as *const _ as *const PyObject) };
+            tracer_fn(typ_obj);
+        }
         self.dict.traverse(tracer_fn);
         // weak_list is inline atomic pointers, no heap allocation, no trace
         self.slots.traverse(tracer_fn);