Add GC infrastructure: tracking bits, tp_clear (#6977)

youknowone · web-flow · commit f71fe9bf3add · 2026-02-03T20:00:34.000+09:00
GC bit operations (_PyObject_GC_TRACK/UNTRACK equivalent):
- Add set_gc_bit() helper for atomic GC bit manipulation
- Add set_gc_tracked() / clear_gc_tracked() methods
- Update is_gc_tracked() to use GcBits::TRACKED flag
- Call set_gc_tracked() in track_object()
- Call clear_gc_tracked() for static types (they are immortal)

tp_clear infrastructure (for breaking reference cycles):
- Add try_clear_obj() function to call payload's try_clear
- Add clear field to PyObjVTable
- Add clear() method to PyType's Traverse impl
diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs
@@ -57,6 +57,33 @@ unsafe impl crate::object::Traverse for PyType {
             .map(|(_, v)| v.traverse(tracer_fn))
             .count();
     }
+
+    /// type_clear: break reference cycles in type objects
+    fn clear(&mut self, out: &mut Vec<crate::PyObjectRef>) {
+        if let Some(base) = self.base.take() {
+            out.push(base.into());
+        }
+        if let Some(mut guard) = self.bases.try_write() {
+            for base in guard.drain(..) {
+                out.push(base.into());
+            }
+        }
+        if let Some(mut guard) = self.mro.try_write() {
+            for typ in guard.drain(..) {
+                out.push(typ.into());
+            }
+        }
+        if let Some(mut guard) = self.subclasses.try_write() {
+            for weak in guard.drain(..) {
+                out.push(weak.into());
+            }
+        }
+        if let Some(mut guard) = self.attributes.try_write() {
+            for (_, val) in guard.drain(..) {
+                out.push(val);
+            }
+        }
+    }
 }
 
 // PyHeapTypeObject in CPython
@@ -393,6 +420,11 @@ impl PyType {
             metaclass,
             None,
         );
+
+        // Static types are not tracked by GC.
+        // They are immortal and never participate in collectable cycles.
+        new_type.as_object().clear_gc_tracked();
+
         new_type.mro.write().insert(0, new_type.clone());
 
         // Note: inherit_slots is called in PyClassImpl::init_class after
diff --git a/crates/vm/src/gc_state.rs b/crates/vm/src/gc_state.rs
@@ -236,6 +236,10 @@ impl GcState {
     pub unsafe fn track_object(&self, obj: NonNull<PyObject>) {
         let gc_ptr = GcObjectPtr(obj);
 
+        // _PyObject_GC_TRACK
+        let obj_ref = unsafe { obj.as_ref() };
+        obj_ref.set_gc_tracked();
+
         // Add to generation 0 tracking first (for correct gc_refs algorithm)
         // Only increment count if we successfully add to the set
         if let Ok(mut gen0) = self.generation_objects[0].write()
diff --git a/crates/vm/src/object/core.rs b/crates/vm/src/object/core.rs
@@ -81,14 +81,28 @@ use core::{
 #[derive(Debug)]
 pub(super) struct Erased;
 
-/// Default dealloc: handles __del__, weakref clearing, and memory free.
+/// Default dealloc: handles __del__, weakref clearing, tp_clear, and memory free.
 /// Equivalent to subtype_dealloc in CPython.
 pub(super) unsafe fn default_dealloc<T: PyPayload>(obj: *mut PyObject) {
     let obj_ref = unsafe { &*(obj as *const PyObject) };
     if let Err(()) = obj_ref.drop_slow_inner() {
         return; // resurrected by __del__
     }
+
+    // Extract child references before deallocation to break circular refs (tp_clear).
+    // This ensures that when edges are dropped after the object is freed,
+    // any pointers back to this object are already gone.
+    let mut edges = Vec::new();
+    if let Some(clear_fn) = obj_ref.0.vtable.clear {
+        unsafe { clear_fn(obj, &mut edges) };
+    }
+
+    // Deallocate the object memory
     drop(unsafe { Box::from_raw(obj as *mut PyInner<T>) });
+
+    // Drop child references - may trigger recursive destruction.
+    // The object is already deallocated, so circular refs are broken.
+    drop(edges);
 }
 pub(super) unsafe fn debug_obj<T: PyPayload + core::fmt::Debug>(
     x: &PyObject,
@@ -105,6 +119,12 @@ pub(super) unsafe fn try_traverse_obj<T: PyPayload>(x: &PyObject, tracer_fn: &mu
     payload.try_traverse(tracer_fn)
 }
 
+/// Call `try_clear` on payload to extract child references (tp_clear)
+pub(super) unsafe fn try_clear_obj<T: PyPayload>(x: *mut PyObject, out: &mut Vec<PyObjectRef>) {
+    let x = unsafe { &mut *(x as *mut PyInner<T>) };
+    x.payload.try_clear(out);
+}
+
 bitflags::bitflags! {
     /// GC bits for free-threading support (like ob_gc_bits in Py_GIL_DISABLED)
     /// These bits are stored in a separate atomic field for lock-free access.
@@ -963,10 +983,27 @@ impl PyObject {
     /// _PyGC_SET_FINALIZED in Py_GIL_DISABLED mode.
     #[inline]
     fn set_gc_finalized(&self) {
-        // Atomic RMW to avoid clobbering other concurrent bit updates.
+        self.set_gc_bit(GcBits::FINALIZED);
+    }
+
+    /// Set a GC bit atomically.
+    #[inline]
+    pub(crate) fn set_gc_bit(&self, bit: GcBits) {
+        self.0.gc_bits.fetch_or(bit.bits(), Ordering::Relaxed);
+    }
+
+    /// _PyObject_GC_TRACK
+    #[inline]
+    pub(crate) fn set_gc_tracked(&self) {
+        self.set_gc_bit(GcBits::TRACKED);
+    }
+
+    /// _PyObject_GC_UNTRACK
+    #[inline]
+    pub(crate) fn clear_gc_tracked(&self) {
         self.0
             .gc_bits
-            .fetch_or(GcBits::FINALIZED.bits(), Ordering::Relaxed);
+            .fetch_and(!GcBits::TRACKED.bits(), Ordering::Relaxed);
     }
 
     #[inline(always)] // the outer function is never inlined
@@ -1046,13 +1083,9 @@ impl PyObject {
         *self.0.slots[offset].write() = value;
     }
 
-    /// Check if this object is tracked by the garbage collector.
-    /// Returns true if the object has a trace function or has an instance dict.
+    /// _PyObject_GC_IS_TRACKED
     pub fn is_gc_tracked(&self) -> bool {
-        if self.0.vtable.trace.is_some() {
-            return true;
-        }
-        self.0.dict.is_some()
+        GcBits::from_bits_retain(self.0.gc_bits.load(Ordering::Relaxed)).contains(GcBits::TRACKED)
     }
 
     /// Get the referents (objects directly referenced) of this object.
@@ -1277,13 +1310,28 @@ impl<T: PyPayload> PyRef<T> {
     }
 }
 
-impl<T: PyPayload + core::fmt::Debug> PyRef<T> {
+impl<T: PyPayload + crate::object::MaybeTraverse + core::fmt::Debug> PyRef<T> {
     #[inline(always)]
     pub fn new_ref(payload: T, typ: crate::builtins::PyTypeRef, dict: Option<PyDictRef>) -> Self {
+        let has_dict = dict.is_some();
+        let is_heaptype = typ.heaptype_ext.is_some();
         let inner = Box::into_raw(PyInner::new(payload, typ, dict));
-        Self {
-            ptr: unsafe { NonNull::new_unchecked(inner.cast::<Py<T>>()) },
+        let ptr = unsafe { NonNull::new_unchecked(inner.cast::<Py<T>>()) };
+
+        // Track object if:
+        // - HAS_TRAVERSE is true (Rust payload implements Traverse), OR
+        // - has instance dict (user-defined class instances), OR
+        // - heap type (all heap type instances are GC-tracked, like Py_TPFLAGS_HAVE_GC)
+        if <T as crate::object::MaybeTraverse>::HAS_TRAVERSE || has_dict || is_heaptype {
+            let gc = crate::gc_state::gc_state();
+            unsafe {
+                gc.track_object(ptr.cast());
+            }
+            // Check if automatic GC should run
+            gc.maybe_collect();
         }
+
+        Self { ptr }
     }
 }
 
@@ -1546,6 +1594,12 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) {
         heaptype_ext: None,
     };
     let weakref_type = PyRef::new_ref(weakref_type, type_type.clone(), None);
+    // Static type: untrack from GC (was tracked by new_ref because PyType has HAS_TRAVERSE)
+    unsafe {
+        crate::gc_state::gc_state()
+            .untrack_object(core::ptr::NonNull::from(weakref_type.as_object()));
+    }
+    weakref_type.as_object().clear_gc_tracked();
     // weakref's mro is [weakref, object]
     weakref_type.mro.write().insert(0, weakref_type.clone());
 
diff --git a/crates/vm/src/object/traverse_object.rs b/crates/vm/src/object/traverse_object.rs
@@ -2,10 +2,10 @@ use alloc::fmt;
 use core::any::TypeId;
 
 use crate::{
-    PyObject,
+    PyObject, PyObjectRef,
     object::{
         Erased, InstanceDict, MaybeTraverse, PyInner, PyObjectPayload, debug_obj, default_dealloc,
-        try_traverse_obj,
+        try_clear_obj, try_traverse_obj,
     },
 };
 
@@ -17,6 +17,9 @@ pub(in crate::object) struct PyObjVTable {
     pub(in crate::object) dealloc: unsafe fn(*mut PyObject),
     pub(in crate::object) debug: unsafe fn(&PyObject, &mut fmt::Formatter<'_>) -> fmt::Result,
     pub(in crate::object) trace: Option<unsafe fn(&PyObject, &mut TraverseFn<'_>)>,
+    /// Clear for circular reference resolution (tp_clear).
+    /// Called just before deallocation to extract child references.
+    pub(in crate::object) clear: Option<unsafe fn(*mut PyObject, &mut Vec<PyObjectRef>)>,
 }
 
 impl PyObjVTable {
@@ -32,6 +35,13 @@ impl PyObjVTable {
                     None
                 }
             },
+            clear: const {
+                if T::HAS_CLEAR {
+                    Some(try_clear_obj::<T>)
+                } else {
+                    None
+                }
+            },
         }
     }
 }
