Use the thread_local crate for local dicts

coolreader18 · coolreader18 · commit fa2f5d640a60 · 2020-06-02T15:32:01.000-05:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -72,6 +72,7 @@ bstr = "0.2.12"
 crossbeam-utils = "0.7"
 generational-arena = "0.2"
 parking_lot = { git = "https://github.com/Amanieu/parking_lot" } # TODO: use published version
+thread_local = "1.0"
 
 ## unicode stuff
 unicode_names2 = "0.4"
diff --git a/vm/src/stdlib/thread.rs b/vm/src/stdlib/thread.rs
@@ -11,13 +11,13 @@ use crate::pyobject::{
 };
 use crate::vm::VirtualMachine;
 
-use crossbeam_utils::atomic::AtomicCell;
 use parking_lot::{
     lock_api::{RawMutex as RawMutexT, RawMutexTimed, RawReentrantMutex},
     RawMutex, RawThreadId,
 };
+use thread_local::ThreadLocal;
+
 use std::cell::RefCell;
-use std::collections::HashMap;
 use std::io::Write;
 use std::time::Duration;
 use std::{fmt, thread};
@@ -267,13 +267,10 @@ fn thread_count(vm: &VirtualMachine) -> usize {
     vm.state.thread_count.load()
 }
 
-static LOCAL_DATA_KEY: AtomicCell<usize> = AtomicCell::new(0);
-thread_local!(static LOCAL_DATA: RefCell<HashMap<usize, PyDictRef>> = RefCell::default());
-
 #[pyclass(name = "_local")]
 #[derive(Debug)]
 struct PyLocal {
-    key: usize,
+    data: ThreadLocal<PyDictRef>,
 }
 
 impl PyValue for PyLocal {
@@ -282,27 +279,16 @@ impl PyValue for PyLocal {
     }
 }
 
-impl Drop for PyLocal {
-    fn drop(&mut self) {
-        LOCAL_DATA.with(|map| map.borrow_mut().remove(&self.key));
-    }
-}
-
 #[pyimpl(flags(BASETYPE))]
 impl PyLocal {
     fn ldict(&self, vm: &VirtualMachine) -> PyDictRef {
-        LOCAL_DATA.with(|map| {
-            map.borrow_mut()
-                .entry(self.key)
-                .or_insert_with(|| vm.ctx.new_dict())
-                .clone()
-        })
+        self.data.get_or(|| vm.ctx.new_dict()).clone()
     }
 
     #[pyslot]
     fn tp_new(cls: PyClassRef, _args: PyFuncArgs, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
         PyLocal {
-            key: LOCAL_DATA_KEY.fetch_add(1),
+            data: ThreadLocal::new(),
         }
         .into_ref_with_type(vm, cls)
     }
