@@ -8,7 +8,7 @@ use super::{ |
8 | 8 | use crate::{ |
9 | 9 | AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, |
10 | 10 | atomic_func, |
11 | | - class::{PyClassDef, PyClassImpl}, |
| 11 | + class::{PyClassDef, PyClassImpl, StaticType}, |
12 | 12 | common::{ |
13 | 13 | ascii, |
14 | 14 | hash::PyHash, |
@@ -22,6 +22,7 @@ use crate::{ |
22 | 22 | ArgIterable, FuncArgs, NameOthers, OptionalArg, PosArgs, PyArithmeticValue, |
23 | 23 | PyComparisonValue, |
24 | 24 | }, |
| 25 | + object::{Traverse, TraverseFn}, |
25 | 26 | protocol::{PyIterReturn, PyNumberMethods, PySequenceMethods}, |
26 | 27 | recursion::ReprGuard, |
27 | 28 | types::AsNumber, |
@@ -40,12 +41,23 @@ use rustpython_common::{ |
40 | 41 |
|
41 | 42 | pub(crate) type SetContentType = dict_inner::Dict<()>; |
42 | 43 |
|
43 | | -#[pyclass(module = false, name = "set", unhashable = true, traverse)] |
| 44 | +#[pyclass(module = false, name = "set", unhashable = true, traverse = "manual")] |
44 | 45 | #[derive(Default)] |
45 | 46 | pub struct PySet { |
46 | 47 | pub(super) inner: PySetInner, |
47 | 48 | } |
48 | 49 |
|
| 50 | +// SAFETY: the inner table visits each exclusively owned reference once. |
| 51 | +unsafe impl Traverse for PySet { |
| 52 | + fn traverse(&self, tracer: &mut TraverseFn<'_>) { |
| 53 | + self.inner.traverse(tracer); |
| 54 | + } |
| 55 | + |
| 56 | + fn clear(&mut self, out: &mut Vec<PyObjectRef>) { |
| 57 | + Traverse::clear(&mut self.inner, out); |
| 58 | + } |
| 59 | +} |
| 60 | + |
49 | 61 | impl PySet { |
50 | 62 | #[deprecated(note = "Use `PySet::default().into_ref(ctx)` instead")] |
51 | 63 | pub fn new_ref(ctx: &Context) -> PyRef<Self> { |
@@ -82,12 +94,28 @@ impl PySet { |
82 | 94 | } |
83 | 95 | } |
84 | 96 |
|
85 | | -#[pyclass(module = false, name = "frozenset", unhashable = true)] |
| 97 | +#[pyclass( |
| 98 | + module = false, |
| 99 | + name = "frozenset", |
| 100 | + unhashable = true, |
| 101 | + traverse = "manual" |
| 102 | +)] |
86 | 103 | pub struct PyFrozenSet { |
87 | 104 | inner: PySetInner, |
88 | 105 | hash: PyAtomic<PyHash>, |
89 | 106 | } |
90 | 107 |
|
| 108 | +// SAFETY: the inner table visits each exclusively owned reference once. |
| 109 | +unsafe impl Traverse for PyFrozenSet { |
| 110 | + fn traverse(&self, tracer: &mut TraverseFn<'_>) { |
| 111 | + self.inner.traverse(tracer); |
| 112 | + } |
| 113 | + |
| 114 | + fn clear(&mut self, out: &mut Vec<PyObjectRef>) { |
| 115 | + Traverse::clear(&mut self.inner, out); |
| 116 | + } |
| 117 | +} |
| 118 | + |
91 | 119 | impl Default for PyFrozenSet { |
92 | 120 | fn default() -> Self { |
93 | 121 | Self { |
@@ -161,13 +189,32 @@ impl fmt::Debug for PyFrozenSet { |
161 | 189 | } |
162 | 190 |
|
163 | 191 | impl PyPayload for PySet { |
| 192 | + // Ordinary destruction can release the table directly after untracking. |
| 193 | + const CLEAR_ON_DEALLOC: bool = false; |
| 194 | + |
164 | 195 | #[inline] |
165 | 196 | fn class(ctx: &Context) -> &'static Py<PyType> { |
166 | 197 | ctx.types.set_type |
167 | 198 | } |
168 | 199 | } |
169 | 200 |
|
170 | 201 | impl PyPayload for PyFrozenSet { |
| 202 | + const CLEAR_ON_DEALLOC: bool = false; |
| 203 | + |
| 204 | + fn needs_gc(&self, class: &Py<PyType>) -> bool { |
| 205 | + // Subclasses can acquire instance state. Shared working tables may |
| 206 | + // still be mutated by another owner, so keep those tracked as well. |
| 207 | + if !class.is(Self::static_type()) || PyRc::strong_count(&self.inner.content) != 1 { |
| 208 | + return true; |
| 209 | + } |
| 210 | + self.inner.content.any_key(|key| { |
| 211 | + // Exact untracked frozensets are immutable and were checked at |
| 212 | + // allocation. Other untracked objects may later acquire GC edges. |
| 213 | + key.gc_may_have_edges() |
| 214 | + && !(key.class().is(Self::static_type()) && !key.is_gc_tracked()) |
| 215 | + }) |
| 216 | + } |
| 217 | + |
171 | 218 | #[inline] |
172 | 219 | fn class(ctx: &Context) -> &'static Py<PyType> { |
173 | 220 | ctx.types.frozenset_type |
@@ -179,10 +226,22 @@ pub(super) struct PySetInner { |
179 | 226 | content: PyRc<SetContentType>, |
180 | 227 | } |
181 | 228 |
|
182 | | -unsafe impl crate::object::Traverse for PySetInner { |
183 | | - fn traverse(&self, tracer_fn: &mut crate::object::TraverseFn<'_>) { |
184 | | - // FIXME(discord9): Rc means shared ref, so should it be traced? |
185 | | - self.content.traverse(tracer_fn) |
| 229 | +// SAFETY: a shared table owns only one reference per entry, regardless of its |
| 230 | +// number of owners. Only attribute those edges to an exclusive owner; otherwise |
| 231 | +// leave them as external roots rather than subtracting the same edges twice. |
| 232 | +unsafe impl Traverse for PySetInner { |
| 233 | + fn traverse(&self, tracer: &mut TraverseFn<'_>) { |
| 234 | + if PyRc::strong_count(&self.content) == 1 { |
| 235 | + self.content.traverse(tracer); |
| 236 | + } |
| 237 | + } |
| 238 | + |
| 239 | + fn clear(&mut self, out: &mut Vec<PyObjectRef>) { |
| 240 | + if let Some(content) = PyRc::get_mut(&mut self.content) { |
| 241 | + // Detach before releasing keys: dropping them may re-enter Python. |
| 242 | + out.reserve(content.len()); |
| 243 | + out.extend(content.drain_entries().map(|(key, ())| key)); |
| 244 | + } |
186 | 245 | } |
187 | 246 | } |
188 | 247 |
|
@@ -1652,6 +1711,13 @@ pub(crate) struct PyAnySet { |
1652 | 1711 | object: PyObject, |
1653 | 1712 | } |
1654 | 1713 |
|
| 1714 | +// SAFETY: this owner contributes exactly one strong reference to the object. |
| 1715 | +unsafe impl Traverse for PyRef<PyAnySet> { |
| 1716 | + fn traverse(&self, tracer: &mut TraverseFn<'_>) { |
| 1717 | + tracer(self.as_object()); |
| 1718 | + } |
| 1719 | +} |
| 1720 | + |
1655 | 1721 | impl Borrow<PyObject> for PyAnySet { |
1656 | 1722 | #[inline(always)] |
1657 | 1723 | fn borrow(&self) -> &PyObject { |
@@ -1747,12 +1813,14 @@ impl TryFromObject for PyRef<PyAnySet> { |
1747 | 1813 | } |
1748 | 1814 | } |
1749 | 1815 |
|
1750 | | -#[pyclass(module = false, name = "set_iterator")] |
| 1816 | +#[pyclass(module = false, name = "set_iterator", traverse, clear = false)] |
1751 | 1817 | pub(crate) struct PySetIterator { |
| 1818 | + #[pytraverse(skip)] |
1752 | 1819 | size: DictSize, |
1753 | 1820 | /// Whether the set was found to have changed, which `setiter_iternext()` |
1754 | 1821 | /// records by writing a size no set can have. Sticky: what it makes the |
1755 | 1822 | /// iterator answer, it answers from then on. |
| 1823 | + #[pytraverse(skip)] |
1756 | 1824 | changed: PyAtomic<bool>, |
1757 | 1825 | internal: PyMutex<PositionIterInternal<PyRef<PyAnySet>>>, |
1758 | 1826 | } |
@@ -1907,6 +1975,115 @@ pub(crate) fn init(context: &'static Context) { |
1907 | 1975 | mod reference_tests { |
1908 | 1976 | use super::*; |
1909 | 1977 |
|
| 1978 | + #[test] |
| 1979 | + fn set_gc_detaches_keys_for_deferred_release() { |
| 1980 | + use crate::object::MaybeTraverse; |
| 1981 | + |
| 1982 | + fn check(mut set: impl MaybeTraverse, key: &PyObjectRef) { |
| 1983 | + let mut visits = 0; |
| 1984 | + set.try_traverse(&mut |child| { |
| 1985 | + assert!(core::ptr::eq(child, &**key)); |
| 1986 | + visits += 1; |
| 1987 | + }); |
| 1988 | + assert_eq!(visits, 1); |
| 1989 | + let mut edges = Vec::new(); |
| 1990 | + set.try_clear(&mut edges); |
| 1991 | + set.try_traverse(&mut |_| panic!("cleared set retained a key")); |
| 1992 | + set.try_clear(&mut edges); |
| 1993 | + assert_eq!(edges.len(), 1); |
| 1994 | + assert!(core::ptr::eq(&*edges[0], &**key)); |
| 1995 | + assert_eq!(key.strong_count(), 2); |
| 1996 | + drop(edges); |
| 1997 | + assert_eq!(key.strong_count(), 1); |
| 1998 | + } |
| 1999 | + |
| 2000 | + crate::Interpreter::without_stdlib(Default::default()).enter(|vm| { |
| 2001 | + for frozen in [false, true] { |
| 2002 | + let key: PyObjectRef = vm.ctx.new_int(1000).into(); |
| 2003 | + let inner = PySetInner::default(); |
| 2004 | + inner.add(&key, vm).unwrap(); |
| 2005 | + if frozen { |
| 2006 | + check( |
| 2007 | + PyFrozenSet { |
| 2008 | + inner, |
| 2009 | + ..Default::default() |
| 2010 | + }, |
| 2011 | + &key, |
| 2012 | + ); |
| 2013 | + } else { |
| 2014 | + check(PySet { inner }, &key); |
| 2015 | + } |
| 2016 | + } |
| 2017 | + }); |
| 2018 | + } |
| 2019 | + |
| 2020 | + #[test] |
| 2021 | + fn set_gc_preserves_shared_table_edges() { |
| 2022 | + crate::Interpreter::without_stdlib(Default::default()).enter(|vm| { |
| 2023 | + let key: PyObjectRef = vm.ctx.new_int(1000).into(); |
| 2024 | + let mut inner = PySetInner::default(); |
| 2025 | + inner.add(&key, vm).unwrap(); |
| 2026 | + let alias = inner.clone(); |
| 2027 | + inner.traverse(&mut |_| panic!("shared table edge counted")); |
| 2028 | + alias.traverse(&mut |_| panic!("shared table edge counted")); |
| 2029 | + let mut edges = Vec::new(); |
| 2030 | + Traverse::clear(&mut inner, &mut edges); |
| 2031 | + assert!(edges.is_empty()); |
| 2032 | + assert!(alias.contains(&key, vm).unwrap()); |
| 2033 | + drop(alias); |
| 2034 | + Traverse::clear(&mut inner, &mut edges); |
| 2035 | + assert_eq!(inner.len(), 0); |
| 2036 | + assert_eq!(edges.len(), 1); |
| 2037 | + }); |
| 2038 | + } |
| 2039 | + |
| 2040 | + #[test] |
| 2041 | + fn frozenset_gc_tracking_covers_native_construction() { |
| 2042 | + crate::Interpreter::without_stdlib(Default::default()).enter(|vm| { |
| 2043 | + let key: PyObjectRef = vm.ctx.new_int(1000).into(); |
| 2044 | + let atomic = PyRef::new_ref( |
| 2045 | + PyFrozenSet::from_iter(vm, [key.clone()]).unwrap(), |
| 2046 | + vm.ctx.types.frozenset_type.to_owned(), |
| 2047 | + None, |
| 2048 | + ); |
| 2049 | + assert!(!atomic.as_object().is_gc_tracked()); |
| 2050 | + let nested = PyFrozenSet::from_iter(vm, [atomic.into()]) |
| 2051 | + .unwrap() |
| 2052 | + .into_ref(&vm.ctx); |
| 2053 | + assert!(!nested.as_object().is_gc_tracked()); |
| 2054 | + |
| 2055 | + // A traversable key must keep the frozen owner visible to GC. |
| 2056 | + let tuple: PyObjectRef = vm.ctx.new_tuple(vec![key.clone()]).into(); |
| 2057 | + let owner = PyFrozenSet::from_iter(vm, [tuple]) |
| 2058 | + .unwrap() |
| 2059 | + .into_ref(&vm.ctx); |
| 2060 | + assert!(owner.as_object().is_gc_tracked()); |
| 2061 | + |
| 2062 | + let subclass = vm.ctx.new_class( |
| 2063 | + None, |
| 2064 | + "FrozenSubclass", |
| 2065 | + vm.ctx.types.frozenset_type.to_owned(), |
| 2066 | + Default::default(), |
| 2067 | + ); |
| 2068 | + let subclass = PyFrozenSet::default() |
| 2069 | + .into_ref_with_type(vm, subclass) |
| 2070 | + .unwrap(); |
| 2071 | + assert!(subclass.as_object().is_gc_tracked()); |
| 2072 | + |
| 2073 | + // A temporary native alias can still add traversable keys. |
| 2074 | + let inner = PySetInner::default(); |
| 2075 | + inner.add(&key, vm).unwrap(); |
| 2076 | + let alias = inner.clone(); |
| 2077 | + let owner = PyFrozenSet { |
| 2078 | + inner, |
| 2079 | + ..Default::default() |
| 2080 | + } |
| 2081 | + .into_ref(&vm.ctx); |
| 2082 | + assert!(owner.as_object().is_gc_tracked()); |
| 2083 | + drop(alias); |
| 2084 | + }); |
| 2085 | + } |
| 2086 | + |
1910 | 2087 | #[test] |
1911 | 2088 | fn anyset_reference_layout_and_ownership() { |
1912 | 2089 | static_assertions::assert_not_impl_any!(PyAnySet: PyPayload); |
|