diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs index 42e383f84..a4bca4e5f 100644 --- a/vm/src/builtins/dict.rs +++ b/vm/src/builtins/dict.rs @@ -718,12 +718,8 @@ macro_rules! dict_iterator { "dictionary changed size during iteration".to_owned(), )); } - let mut position = zelf.position.load(); - match zelf.dict.entries.next_entry(&mut position) { - Some((key, value)) => { - zelf.position.store(position); - Ok(($result_fn)(vm, key, value)) - } + match zelf.dict.entries.next_entry_atomic(&zelf.position) { + Some((key, value)) => Ok(($result_fn)(vm, key, value)), None => { zelf.status.store(IterStatus::Exhausted); Err(vm.new_stop_iteration()) @@ -782,7 +778,7 @@ macro_rules! dict_iterator { "dictionary changed size during iteration".to_owned(), )); } - match zelf.dict.entries.next_entry_reversed(&zelf.position) { + match zelf.dict.entries.next_entry_atomic_reversed(&zelf.position) { Some((key, value)) => Ok(($result_fn)(vm, key, value)), None => { zelf.status.store(IterStatus::Exhausted); diff --git a/vm/src/dictdatatype.rs b/vm/src/dictdatatype.rs index 86d80b722..78c3c2cb1 100644 --- a/vm/src/dictdatatype.rs +++ b/vm/src/dictdatatype.rs @@ -488,7 +488,21 @@ impl Dict { } } - pub fn next_entry_reversed(&self, position: &AtomicCell) -> Option<(PyObjectRef, T)> { + pub fn next_entry_atomic(&self, position: &AtomicCell) -> Option<(PyObjectRef, T)> { + let inner = self.read(); + loop { + let position_usize = position.fetch_add(1); + let entry = inner.entries.get(position_usize)?; + if let Some(entry) = entry { + break Some((entry.key.clone(), entry.value.clone())); + } + } + } + + pub fn next_entry_atomic_reversed( + &self, + position: &AtomicCell, + ) -> Option<(PyObjectRef, T)> { let inner = self.read(); loop { let position_usize = position.fetch_add(1);