Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion Lib/test/test_set.py
Original file line number Diff line number Diff line change
Expand Up @@ -352,7 +352,6 @@ def test_do_not_rehash_dict_keys(self):
self.assertEqual(sum(elem.hash_count for elem in d), n)
self.assertEqual(d3, dict.fromkeys(d, 123))

@unittest.expectedFailure # TODO: RUSTPYTHON
def test_container_iterator(self):
# Bug #3680: tp_traverse was not implemented for set iterator object
class C(object):
Expand Down
193 changes: 185 additions & 8 deletions crates/vm/src/builtins/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use super::{
use crate::{
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
atomic_func,
class::{PyClassDef, PyClassImpl},
class::{PyClassDef, PyClassImpl, StaticType},
Comment thread
1ndahous3 marked this conversation as resolved.
common::{
ascii,
hash::PyHash,
Expand All @@ -22,6 +22,7 @@ use crate::{
ArgIterable, FuncArgs, NameOthers, OptionalArg, PosArgs, PyArithmeticValue,
PyComparisonValue,
},
object::{Traverse, TraverseFn},
protocol::{PyIterReturn, PyNumberMethods, PySequenceMethods},
recursion::ReprGuard,
types::AsNumber,
Expand All @@ -40,12 +41,23 @@ use rustpython_common::{

pub(crate) type SetContentType = dict_inner::Dict<()>;

#[pyclass(module = false, name = "set", unhashable = true, traverse)]
#[pyclass(module = false, name = "set", unhashable = true, traverse = "manual")]
#[derive(Default)]
pub struct PySet {
pub(super) inner: PySetInner,
}

// SAFETY: the inner table visits each exclusively owned reference once.
unsafe impl Traverse for PySet {
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
self.inner.traverse(tracer);
}

fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
Traverse::clear(&mut self.inner, out);
}
}

impl PySet {
#[deprecated(note = "Use `PySet::default().into_ref(ctx)` instead")]
pub fn new_ref(ctx: &Context) -> PyRef<Self> {
Expand Down Expand Up @@ -82,12 +94,28 @@ impl PySet {
}
}

#[pyclass(module = false, name = "frozenset", unhashable = true)]
#[pyclass(
module = false,
name = "frozenset",
unhashable = true,
traverse = "manual"
)]
pub struct PyFrozenSet {
inner: PySetInner,
hash: PyAtomic<PyHash>,
}

// SAFETY: the inner table visits each exclusively owned reference once.
unsafe impl Traverse for PyFrozenSet {
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
self.inner.traverse(tracer);
}

fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
Traverse::clear(&mut self.inner, out);
}
}

impl Default for PyFrozenSet {
fn default() -> Self {
Self {
Expand Down Expand Up @@ -161,13 +189,32 @@ impl fmt::Debug for PyFrozenSet {
}

impl PyPayload for PySet {
// Ordinary destruction can release the table directly after untracking.
const CLEAR_ON_DEALLOC: bool = false;

#[inline]
fn class(ctx: &Context) -> &'static Py<PyType> {
ctx.types.set_type
}
}

impl PyPayload for PyFrozenSet {
const CLEAR_ON_DEALLOC: bool = false;

fn needs_gc(&self, class: &Py<PyType>) -> bool {
// Subclasses can acquire instance state. Shared working tables may
// still be mutated by another owner, so keep those tracked as well.
if !class.is(Self::static_type()) || PyRc::strong_count(&self.inner.content) != 1 {
return true;
}
self.inner.content.any_key(|key| {
// Exact untracked frozensets are immutable and were checked at
// allocation. Other untracked objects may later acquire GC edges.
key.gc_may_have_edges()
&& !(key.class().is(Self::static_type()) && !key.is_gc_tracked())
})
}

#[inline]
fn class(ctx: &Context) -> &'static Py<PyType> {
ctx.types.frozenset_type
Expand All @@ -179,10 +226,22 @@ pub(super) struct PySetInner {
content: PyRc<SetContentType>,
}

unsafe impl crate::object::Traverse for PySetInner {
fn traverse(&self, tracer_fn: &mut crate::object::TraverseFn<'_>) {
// FIXME(discord9): Rc means shared ref, so should it be traced?
self.content.traverse(tracer_fn)
// SAFETY: a shared table owns only one reference per entry, regardless of its
// number of owners. Only attribute those edges to an exclusive owner; otherwise
// leave them as external roots rather than subtracting the same edges twice.
unsafe impl Traverse for PySetInner {
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
if PyRc::strong_count(&self.content) == 1 {
self.content.traverse(tracer);
}
}

fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
if let Some(content) = PyRc::get_mut(&mut self.content) {
// Detach before releasing keys: dropping them may re-enter Python.
out.reserve(content.len());
out.extend(content.drain_entries().map(|(key, ())| key));
}
}
}

Expand Down Expand Up @@ -1652,6 +1711,13 @@ pub(crate) struct PyAnySet {
object: PyObject,
}

// SAFETY: this owner contributes exactly one strong reference to the object.
unsafe impl Traverse for PyRef<PyAnySet> {
fn traverse(&self, tracer: &mut TraverseFn<'_>) {
tracer(self.as_object());
}
}

impl Borrow<PyObject> for PyAnySet {
#[inline(always)]
fn borrow(&self) -> &PyObject {
Expand Down Expand Up @@ -1747,12 +1813,14 @@ impl TryFromObject for PyRef<PyAnySet> {
}
}

#[pyclass(module = false, name = "set_iterator")]
#[pyclass(module = false, name = "set_iterator", traverse, clear = false)]
pub(crate) struct PySetIterator {
#[pytraverse(skip)]
size: DictSize,
/// Whether the set was found to have changed, which `setiter_iternext()`
/// records by writing a size no set can have. Sticky: what it makes the
/// iterator answer, it answers from then on.
#[pytraverse(skip)]
changed: PyAtomic<bool>,
internal: PyMutex<PositionIterInternal<PyRef<PyAnySet>>>,
}
Expand Down Expand Up @@ -1907,6 +1975,115 @@ pub(crate) fn init(context: &'static Context) {
mod reference_tests {
use super::*;

#[test]
fn set_gc_detaches_keys_for_deferred_release() {
use crate::object::MaybeTraverse;

fn check(mut set: impl MaybeTraverse, key: &PyObjectRef) {
let mut visits = 0;
set.try_traverse(&mut |child| {
assert!(core::ptr::eq(child, &**key));
visits += 1;
});
assert_eq!(visits, 1);
let mut edges = Vec::new();
set.try_clear(&mut edges);
set.try_traverse(&mut |_| panic!("cleared set retained a key"));
set.try_clear(&mut edges);
assert_eq!(edges.len(), 1);
assert!(core::ptr::eq(&*edges[0], &**key));
assert_eq!(key.strong_count(), 2);
drop(edges);
assert_eq!(key.strong_count(), 1);
}

crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
for frozen in [false, true] {
let key: PyObjectRef = vm.ctx.new_int(1000).into();
let inner = PySetInner::default();
inner.add(&key, vm).unwrap();
if frozen {
check(
PyFrozenSet {
inner,
..Default::default()
},
&key,
);
} else {
check(PySet { inner }, &key);
}
}
});
}

#[test]
fn set_gc_preserves_shared_table_edges() {
crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
let key: PyObjectRef = vm.ctx.new_int(1000).into();
let mut inner = PySetInner::default();
inner.add(&key, vm).unwrap();
let alias = inner.clone();
inner.traverse(&mut |_| panic!("shared table edge counted"));
alias.traverse(&mut |_| panic!("shared table edge counted"));
let mut edges = Vec::new();
Traverse::clear(&mut inner, &mut edges);
assert!(edges.is_empty());
assert!(alias.contains(&key, vm).unwrap());
drop(alias);
Traverse::clear(&mut inner, &mut edges);
assert_eq!(inner.len(), 0);
assert_eq!(edges.len(), 1);
});
}

#[test]
fn frozenset_gc_tracking_covers_native_construction() {
crate::Interpreter::without_stdlib(Default::default()).enter(|vm| {
let key: PyObjectRef = vm.ctx.new_int(1000).into();
let atomic = PyRef::new_ref(
PyFrozenSet::from_iter(vm, [key.clone()]).unwrap(),
vm.ctx.types.frozenset_type.to_owned(),
None,
);
assert!(!atomic.as_object().is_gc_tracked());
let nested = PyFrozenSet::from_iter(vm, [atomic.into()])
.unwrap()
.into_ref(&vm.ctx);
assert!(!nested.as_object().is_gc_tracked());

// A traversable key must keep the frozen owner visible to GC.
let tuple: PyObjectRef = vm.ctx.new_tuple(vec![key.clone()]).into();
let owner = PyFrozenSet::from_iter(vm, [tuple])
.unwrap()
.into_ref(&vm.ctx);
assert!(owner.as_object().is_gc_tracked());

let subclass = vm.ctx.new_class(
None,
"FrozenSubclass",
vm.ctx.types.frozenset_type.to_owned(),
Default::default(),
);
let subclass = PyFrozenSet::default()
.into_ref_with_type(vm, subclass)
.unwrap();
assert!(subclass.as_object().is_gc_tracked());

// A temporary native alias can still add traversable keys.
let inner = PySetInner::default();
inner.add(&key, vm).unwrap();
let alias = inner.clone();
let owner = PyFrozenSet {
inner,
..Default::default()
}
.into_ref(&vm.ctx);
assert!(owner.as_object().is_gc_tracked());
drop(alias);
});
}

#[test]
fn anyset_reference_layout_and_ownership() {
static_assertions::assert_not_impl_any!(PyAnySet: PyPayload);
Expand Down
10 changes: 10 additions & 0 deletions crates/vm/src/dict_inner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1165,6 +1165,16 @@ impl<T: Clone> Dict<T> {
current != *old
}

/// Inspect keys without allocating or cloning their references.
/// The predicate must not run Python or mutate this table.
pub(crate) fn any_key(&self, mut predicate: impl FnMut(&PyObject) -> bool) -> bool {
self.read()
.entries
.iter()
.flatten()
.any(|entry| predicate(&entry.key))
}

pub(crate) fn keys(&self) -> Vec<PyObjectRef> {
self.read()
.entries
Expand Down
21 changes: 16 additions & 5 deletions crates/vm/src/object/core.rs
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,9 @@ pub(super) unsafe fn default_dealloc<T: PyPayload>(obj: *mut PyObject) {
// them. Some payloads (e.g. FrameObject) drop children in place inside clear_fn
// instead of extracting them, so user code (`__del__`) may run here.
let mut edges = Vec::new();
if let Some(clear_fn) = vtable.clear {
if (T::CLEAR_ON_DEALLOC || T::HAS_FREELIST)
&& let Some(clear_fn) = vtable.clear
{
unsafe { clear_fn(obj, &mut edges) };
}
// Drop extracted child references - may trigger recursive destruction.
Expand Down Expand Up @@ -2290,6 +2292,15 @@ impl PyObject {
GcBits::from_bits_retain(self.0.gc_bits.load(Ordering::Relaxed)).contains(GcBits::TRACKED)
}

/// Whether this object's payload or instance storage can own GC edges.
/// Unlike is_gc_tracked, this also covers temporarily untracked objects.
#[inline]
pub(crate) fn gc_may_have_edges(&self) -> bool {
self.0.vtable.trace.is_some()
|| self.class().heaptype_ext.is_some()
|| self.0.ext_ref().is_some()
}

/// Get the referents (objects directly referenced) of this object.
/// Uses the full traverse including dict and slots.
pub fn gc_get_referents(&self) -> Vec<PyObjectRef> {
Expand Down Expand Up @@ -2937,6 +2948,8 @@ impl<T: PyPayload + crate::object::MaybeTraverse + core::fmt::Debug> PyRef<T> {
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 track = !T::NEW_REF_UNTRACKED
&& (has_dict || is_heaptype || (T::HAS_TRAVERSE && payload.needs_gc(&typ)));

// Try to reuse from freelist (no dict, no heaptype)
let cached = if !has_dict && !is_heaptype {
Expand Down Expand Up @@ -2969,14 +2982,12 @@ impl<T: PyPayload + crate::object::MaybeTraverse + core::fmt::Debug> PyRef<T> {
};

// Track object if:
// - HAS_TRAVERSE is true (Rust payload implements Traverse), OR
// - HAS_TRAVERSE is true and the payload needs cyclic GC, OR
// - has instance dict (user-defined class instances), OR
// - heap type (all heap type instances are GC-tracked, like Py_TPFLAGS_HAVE_GC)
// unless the payload opts out via NEW_REF_UNTRACKED (e.g. call frames,
// which are tracked lazily only on escape).
if (<T as crate::object::MaybeTraverse>::HAS_TRAVERSE || has_dict || is_heaptype)
&& !T::NEW_REF_UNTRACKED
{
if track {
// Tracks under the interpreter running now and collects if this
// allocation pushed gen0 past its threshold.
unsafe {
Expand Down
15 changes: 15 additions & 0 deletions crates/vm/src/object/payload.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,21 @@ pub trait PyPayload: MaybeTraverse + PyThreadingConstraint + Sized + 'static {
/// which is created untracked and tracked lazily only on escape.
const NEW_REF_UNTRACKED: bool = false;

/// Whether this payload needs cyclic GC when first wrapped in an object.
/// Returning false promises that this instance cannot participate in a
/// reference cycle for its entire lifetime, not just at allocation time.
/// The check must not run Python. Instance dictionaries and heap types are
/// still tracked, unless NEW_REF_UNTRACKED opts out explicitly.
#[inline]
fn needs_gc(&self, _class: &Py<PyType>) -> bool {
true
}

/// Run tp_clear before ordinary reference-counted destruction. Payloads
/// whose Rust destructor releases their children directly may opt out.
/// Cyclic GC still calls tp_clear; freelists also require a cleared payload.
const CLEAR_ON_DEALLOC: bool = true;

/// Whether this type has a freelist. Types with freelists require
/// immediate (non-deferred) GC untracking during dealloc to prevent
/// race conditions when the object is reused.
Expand Down
Loading