From 391392819bf716573a5666a7695c9d027be66bc0 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Fri, 6 Mar 2026 06:18:07 +0900 Subject: [PATCH] Separate WeakRefList from ObjExt as independent prefix - Remove weak_list from ObjExt, allocate WeakRefList as its own prefix slot before PyInner - Add MANAGED_WEAKREF flag (1 << 3) to PyTypeFlags - Normalize MANAGED_WEAKREF from HAS_WEAKREF after flag assembly - Use Layout::extend offsets in bootstrap allocator --- crates/vm/src/builtins/type.rs | 18 +- crates/vm/src/object/core.rs | 215 +++++++++++++++++------- crates/vm/src/object/traverse_object.rs | 1 - crates/vm/src/types/slot.rs | 1 + 4 files changed, 170 insertions(+), 65 deletions(-) diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs index 9fa92903f37..f383441e4ef 100644 --- a/crates/vm/src/builtins/type.rs +++ b/crates/vm/src/builtins/type.rs @@ -585,12 +585,12 @@ impl PyType { slots.flags |= PyTypeFlags::HAS_DICT } - // Inherit HAS_WEAKREF from any base in MRO that has it + // Inherit HAS_WEAKREF/MANAGED_WEAKREF from any base in MRO that has it if mro .iter() .any(|b| b.slots.flags.has_feature(PyTypeFlags::HAS_WEAKREF)) { - slots.flags |= PyTypeFlags::HAS_WEAKREF + slots.flags |= PyTypeFlags::HAS_WEAKREF | PyTypeFlags::MANAGED_WEAKREF } // Inherit SEQUENCE and MAPPING flags from base classes @@ -605,6 +605,11 @@ impl PyType { Self::inherit_readonly_slots(&mut slots, &base); + // Normalize: any type with HAS_WEAKREF gets MANAGED_WEAKREF + if slots.flags.has_feature(PyTypeFlags::HAS_WEAKREF) { + slots.flags |= PyTypeFlags::MANAGED_WEAKREF; + } + if let Some(qualname) = attrs.get(identifier!(ctx, __qualname__)) && !qualname.fast_isinstance(ctx.types.str_type) { @@ -655,7 +660,7 @@ impl PyType { slots.flags |= PyTypeFlags::HAS_DICT } if base.slots.flags.has_feature(PyTypeFlags::HAS_WEAKREF) { - slots.flags |= PyTypeFlags::HAS_WEAKREF + slots.flags |= PyTypeFlags::HAS_WEAKREF | PyTypeFlags::MANAGED_WEAKREF } // Inherit SEQUENCE and MAPPING flags from base class @@ -668,6 +673,11 @@ impl PyType { Self::inherit_readonly_slots(&mut slots, &base); + // Normalize: any type with HAS_WEAKREF gets MANAGED_WEAKREF + if slots.flags.has_feature(PyTypeFlags::HAS_WEAKREF) { + slots.flags |= PyTypeFlags::MANAGED_WEAKREF; + } + let bases = PyRwLock::new(vec![base.clone()]); let mro = base.mro_map_collect(|x| x.to_owned()); @@ -1976,7 +1986,7 @@ impl Constructor for PyType { // 2. __weakref__ is in __slots__ let may_add_weakref = !base.slots.flags.has_feature(PyTypeFlags::HAS_WEAKREF); if (heaptype_slots.is_none() && may_add_weakref) || add_weakref { - flags |= PyTypeFlags::HAS_WEAKREF; + flags |= PyTypeFlags::HAS_WEAKREF | PyTypeFlags::MANAGED_WEAKREF; } let (slots, heaptype_ext) = { diff --git a/crates/vm/src/object/core.rs b/crates/vm/src/object/core.rs index a7e5b519f9a..e31470fd12e 100644 --- a/crates/vm/src/object/core.rs +++ b/crates/vm/src/object/core.rs @@ -287,7 +287,7 @@ unsafe impl Link for GcLink { } } -/// Extension fields for objects that need dict, weakref list, or member slots. +/// Extension fields for objects that need dict or member slots. /// Allocated as a prefix before PyInner when needed (prefix allocation pattern). /// Access via `PyInner::ext_ref()` using negative offset from the object pointer. /// @@ -298,7 +298,6 @@ unsafe impl Link for GcLink { #[repr(C, align(8))] pub(super) struct ObjExt { pub(super) dict: Option, - pub(super) weak_list: WeakRefList, pub(super) slots: Box<[PyRwLock>]>, } @@ -306,7 +305,6 @@ impl ObjExt { fn new(dict: Option, member_count: usize) -> Self { Self { dict: dict.map(InstanceDict::new), - weak_list: WeakRefList::new(), slots: core::iter::repeat_with(|| PyRwLock::new(None)) .take(member_count) .collect_vec() @@ -321,16 +319,20 @@ impl fmt::Debug for ObjExt { } } -/// Precomputed offset from PyInner pointer back to ObjExt prefix. -/// ObjExt is #[repr(C, align(8))] and PyInner is #[repr(C)], so as long as -/// ObjExt's alignment >= PyInner's alignment, Layout::extend adds no padding -/// and the offset equals size_of::(). +/// Precomputed offset constants for prefix allocation. +/// All prefix components are align(8) and their sizes are multiples of 8, +/// so Layout::extend adds no inter-padding. const EXT_OFFSET: usize = core::mem::size_of::(); -// Guarantee: ObjExt size is a multiple of its alignment, and its alignment -// is >= any PyInner alignment, so Layout::extend produces no inter-padding. +const WEAKREF_OFFSET: usize = core::mem::size_of::(); + const _: () = assert!(core::mem::size_of::().is_multiple_of(core::mem::align_of::())); const _: () = assert!(core::mem::align_of::() >= core::mem::align_of::>()); +const _: () = assert!( + core::mem::size_of::().is_multiple_of(core::mem::align_of::()) +); +const _: () = + assert!(core::mem::align_of::() >= core::mem::align_of::>()); /// This is an actual python object. It consists of a `typ` which is the /// python class, and carries some rust payload optionally. This rust @@ -354,38 +356,59 @@ pub(super) struct PyInner { pub(crate) const SIZEOF_PYOBJECT_HEAD: usize = core::mem::size_of::>(); impl PyInner { - /// Check if this object has an ObjExt prefix based on type flags. - /// Uses raw pointer reads to avoid Stacked Borrows violations during bootstrap, - /// where type objects have self-referential typ pointers that may be mutated. + /// Read type flags and member_count via raw pointers to avoid Stacked Borrows + /// violations during bootstrap, where type objects have self-referential typ pointers. #[inline(always)] - fn has_ext(&self) -> bool { - // Read slots via raw pointers only — creating a &Py reference - // would retag the entire object, conflicting with &mut writes during bootstrap. + fn read_type_flags(&self) -> (crate::types::PyTypeFlags, usize) { let typ_ptr = self.typ.load_raw(); let slots = unsafe { core::ptr::addr_of!((*typ_ptr).0.payload.slots) }; let flags = unsafe { core::ptr::addr_of!((*slots).flags).read() }; let member_count = unsafe { core::ptr::addr_of!((*slots).member_count).read() }; - flags.has_feature(crate::types::PyTypeFlags::HAS_DICT) - || flags.has_feature(crate::types::PyTypeFlags::HAS_WEAKREF) - || member_count > 0 + (flags, member_count) } /// Access the ObjExt prefix at a negative offset from this PyInner. - /// Returns None if this object was allocated without the prefix. + /// Returns None if this object was allocated without dict/slots. /// - /// Uses exposed provenance to reconstruct a pointer covering the entire - /// allocation (ObjExt prefix + PyInner). The allocation pointer's provenance - /// is exposed at allocation time via `expose_provenance()`. + /// Layout: [ObjExt?][WeakRefList?][PyInner] + /// ObjExt offset depends on whether WeakRefList is also present. #[inline(always)] pub(super) fn ext_ref(&self) -> Option<&ObjExt> { - if !self.has_ext() { + let (flags, member_count) = self.read_type_flags(); + let has_ext = + flags.has_feature(crate::types::PyTypeFlags::HAS_DICT) || member_count > 0; + if !has_ext { return None; } + let has_weakref = flags.has_feature(crate::types::PyTypeFlags::HAS_WEAKREF); + let offset = if has_weakref { + WEAKREF_OFFSET + EXT_OFFSET + } else { + EXT_OFFSET + }; let self_addr = (self as *const Self as *const u8).addr(); let ext_ptr = - core::ptr::with_exposed_provenance::(self_addr.wrapping_sub(EXT_OFFSET)); + core::ptr::with_exposed_provenance::(self_addr.wrapping_sub(offset)); Some(unsafe { &*ext_ptr }) } + + /// Access the WeakRefList prefix at a fixed negative offset from this PyInner. + /// Returns None if the type does not support weakrefs. + /// + /// Layout: [ObjExt?][WeakRefList?][PyInner] + /// WeakRefList is always immediately before PyInner (fixed WEAKREF_OFFSET). + #[inline(always)] + pub(super) fn weakref_list_ref(&self) -> Option<&WeakRefList> { + let (flags, _) = self.read_type_flags(); + if !flags.has_feature(crate::types::PyTypeFlags::HAS_WEAKREF) { + return None; + } + let self_addr = (self as *const Self as *const u8).addr(); + let ptr = core::ptr::with_exposed_provenance::( + self_addr.wrapping_sub(WEAKREF_OFFSET), + ); + Some(unsafe { &*ptr }) + } } impl fmt::Debug for PyInner { @@ -475,6 +498,7 @@ pub(crate) fn reset_weakref_locks_after_fork() { // === WeakRefList: inline on every object (tp_weaklist) === +#[repr(C)] pub(super) struct WeakRefList { /// Head of the intrusive doubly-linked list of weakrefs. head: PyAtomic<*mut Py>, @@ -842,8 +866,8 @@ impl PyWeak { } let obj = unsafe { &*obj_ptr }; - // Safety: if a weakref exists pointing to this object, ext must be present - let wrl = &obj.0.ext_ref().unwrap().weak_list; + // Safety: if a weakref exists pointing to this object, weakref prefix must be present + let wrl = obj.0.weakref_list_ref().unwrap(); // Compute our Py node pointer from payload address let offset = std::mem::offset_of!(PyInner, payload); @@ -925,22 +949,49 @@ impl InstanceDict { } impl PyInner { - /// Deallocate a PyInner, handling the optional ObjExt prefix. + /// Deallocate a PyInner, handling optional prefix(es). + /// Layout: [ObjExt?][WeakRefList?][PyInner] /// /// # Safety /// `ptr` must be a valid pointer from `PyInner::new` and must not be used after this call. unsafe fn dealloc(ptr: *mut Self) { unsafe { - if (*ptr).has_ext() { - let ext_layout = core::alloc::Layout::new::(); - let inner_layout = core::alloc::Layout::new::(); - let (combined, inner_offset) = ext_layout.extend(inner_layout).unwrap(); + let (flags, member_count) = (*ptr).read_type_flags(); + let has_ext = flags.has_feature(crate::types::PyTypeFlags::HAS_DICT) + || member_count > 0; + let has_weakref = flags.has_feature(crate::types::PyTypeFlags::HAS_WEAKREF); + + if has_ext || has_weakref { + // Reconstruct the same layout used in new() + let mut layout = core::alloc::Layout::from_size_align(0, 1).unwrap(); + + if has_ext { + layout = layout + .extend(core::alloc::Layout::new::()) + .unwrap() + .0; + } + if has_weakref { + layout = layout + .extend(core::alloc::Layout::new::()) + .unwrap() + .0; + } + let (combined, inner_offset) = layout + .extend(core::alloc::Layout::new::()) + .unwrap(); let combined = combined.pad_to_align(); let alloc_ptr = (ptr as *mut u8).sub(inner_offset); + // Drop PyInner (payload, typ, etc.) core::ptr::drop_in_place(ptr); - core::ptr::drop_in_place(alloc_ptr as *mut ObjExt); + + // Drop ObjExt if present (dict, slots) + if has_ext { + core::ptr::drop_in_place(alloc_ptr as *mut ObjExt); + } + // WeakRefList has no Drop (just raw pointers), no drop_in_place needed alloc::alloc::dealloc(alloc_ptr, combined); } else { @@ -951,42 +1002,71 @@ impl PyInner { } impl PyInner { - /// Allocate a new PyInner, optionally with an ObjExt prefix. + /// Allocate a new PyInner, optionally with prefix(es). /// Returns a raw pointer to the PyInner (NOT the allocation start). - /// For objects with ext, the allocation layout is: [ObjExt][PyInner] + /// Layout: [ObjExt?][WeakRefList?][PyInner] fn new(payload: T, typ: PyTypeRef, dict: Option) -> *mut Self { let member_count = typ.slots.member_count; let needs_ext = typ .slots .flags .has_feature(crate::types::PyTypeFlags::HAS_DICT) - || typ - .slots - .flags - .has_feature(crate::types::PyTypeFlags::HAS_WEAKREF) || member_count > 0; + let needs_weakref = typ + .slots + .flags + .has_feature(crate::types::PyTypeFlags::HAS_WEAKREF); debug_assert!( needs_ext || dict.is_none(), "dict passed to type '{}' without HAS_DICT flag", typ.name() ); - if needs_ext { - let ext_layout = core::alloc::Layout::new::(); - let inner_layout = core::alloc::Layout::new::(); - let (combined, inner_offset) = ext_layout.extend(inner_layout).unwrap(); + if needs_ext || needs_weakref { + // Build layout left-to-right: [ObjExt?][WeakRefList?][PyInner] + let mut layout = core::alloc::Layout::from_size_align(0, 1).unwrap(); + + let ext_start = if needs_ext { + let (combined, offset) = + layout.extend(core::alloc::Layout::new::()).unwrap(); + layout = combined; + Some(offset) + } else { + None + }; + + let weakref_start = if needs_weakref { + let (combined, offset) = layout + .extend(core::alloc::Layout::new::()) + .unwrap(); + layout = combined; + Some(offset) + } else { + None + }; + + let (combined, inner_offset) = layout + .extend(core::alloc::Layout::new::()) + .unwrap(); let combined = combined.pad_to_align(); let alloc_ptr = unsafe { alloc::alloc::alloc(combined) }; if alloc_ptr.is_null() { alloc::alloc::handle_alloc_error(combined); } - // Expose provenance so ext_ref() can reconstruct via with_exposed_provenance + // Expose provenance so ext_ref()/weakref_list_ref() can reconstruct alloc_ptr.expose_provenance(); unsafe { - let ext_ptr = alloc_ptr as *mut ObjExt; - ext_ptr.write(ObjExt::new(dict, member_count)); + if let Some(offset) = ext_start { + let ext_ptr = alloc_ptr.add(offset) as *mut ObjExt; + ext_ptr.write(ObjExt::new(dict, member_count)); + } + + if let Some(offset) = weakref_start { + let weakref_ptr = alloc_ptr.add(offset) as *mut WeakRefList; + weakref_ptr.write(WeakRefList::new()); + } let inner_ptr = alloc_ptr.add(inner_offset) as *mut Self; inner_ptr.write(Self { @@ -1229,13 +1309,12 @@ impl PyObjectRef { } impl PyObject { - /// Returns the WeakRefList if the object has ext (prefix allocation). - /// Note: This does NOT check HAS_WEAKREF flag. Callers creating weakrefs - /// must check HAS_WEAKREF themselves. This method is used in drop/GC paths - /// where the type may already be deallocated. + /// Returns the WeakRefList if the type supports weakrefs (HAS_WEAKREF). + /// The WeakRefList is stored as a separate prefix before PyInner, + /// independent from ObjExt (dict/slots). #[inline(always)] fn weak_ref_list(&self) -> Option<&WeakRefList> { - self.0.ext_ref().map(|ext| &ext.weak_list) + self.0.weakref_list_ref() } /// Returns the first weakref in the weakref list, if any. @@ -1677,10 +1756,19 @@ impl PyObject { // Detach the dict via Py_CLEAR(*_PyObject_GetDictPtr(self)) — NULL // the pointer without clearing dict contents. The dict may still be // referenced by other live objects (e.g. function.__globals__). - if obj.0.has_ext() { + let (flags, member_count) = obj.0.read_type_flags(); + let has_ext = flags.has_feature(crate::types::PyTypeFlags::HAS_DICT) + || member_count > 0; + if has_ext { + let has_weakref = flags.has_feature(crate::types::PyTypeFlags::HAS_WEAKREF); + let offset = if has_weakref { + WEAKREF_OFFSET + EXT_OFFSET + } else { + EXT_OFFSET + }; let self_addr = (ptr as *const u8).addr(); let ext_ptr = core::ptr::with_exposed_provenance_mut::( - self_addr.wrapping_sub(EXT_OFFSET), + self_addr.wrapping_sub(offset), ); let ext = unsafe { &mut *ext_ptr }; if let Some(old_dict) = ext.dict.take() { @@ -2328,29 +2416,36 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) { heaptype_ext: None, tp_version_tag: core::sync::atomic::AtomicU32::new(0), }; - // Both type_type and object_type are instances of `type`, which has HAS_WEAKREF, - // so they need prefix allocation with ObjExt. - let alloc_type_with_ext = || -> *mut MaybeUninit> { - let ext_layout = core::alloc::Layout::new::(); + // Both type_type and object_type are instances of `type`, which has + // HAS_DICT and HAS_WEAKREF, so they need both ObjExt and WeakRefList prefixes. + // Layout: [ObjExt][WeakRefList][PyInner] + let alloc_type_with_prefixes = || -> *mut MaybeUninit> { let inner_layout = core::alloc::Layout::new::>>(); - let (combined, inner_offset) = ext_layout.extend(inner_layout).unwrap(); + let ext_layout = core::alloc::Layout::new::(); + let weakref_layout = core::alloc::Layout::new::(); + + let (layout, weakref_offset) = ext_layout.extend(weakref_layout).unwrap(); + let (combined, inner_offset) = layout.extend(inner_layout).unwrap(); let combined = combined.pad_to_align(); let alloc_ptr = unsafe { alloc::alloc::alloc(combined) }; if alloc_ptr.is_null() { alloc::alloc::handle_alloc_error(combined); } - // Expose provenance so ext_ref() can reconstruct via with_exposed_provenance alloc_ptr.expose_provenance(); unsafe { let ext_ptr = alloc_ptr as *mut ObjExt; ext_ptr.write(ObjExt::new(None, 0)); + + let weakref_ptr = alloc_ptr.add(weakref_offset) as *mut WeakRefList; + weakref_ptr.write(WeakRefList::new()); + alloc_ptr.add(inner_offset) as *mut MaybeUninit> } }; - let type_type_ptr = alloc_type_with_ext(); + let type_type_ptr = alloc_type_with_prefixes(); unsafe { type_type_ptr.write(partially_init!( PyInner:: { @@ -2365,7 +2460,7 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) { )); } - let object_type_ptr = alloc_type_with_ext(); + let object_type_ptr = alloc_type_with_prefixes(); unsafe { object_type_ptr.write(partially_init!( PyInner:: { diff --git a/crates/vm/src/object/traverse_object.rs b/crates/vm/src/object/traverse_object.rs index de8d2d5f53e..f5614b3502a 100644 --- a/crates/vm/src/object/traverse_object.rs +++ b/crates/vm/src/object/traverse_object.rs @@ -68,7 +68,6 @@ unsafe impl Traverse for PyInner { // Traverse ObjExt prefix fields (dict and slots) if present if let Some(ext) = self.ext_ref() { ext.dict.traverse(tracer_fn); - // weak_list is atomic pointers, no trace needed ext.slots.traverse(tracer_fn); } diff --git a/crates/vm/src/types/slot.rs b/crates/vm/src/types/slot.rs index bd390af9546..af404d5c956 100644 --- a/crates/vm/src/types/slot.rs +++ b/crates/vm/src/types/slot.rs @@ -216,6 +216,7 @@ bitflags! { #[derive(Copy, Clone, Debug, PartialEq)] #[non_exhaustive] pub struct PyTypeFlags: u64 { + const MANAGED_WEAKREF = 1 << 3; const MANAGED_DICT = 1 << 4; const SEQUENCE = 1 << 5; const MAPPING = 1 << 6;