Convert Dict to ThreadSafe

vm/src/dictdatatype.rs

@@ -1,6 +1,6 @@
 use crate::obj::objstr::PyString;
 use crate::pyhash;
-use crate::pyobject::{IdProtocol, IntoPyObject, PyObjectRef, PyResult};
+use crate::pyobject::{IdProtocol, IntoPyObject, PyObjectRef, PyResult, ThreadSafe};
 use crate::vm::VirtualMachine;
 use num_bigint::ToBigInt;
 /// Ordered dictionary implementation.
@@ -10,6 +10,7 @@ use num_bigint::ToBigInt;
 use std::collections::{hash_map::DefaultHasher, HashMap};
 use std::hash::{Hash, Hasher};
 use std::mem::size_of;
+use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 /// hash value of an object returned by __hash__
 type HashValue = pyhash::PyHash;
@@ -18,30 +19,64 @@ type HashIndex = pyhash::PyHash;
 /// entry index mapped in indices
 type EntryIndex = usize;
 
-#[derive(Clone)]
 pub struct Dict<T = PyObjectRef> {
+    inner: RwLock<InnerDict<T>>,
+}
+
+impl ThreadSafe for Dict {}
+
+struct InnerDict<T> {
     size: usize,
     indices: HashMap<HashIndex, EntryIndex>,
     entries: Vec<Option<DictEntry<T>>>,
 }
 
+impl<T: Clone> Clone for InnerDict<T> {
+    fn clone(&self) -> Self {
+        InnerDict {
+            size: self.size,
+            indices: self.indices.clone(),
+            entries: self.entries.clone(),
+        }
+    }
+}
+
+impl<T: Clone> Clone for Dict<T> {
+    fn clone(&self) -> Self {
+        Dict {
+            inner: RwLock::new(self.inner.read().unwrap().clone()),
+        }
+    }
+}
+
 impl<T> Default for Dict<T> {
     fn default() -> Self {
         Dict {
-            size: 0,
-            indices: HashMap::new(),
-            entries: Vec::new(),
+            inner: RwLock::new(InnerDict {
+                size: 0,
+                indices: HashMap::new(),
+                entries: Vec::new(),
+            }),
         }
     }
 }
 
-#[derive(Clone)]
 struct DictEntry<T> {
     hash: HashValue,
     key: PyObjectRef,
     value: T,
 }
 
+impl<T: Clone> Clone for DictEntry<T> {
+    fn clone(&self) -> Self {
+        DictEntry {
+            hash: self.hash,
+            key: self.key.clone(),
+            value: self.value.clone(),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct DictSize {
     size: usize,
@@ -49,10 +84,19 @@ pub struct DictSize {
 }
 
 impl<T: Clone> Dict<T> {
-    fn resize(&mut self) {
-        let mut new_indices = HashMap::with_capacity(self.size);
-        let mut new_entries = Vec::with_capacity(self.size);
-        for maybe_entry in self.entries.drain(0..) {
+    fn borrow_value(&self) -> RwLockReadGuard<'_, InnerDict<T>> {
+        self.inner.read().unwrap()
+    }
+
+    fn borrow_value_mut(&self) -> RwLockWriteGuard<'_, InnerDict<T>> {
+        self.inner.write().unwrap()
+    }
+
+    fn resize(&self) {
+        let mut inner = self.borrow_value_mut();
+        let mut new_indices = HashMap::with_capacity(inner.size);
+        let mut new_entries = Vec::with_capacity(inner.size);
+        for maybe_entry in inner.entries.drain(0..) {
             if let Some(entry) = maybe_entry {
                 let mut hash_index = entry.hash;
                 // Faster version of lookup. No equality checks here.
@@ -64,8 +108,8 @@ impl<T: Clone> Dict<T> {
                 new_entries.push(Some(entry));
             }
         }
-        self.indices = new_indices;
-        self.entries = new_entries;
+        inner.indices = new_indices;
+        inner.entries = new_entries;
     }
 
     fn unchecked_push(
@@ -80,15 +124,11 @@ impl<T: Clone> Dict<T> {
             key,
             value,
         };
-        let entry_index = self.entries.len();
-        self.entries.push(Some(entry));
-        self.indices.insert(hash_index, entry_index);
-        self.size += 1;
-    }
-
-    fn unchecked_delete(&mut self, entry_index: EntryIndex) {
-        self.entries[entry_index] = None;
-        self.size -= 1;
+        let mut inner = self.borrow_value_mut();
+        let entry_index = inner.entries.len();
+        inner.entries.push(Some(entry));
+        inner.indices.insert(hash_index, entry_index);
+        inner.size += 1;
     }
 
     /// Store a key
@@ -98,26 +138,32 @@ impl<T: Clone> Dict<T> {
         key: K,
         value: T,
     ) -> PyResult<()> {
-        if self.indices.len() > 2 * self.size {
+        // This does not need to be accurate so we can take the lock mutiple times.
+        if self.borrow_value().indices.len() > 2 * self.borrow_value().size {
             self.resize();
         }
-        match self.lookup(vm, key)? {
-            LookupResult::Existing(index) => {
-                // Update existing key
-                if let Some(ref mut entry) = self.entries[index] {
-                    entry.value = value;
-                    Ok(())
-                } else {
-                    panic!("Lookup returned invalid index into entries!");
+        loop {
+            match self.lookup(vm, key)? {
+                LookupResult::Existing(index) => {
+                    let mut inner = self.borrow_value_mut();
+                    // Update existing key
+                    if let Some(ref mut entry) = inner.entries[index] {
+                        // They entry might have changed since we did lookup. Should we update the key?
+                        entry.value = value;
+                        break Ok(());
+                    } else {
+                        // The dict was changed since we did lookup. Let's try again.
+                        continue;
+                    }
+                }
+                LookupResult::NewIndex {
+                    hash_index,
+                    hash_value,
+                } => {
+                    // New key:
+                    self.unchecked_push(hash_index, hash_value, key.into_pyobject(vm)?, value);
+                    break Ok(());
                 }
-            }
-            LookupResult::NewIndex {
-                hash_index,
-                hash_value,
-            } => {
-                // New key:
-                self.unchecked_push(hash_index, hash_value, key.into_pyobject(vm)?, value);
-                Ok(())
             }
         }
     }
@@ -130,28 +176,28 @@ impl<T: Clone> Dict<T> {
         }
     }
 
-    fn unchecked_get(&self, index: EntryIndex) -> T {
-        if let Some(entry) = &self.entries[index] {
-            entry.value.clone()
-        } else {
-            panic!("Lookup returned invalid index into entries!");
-        }
-    }
-
     /// Retrieve a key
     #[cfg_attr(feature = "flame-it", flame("Dict"))]
     pub fn get<K: DictKey + Copy>(&self, vm: &VirtualMachine, key: K) -> PyResult<Option<T>> {
-        if let LookupResult::Existing(index) = self.lookup(vm, key)? {
-            Ok(Some(self.unchecked_get(index)))
-        } else {
-            Ok(None)
+        loop {
+            if let LookupResult::Existing(index) = self.lookup(vm, key)? {
+                if let Some(entry) = &self.borrow_value().entries[index] {
+                    break Ok(Some(entry.value.clone()));
+                } else {
+                    // The dict was changed since we did lookup. Let's try again.
+                    continue;
+                }
+            } else {
+                break Ok(None);
+            }
         }
     }
 
     pub fn clear(&mut self) {
-        self.entries.clear();
-        self.indices.clear();
-        self.size = 0
+        let mut inner = self.borrow_value_mut();
+        inner.entries.clear();
+        inner.indices.clear();
+        inner.size = 0
     }
 
     /// Delete a key
@@ -164,11 +210,20 @@ impl<T: Clone> Dict<T> {
     }
 
     pub fn delete_if_exists(&mut self, vm: &VirtualMachine, key: &PyObjectRef) -> PyResult<bool> {
-        if let LookupResult::Existing(entry_index) = self.lookup(vm, key)? {
-            self.unchecked_delete(entry_index);
-            Ok(true)
-        } else {
-            Ok(false)
+        loop {
+            if let LookupResult::Existing(entry_index) = self.lookup(vm, key)? {
+                let mut inner = self.borrow_value_mut();
+                if let Some(_) = &inner.entries[entry_index] {
+                    inner.entries[entry_index] = None;
+                    inner.size -= 1;
+                    break Ok(true);
+                } else {
+                    // The dict was changed since we did lookup. Let's try again.
+                    continue;
+                }
+            } else {
+                break Ok(false);
+            }
         }
     }
 
@@ -178,18 +233,32 @@ impl<T: Clone> Dict<T> {
         key: &PyObjectRef,
         value: T,
     ) -> PyResult<()> {
-        match self.lookup(vm, key)? {
-            LookupResult::Existing(entry_index) => self.unchecked_delete(entry_index),
-            LookupResult::NewIndex {
-                hash_value,
-                hash_index,
-            } => self.unchecked_push(hash_index, hash_value, key.clone(), value),
-        };
-        Ok(())
+        loop {
+            match self.lookup(vm, key)? {
+                LookupResult::Existing(entry_index) => {
+                    let mut inner = self.borrow_value_mut();
+                    if let Some(_) = &inner.entries[entry_index] {
+                        inner.entries[entry_index] = None;
+                        inner.size -= 1;
+                        break Ok(());
+                    } else {
+                        // The dict was changed since we did lookup. Let's try again.
+                        continue;
+                    }
+                }
+                LookupResult::NewIndex {
+                    hash_value,
+                    hash_index,
+                } => {
+                    self.unchecked_push(hash_index, hash_value, key.clone(), value);
+                    break Ok(());
+                }
+            };
+        }
     }
 
     pub fn len(&self) -> usize {
-        self.size
+        self.borrow_value().size
     }
 
     pub fn is_empty(&self) -> bool {
@@ -197,33 +266,42 @@ impl<T: Clone> Dict<T> {
     }
 
     pub fn size(&self) -> DictSize {
+        let inner = self.borrow_value();
         DictSize {
-            size: self.size,
-            entries_size: self.entries.len(),
+            size: inner.size,
+            entries_size: inner.entries.len(),
         }
     }
 
-    pub fn next_entry(&self, position: &mut EntryIndex) -> Option<(&PyObjectRef, &T)> {
-        self.entries[*position..].iter().find_map(|entry| {
-            *position += 1;
-            entry
-                .as_ref()
-                .map(|DictEntry { key, value, .. }| (key, value))
-        })
+    pub fn next_entry(&self, position: &mut EntryIndex) -> Option<(PyObjectRef, T)> {
+        self.borrow_value().entries[*position..]
+            .iter()
+            .find_map(|entry| {
+                *position += 1;
+                entry
+                    .as_ref()
+                    .map(|DictEntry { key, value, .. }| (key.clone(), value.clone()))
+            })
     }
 
     pub fn len_from_entry_index(&self, position: EntryIndex) -> usize {
-        self.entries[position..].iter().flatten().count()
+        self.borrow_value().entries[position..]
+            .iter()
+            .flatten()
+            .count()
     }
 
     pub fn has_changed_size(&self, position: &DictSize) -> bool {
-        position.size != self.size || self.entries.len() != position.entries_size
+        let inner = self.borrow_value();
+        position.size != inner.size || inner.entries.len() != position.entries_size
     }
 
-    pub fn keys<'a>(&'a self) -> impl Iterator<Item = PyObjectRef> + 'a {
-        self.entries
+    pub fn keys(&self) -> Vec<PyObjectRef> {
+        self.borrow_value()
+            .entries
             .iter()
             .filter_map(|v| v.as_ref().map(|v| v.key.clone()))
+            .collect()
     }
 
     /// Lookup the index for the given key.
@@ -232,33 +310,41 @@ impl<T: Clone> Dict<T> {
         let hash_value = key.do_hash(vm)?;
         let perturb = hash_value;
         let mut hash_index: HashIndex = hash_value;
-        loop {
-            if self.indices.contains_key(&hash_index) {
-                // Now we have an index, lets check the key.
-                let index = self.indices[&hash_index];
-                if let Some(entry) = &self.entries[index] {
-                    // Okay, we have an entry at this place
-                    if key.do_is(&entry.key) {
-                        // Literally the same object
-                        break Ok(LookupResult::Existing(index));
-                    } else if entry.hash == hash_value {
-                        if key.do_eq(vm, &entry.key)? {
-                            break Ok(LookupResult::Existing(index));
+        'outer: loop {
+            let (entry, index) = loop {
+                let inner = self.borrow_value();
+                if inner.indices.contains_key(&hash_index) {
+                    // Now we have an index, lets check the key.
+                    let index = inner.indices[&hash_index];
+                    if let Some(entry) = &inner.entries[index] {
+                        // Okay, we have an entry at this place
+                        if key.do_is(&entry.key) {
+                            // Literally the same object
+                            break 'outer Ok(LookupResult::Existing(index));
+                        } else if entry.hash == hash_value {
+                            break (entry.clone(), index);
                         } else {
                             // entry mismatch.
                         }
                     } else {
-                        // entry mismatch.
+                        // Removed entry, continue search...
                     }
                 } else {
-                    // Removed entry, continue search...
+                    // Hash not in table, we are at free slot now.
+                    break 'outer Ok(LookupResult::NewIndex {
+                        hash_value,
+                        hash_index,
+                    });
                 }
+                // Update i to next probe location:
+                hash_index = Self::next_index(perturb, hash_index)
+                // warn!("Perturb value: {}", i);
+            };
+            // This comparison needs to be done outside the lock.
+            if key.do_eq(vm, &entry.key)? {
+                break Ok(LookupResult::Existing(index));
             } else {
-                // Hash not in table, we are at free slot now.
-                break Ok(LookupResult::NewIndex {
-                    hash_value,
-                    hash_index,
-                });
+                // entry mismatch.
             }
 
             // Update i to next probe location:
@@ -276,29 +362,44 @@ impl<T: Clone> Dict<T> {
 
     /// Retrieve and delete a key
     pub fn pop<K: DictKey + Copy>(&mut self, vm: &VirtualMachine, key: K) -> PyResult<Option<T>> {
-        if let LookupResult::Existing(index) = self.lookup(vm, key)? {
-            let value = self.unchecked_get(index);
-            self.unchecked_delete(index);
-            Ok(Some(value))
-        } else {
-            Ok(None)
+        loop {
+            if let LookupResult::Existing(index) = self.lookup(vm, key)? {
+                let mut inner = self.borrow_value_mut();
+                if let Some(entry) = &inner.entries[index] {
+                    let value = entry.value.clone();
+                    inner.entries[index] = None;
+                    inner.size -= 1;
+                    break Ok(Some(value));
+                } else {
+                    // The dict was changed since we did lookup. Let's try again.
+                    continue;
+                }
+            } else {
+                break Ok(None);
+            }
         }
     }
 
     pub fn pop_front(&mut self) -> Option<(PyObjectRef, T)> {
-        let mut entry_index = 0;
-        match self.next_entry(&mut entry_index) {
-            Some((key, value)) => {
-                let item = (key.clone(), value.clone());
-                self.unchecked_delete(entry_index - 1);
-                Some(item)
-            }
-            None => None,
+        let mut position = 0;
+        let mut inner = self.borrow_value_mut();
+        let first_item = inner.entries.iter().find_map(|entry| {
+            position += 1;
+            entry
+                .as_ref()
+                .map(|DictEntry { key, value, .. }| (key.clone(), value.clone()))
+        });
+        if let Some(item) = first_item {
+            inner.entries[position - 1] = None;
+            inner.size -= 1;
+            Some(item)
+        } else {
+            None
         }
     }
 
     pub fn sizeof(&self) -> usize {
-        size_of::<Self>() + self.size * size_of::<DictEntry<T>>()
+        size_of::<Self>() + self.borrow_value().size * size_of::<DictEntry<T>>()
     }
 }
 

vm/src/obj/objdict.rs

@@ -488,7 +488,7 @@ impl Iterator for DictIter {
 
     fn next(&mut self) -> Option<Self::Item> {
         match self.dict.entries.borrow().next_entry(&mut self.position) {
-            Some((key, value)) => Some((key.clone(), value.clone())),
+            Some((key, value)) => Some((key, value)),
             None => None,
         }
     }
@@ -524,7 +524,7 @@ macro_rules! dict_iterator {
                 let s = if let Some(_guard) = ReprGuard::enter(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))?;
+                        let s = vm.to_repr(&$result_fn(vm, key, value))?;
                         str_parts.push(s.as_str().to_owned());
                     }
                     format!("{}([{}])", $class_name, str_parts.join(", "))
@@ -607,7 +607,7 @@ dict_iterator! {
     dictkeyiterator_type,
     "dict_keys",
     "dictkeyiterator",
-    |_vm: &VirtualMachine, key: &PyObjectRef, _value: &PyObjectRef| key.clone()
+    |_vm: &VirtualMachine, key: PyObjectRef, _value: PyObjectRef| key
 }
 
 dict_iterator! {
@@ -617,7 +617,7 @@ dict_iterator! {
     dictvalueiterator_type,
     "dict_values",
     "dictvalueiterator",
-    |_vm: &VirtualMachine, _key: &PyObjectRef, value: &PyObjectRef| value.clone()
+    |_vm: &VirtualMachine, _key: PyObjectRef, value: PyObjectRef| value
 }
 
 dict_iterator! {
@@ -627,8 +627,8 @@ dict_iterator! {
     dictitemiterator_type,
     "dict_items",
     "dictitemiterator",
-    |vm: &VirtualMachine, key: &PyObjectRef, value: &PyObjectRef|
-        vm.ctx.new_tuple(vec![key.clone(), value.clone()])
+    |vm: &VirtualMachine, key: PyObjectRef, value: PyObjectRef|
+        vm.ctx.new_tuple(vec![key, value])
 }
 
 pub(crate) fn init(context: &PyContext) {

vm/src/obj/objset.rs

@@ -241,8 +241,7 @@ impl PySetInner {
     }
 
     fn iter(&self, vm: &VirtualMachine) -> PyListIterator {
-        let items = self.content.keys().collect();
-        let set_list = vm.ctx.new_list(items);
+        let set_list = vm.ctx.new_list(self.content.keys());
         PyListIterator {
             position: crossbeam_utils::atomic::AtomicCell::new(0),
             list: set_list.downcast().unwrap(),
@@ -337,7 +336,7 @@ impl PySetInner {
     }
 
     fn hash(&self, vm: &VirtualMachine) -> PyResult<pyhash::PyHash> {
-        pyhash::hash_iter_unordered(self.content.keys(), vm)
+        pyhash::hash_iter_unordered(self.content.keys().iter(), vm)
     }
 }
 

vm/src/pyhash.rs

@@ -84,13 +84,13 @@ pub fn hash_iter<'a, I: std::iter::Iterator<Item = &'a PyObjectRef>>(
     Ok(hasher.finish() as PyHash)
 }
 
-pub fn hash_iter_unordered<I: std::iter::Iterator<Item = PyObjectRef>>(
+pub fn hash_iter_unordered<'a, I: std::iter::Iterator<Item = &'a PyObjectRef>>(
     iter: I,
     vm: &VirtualMachine,
 ) -> PyResult<PyHash> {
     let mut hash: PyHash = 0;
     for element in iter {
-        let item_hash = vm._hash(&element)?;
+        let item_hash = vm._hash(element)?;
         // xor is commutative and hash should be independent of order
         hash ^= item_hash;
     }