Make ReprGuard vm-local

Author: coolreader18

vm/src/obj/objdict.rs

@@ -194,7 +194,7 @@ impl PyDictRef {
 
     #[pymethod(magic)]
     fn repr(self, vm: &VirtualMachine) -> PyResult<String> {
-        let s = if let Some(_guard) = ReprGuard::enter(self.as_object()) {
+        let s = if let Some(_guard) = ReprGuard::enter(vm, self.as_object()) {
             let mut str_parts = vec![];
             for (key, value) in self {
                 let key_repr = vm.to_repr(&key)?;
@@ -571,7 +571,7 @@ macro_rules! dict_iterator {
             #[pymethod(name = "__repr__")]
             #[allow(clippy::redundant_closure_call)]
             fn repr(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<String> {
-                let s = if let Some(_guard) = ReprGuard::enter(zelf.as_object()) {
+                let s = if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
                     let mut str_parts = vec![];
                     for (key, value) in zelf.dict.clone() {
                         let s = vm.to_repr(&$result_fn(vm, key, value))?;

vm/src/obj/objlist.rs

@@ -420,7 +420,7 @@ impl PyList {
 
     #[pymethod(name = "__repr__")]
     fn repr(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<String> {
-        let s = if let Some(_guard) = ReprGuard::enter(zelf.as_object()) {
+        let s = if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
             let elements = zelf.borrow_value().clone();
             let mut str_parts = Vec::with_capacity(elements.len());
             for elem in elements.iter() {

vm/src/obj/objset.rs

@@ -513,7 +513,7 @@ impl PySet {
     fn repr(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
         let s = if zelf.inner.len() == 0 {
             "set()".to_owned()
-        } else if let Some(_guard) = ReprGuard::enter(zelf.as_object()) {
+        } else if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
             zelf.inner.repr(vm)?
         } else {
             "set(...)".to_owned()
@@ -783,7 +783,7 @@ impl PyFrozenSet {
         let inner = &zelf.inner;
         let s = if inner.len() == 0 {
             "frozenset()".to_owned()
-        } else if let Some(_guard) = ReprGuard::enter(zelf.as_object()) {
+        } else if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
             format!("frozenset({})", inner.repr(vm)?)
         } else {
             "frozenset(...)".to_owned()

vm/src/obj/objtuple.rs

@@ -179,7 +179,7 @@ impl PyTuple {
 
     #[pymethod(name = "__repr__")]
     fn repr(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<String> {
-        let s = if let Some(_guard) = ReprGuard::enter(zelf.as_object()) {
+        let s = if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
             let mut str_parts = Vec::with_capacity(zelf.elements.len());
             for elem in zelf.elements.iter() {
                 let s = vm.to_repr(elem)?;

vm/src/pyobject.rs

@@ -1454,7 +1454,7 @@ pub trait PyStructSequenceImpl: PyClassImpl {
             Ok(format!("{}: {}", name, s))
         };
         let (body, suffix) =
-            if let Some(_guard) = rustpython_vm::vm::ReprGuard::enter(zelf.as_object()) {
+            if let Some(_guard) = rustpython_vm::vm::ReprGuard::enter(vm, zelf.as_object()) {
                 if Self::FIELD_NAMES.len() == 1 {
                     let value = zelf.borrow_value().first().unwrap();
                     let formatted = format_field((value, Self::FIELD_NAMES[0]))?;

vm/src/stdlib/collections.rs

@@ -248,7 +248,7 @@ mod _collections {
 
         #[pymethod(name = "__repr__")]
         fn repr(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<String> {
-            let repr = if let Some(_guard) = ReprGuard::enter(zelf.as_object()) {
+            let repr = if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
                 let elements = zelf
                     .borrow_deque()
                     .iter()

vm/src/vm.rs

@@ -63,6 +63,7 @@ pub struct VirtualMachine {
     pub use_tracing: Cell<bool>,
     pub recursion_limit: Cell<usize>,
     pub signal_handlers: Option<Box<RefCell<[PyObjectRef; NSIG]>>>,
+    pub repr_guards: RefCell<HashSet<usize>>,
     pub state: PyRc<PyGlobalState>,
     pub initialized: bool,
 }
@@ -199,6 +200,7 @@ impl VirtualMachine {
             use_tracing: Cell::new(false),
             recursion_limit: Cell::new(if cfg!(debug_assertions) { 256 } else { 512 }),
             signal_handlers: Some(Box::new(signal_handlers)),
+            repr_guards: RefCell::default(),
             state: PyRc::new(PyGlobalState {
                 settings,
                 stdlib_inits,
@@ -293,6 +295,7 @@ impl VirtualMachine {
             use_tracing: Cell::new(false),
             recursion_limit: self.recursion_limit.clone(),
             signal_handlers: None,
+            repr_guards: RefCell::default(),
             state: self.state.clone(),
             initialized: self.initialized,
         }
@@ -1489,37 +1492,32 @@ impl Default for VirtualMachine {
     }
 }
 
-thread_local! {
-    static REPR_GUARDS: RefCell<HashSet<usize>> = RefCell::default();
-}
-
-pub struct ReprGuard {
+pub struct ReprGuard<'vm> {
+    vm: &'vm VirtualMachine,
     id: usize,
 }
 
 /// A guard to protect repr methods from recursion into itself,
-impl ReprGuard {
+impl<'vm> ReprGuard<'vm> {
     /// Returns None if the guard against 'obj' is still held otherwise returns the guard. The guard
     /// which is released if dropped.
-    pub fn enter(obj: &PyObjectRef) -> Option<ReprGuard> {
-        REPR_GUARDS.with(|guards| {
-            let mut guards = guards.borrow_mut();
-
-            // Should this be a flag on the obj itself? putting it in a global variable for now until it
-            // decided the form of the PyObject. https://github.com/RustPython/RustPython/issues/371
-            let id = obj.get_id();
-            if guards.contains(&id) {
-                return None;
-            }
-            guards.insert(id);
-            Some(ReprGuard { id })
-        })
+    pub fn enter(vm: &'vm VirtualMachine, obj: &PyObjectRef) -> Option<Self> {
+        let mut guards = vm.repr_guards.borrow_mut();
+
+        // Should this be a flag on the obj itself? putting it in a global variable for now until it
+        // decided the form of the PyObject. https://github.com/RustPython/RustPython/issues/371
+        let id = obj.get_id();
+        if guards.contains(&id) {
+            return None;
+        }
+        guards.insert(id);
+        Some(ReprGuard { vm, id })
     }
 }
 
-impl Drop for ReprGuard {
+impl<'vm> Drop for ReprGuard<'vm> {
     fn drop(&mut self) {
-        REPR_GUARDS.with(|guards| guards.borrow_mut().remove(&self.id));
+        self.vm.repr_guards.borrow_mut().remove(&self.id);
     }
 }