From 131736eecfcc97fbcc047051ba048fbe107690b7 Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Thu, 24 Sep 2020 23:51:33 -0500
Subject: [PATCH] Make setdefault atomic

---
 Lib/test/test_dict.py  |  2 --
 vm/src/dictdatatype.rs | 47 +++++++++++++++++++++++++++++++++++-------
 vm/src/obj/objdict.rs  | 10 ++-------
 3 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/Lib/test/test_dict.py b/Lib/test/test_dict.py
index 05b8edf97..1782b89c1 100644
--- a/Lib/test/test_dict.py
+++ b/Lib/test/test_dict.py
@@ -368,8 +368,6 @@ class DictTest(unittest.TestCase):
         x.fail = True
         self.assertRaises(Exc, d.setdefault, x, [])
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_setdefault_atomic(self):
         # Issue #13521: setdefault() calls __hash__ and __eq__ only once.
         class Hashed(object):
diff --git a/vm/src/dictdatatype.rs b/vm/src/dictdatatype.rs
index d02dddb3b..535330eb3 100644
--- a/vm/src/dictdatatype.rs
+++ b/vm/src/dictdatatype.rs
@@ -157,6 +157,7 @@ impl<T> DictInner<T> {
         hash_value: HashValue,
         key: PyObjectRef,
         value: T,
+        index_entry: IndexEntry,
     ) {
         let entry = DictEntry {
             hash: hash_value,
@@ -168,6 +169,12 @@ impl<T> DictInner<T> {
         self.entries.push(entry);
         self.indices[index] = entry_index as i64;
         self.used += 1;
+        if let IndexEntry::Free = index_entry {
+            self.filled += 1;
+            if let Some(new_size) = self.should_resize() {
+                self.resize(new_size)
+            }
+        }
     }
 
     fn size(&self) -> DictSize {
@@ -223,13 +230,7 @@ impl<T: Clone> Dict<T> {
             } else {
                 // New key:
                 let mut inner = self.borrow_value_mut();
-                inner.unchecked_push(index_index, hash, key.into_pyobject(vm), value);
-                if let IndexEntry::Free = entry_index {
-                    inner.filled += 1;
-                    if let Some(new_size) = inner.should_resize() {
-                        inner.resize(new_size)
-                    }
-                }
+                inner.unchecked_push(index_index, hash, key.into_pyobject(vm), value, entry_index);
                 break None;
             }
         };
@@ -328,13 +329,43 @@ impl<T: Clone> Dict<T> {
                 }
             } else {
                 let mut inner = self.borrow_value_mut();
-                inner.unchecked_push(index_index, hash, key.clone(), value);
+                inner.unchecked_push(index_index, hash, key.clone(), value, entry);
                 break None;
             }
         };
         Ok(())
     }
 
+    pub fn setdefault<F>(&self, vm: &VirtualMachine, key: PyObjectRef, default: F) -> PyResult<T>
+    where
+        F: FnOnce() -> T,
+    {
+        let hash = key.key_hash(vm)?;
+        let res = loop {
+            let lookup = self.lookup(vm, &key, hash, None)?;
+            let (entry, index_index) = lookup;
+            if let IndexEntry::Index(index) = entry {
+                let inner = self.borrow_value();
+                if let Some(entry) = inner.entries.get(index) {
+                    if entry.index == index_index {
+                        break entry.value.clone();
+                    } else {
+                        // stuff shifted around, let's try again
+                    }
+                } else {
+                    // The dict was changed since we did lookup, let's try again.
+                    continue;
+                }
+            } else {
+                let value = default();
+                let mut inner = self.borrow_value_mut();
+                inner.unchecked_push(index_index, hash, key, value.clone(), entry);
+                break value;
+            }
+        };
+        Ok(res)
+    }
+
     pub fn len(&self) -> usize {
         self.borrow_value().used
     }
diff --git a/vm/src/obj/objdict.rs b/vm/src/obj/objdict.rs
index d703836ae..c45997257 100644
--- a/vm/src/obj/objdict.rs
+++ b/vm/src/obj/objdict.rs
@@ -299,14 +299,8 @@ impl PyDict {
         default: OptionalArg<PyObjectRef>,
         vm: &VirtualMachine,
     ) -> PyResult {
-        match self.entries.get(vm, &key)? {
-            Some(value) => Ok(value),
-            None => {
-                let set_value = default.unwrap_or_none(vm);
-                self.entries.insert(vm, key, set_value.clone())?;
-                Ok(set_value)
-            }
-        }
+        self.entries
+            .setdefault(vm, key, || default.unwrap_or_none(vm))
     }
 
     #[pymethod]