@@ -8,7 +8,7 @@ use super::{
88use crate :: {
99 AsObject , Context , Py , PyObject , PyObjectRef , PyPayload , PyRef , PyResult , TryFromObject ,
1010 atomic_func,
11- class:: { PyClassDef , PyClassImpl } ,
11+ class:: { PyClassDef , PyClassImpl , StaticType } ,
1212 common:: {
1313 ascii,
1414 hash:: PyHash ,
@@ -22,6 +22,7 @@ use crate::{
2222 ArgIterable , FuncArgs , NameOthers , OptionalArg , PosArgs , PyArithmeticValue ,
2323 PyComparisonValue ,
2424 } ,
25+ object:: { Traverse , TraverseFn } ,
2526 protocol:: { PyIterReturn , PyNumberMethods , PySequenceMethods } ,
2627 recursion:: ReprGuard ,
2728 types:: AsNumber ,
@@ -40,12 +41,23 @@ use rustpython_common::{
4041
4142pub ( crate ) type SetContentType = dict_inner:: Dict < ( ) > ;
4243
43- #[ pyclass( module = false , name = "set" , unhashable = true , traverse) ]
44+ #[ pyclass( module = false , name = "set" , unhashable = true , traverse = "manual" ) ]
4445#[ derive( Default ) ]
4546pub struct PySet {
4647 pub ( super ) inner : PySetInner ,
4748}
4849
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+
4961impl PySet {
5062 #[ deprecated( note = "Use `PySet::default().into_ref(ctx)` instead" ) ]
5163 pub fn new_ref ( ctx : & Context ) -> PyRef < Self > {
@@ -82,12 +94,28 @@ impl PySet {
8294 }
8395}
8496
85- #[ pyclass( module = false , name = "frozenset" , unhashable = true ) ]
97+ #[ pyclass(
98+ module = false ,
99+ name = "frozenset" ,
100+ unhashable = true ,
101+ traverse = "manual"
102+ ) ]
86103pub struct PyFrozenSet {
87104 inner : PySetInner ,
88105 hash : PyAtomic < PyHash > ,
89106}
90107
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+
91119impl Default for PyFrozenSet {
92120 fn default ( ) -> Self {
93121 Self {
@@ -161,13 +189,32 @@ impl fmt::Debug for PyFrozenSet {
161189}
162190
163191impl PyPayload for PySet {
192+ // Ordinary destruction can release the table directly after untracking.
193+ const CLEAR_ON_DEALLOC : bool = false ;
194+
164195 #[ inline]
165196 fn class ( ctx : & Context ) -> & ' static Py < PyType > {
166197 ctx. types . set_type
167198 }
168199}
169200
170201impl 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+
171218 #[ inline]
172219 fn class ( ctx : & Context ) -> & ' static Py < PyType > {
173220 ctx. types . frozenset_type
@@ -179,10 +226,22 @@ pub(super) struct PySetInner {
179226 content : PyRc < SetContentType > ,
180227}
181228
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+ }
186245 }
187246}
188247
@@ -1652,6 +1711,13 @@ pub(crate) struct PyAnySet {
16521711 object : PyObject ,
16531712}
16541713
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+
16551721impl Borrow < PyObject > for PyAnySet {
16561722 #[ inline( always) ]
16571723 fn borrow ( & self ) -> & PyObject {
@@ -1747,12 +1813,14 @@ impl TryFromObject for PyRef<PyAnySet> {
17471813 }
17481814}
17491815
1750- #[ pyclass( module = false , name = "set_iterator" ) ]
1816+ #[ pyclass( module = false , name = "set_iterator" , traverse , clear = false ) ]
17511817pub ( crate ) struct PySetIterator {
1818+ #[ pytraverse( skip) ]
17521819 size : DictSize ,
17531820 /// Whether the set was found to have changed, which `setiter_iternext()`
17541821 /// records by writing a size no set can have. Sticky: what it makes the
17551822 /// iterator answer, it answers from then on.
1823+ #[ pytraverse( skip) ]
17561824 changed : PyAtomic < bool > ,
17571825 internal : PyMutex < PositionIterInternal < PyRef < PyAnySet > > > ,
17581826}
@@ -1907,6 +1975,115 @@ pub(crate) fn init(context: &'static Context) {
19071975mod reference_tests {
19081976 use super :: * ;
19091977
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+
19102087 #[ test]
19112088 fn anyset_reference_layout_and_ownership ( ) {
19122089 static_assertions:: assert_not_impl_any!( PyAnySet : PyPayload ) ;
0 commit comments