Guard for changes in dictionary size during iteration.

Author: cthulahoops

tests/snippets/dict.py

@@ -69,6 +69,18 @@
 data[3] = "changed"
 assert (3, "changed") == next(items)
 
+# But we can't add or delete items during iteration.
+d = {}
+a = iter(d.items())
+d['a'] = 2
+b = iter(d.items())
+assert ('a', 2) == next(b)
+with assertRaises(RuntimeError):
+    next(a)
+del d['a']
+with assertRaises(RuntimeError):
+    next(b)
+
 # View isn't itself an iterator.
 with assertRaises(TypeError):
     next(data.keys())

vm/src/dictdatatype.rs

@@ -33,6 +33,12 @@ struct DictEntry<T> {
     value: T,
 }
 
+#[derive(Debug)]
+pub struct DictSize {
+    size: usize,
+    entries_size: usize,
+}
+
 impl<T: Clone> Dict<T> {
     pub fn new() -> Self {
         Dict {
@@ -120,17 +126,28 @@ impl<T: Clone> Dict<T> {
         self.len() == 0
     }
 
-    pub fn next_entry(&self, mut position: usize) -> Option<(usize, &PyObjectRef, &T)> {
-        while position < self.entries.len() {
-            if let Some(DictEntry { key, value, .. }) = &self.entries[position] {
-                return Some((position + 1, key, value));
-            } else {
-                position += 1;
+    pub fn size(&self) -> DictSize {
+        DictSize {
+            size: self.size,
+            entries_size: self.entries.len(),
+        }
+    }
+
+    pub fn next_entry(&self, position: &mut usize) -> Option<(&PyObjectRef, &T)> {
+        while *position < self.entries.len() {
+            if let Some(DictEntry { key, value, .. }) = &self.entries[*position] {
+                *position += 1;
+                return Some((key, value));
             }
+            *position += 1;
         }
         None
     }
 
+    pub fn has_changed_size(&self, position: &DictSize) -> bool {
+        position.size != self.size || self.entries.len() != position.entries_size
+    }
+
     /// Lookup the index for the given key.
     fn lookup(&self, vm: &VirtualMachine, key: &PyObjectRef) -> PyResult<LookupResult> {
         let hash_value = calc_hash(vm, key)?;

vm/src/obj/objdict.rs

@@ -260,6 +260,10 @@ impl PyDictRef {
         let key = key.into_pyobject(vm).unwrap();
         self.entries.borrow().contains(vm, &key).unwrap()
     }
+
+    pub fn size(&self) -> dictdatatype::DictSize {
+        self.entries.borrow().size()
+    }
 }
 
 impl ItemProtocol for PyDictRef {
@@ -315,11 +319,8 @@ impl Iterator for DictIter {
     type Item = (PyObjectRef, PyObjectRef);
 
     fn next(&mut self) -> Option<Self::Item> {
-        match self.dict.entries.borrow().next_entry(self.position) {
-            Some((new_position, key, value)) => {
-                self.position = new_position;
-                Some((key.clone(), value.clone()))
-            }
+        match self.dict.entries.borrow().next_entry(&mut self.position) {
+            Some((key, value)) => Some((key.clone(), value.clone())),
             None => None,
         }
     }
@@ -360,6 +361,7 @@ macro_rules! dict_iterator {
         #[derive(Debug)]
         struct $iter_name {
             pub dict: PyDictRef,
+            pub size: dictdatatype::DictSize,
             pub position: Cell<usize>,
         }
 
@@ -368,15 +370,24 @@ macro_rules! dict_iterator {
             fn new(dict: PyDictRef) -> Self {
                 $iter_name {
                     position: Cell::new(0),
+                    size: dict.size(),
                     dict,
                 }
             }
 
             #[pymethod(name = "__next__")]
             fn next(&self, vm: &VirtualMachine) -> PyResult {
-                match self.dict.entries.borrow().next_entry(self.position.get()) {
-                    Some((new_position, key, value)) => {
-                        self.position.set(new_position);
+                let mut position = self.position.get();
+                let dict = self.dict.entries.borrow();
+                if dict.has_changed_size(&self.size) {
+                    return Err(vm.new_exception(
+                        vm.ctx.exceptions.runtime_error.clone(),
+                        "dictionary changed size during iteration".to_string(),
+                    ));
+                }
+                match dict.next_entry(&mut position) {
+                    Some((key, value)) => {
+                        self.position.set(position);
                         Ok($result_fn(vm, key, value))
                     }
                     None => Err(objiter::new_stop_iteration(vm)),