From 3e656c0b6eb2128f7c8f4a0eacebc85ced6372da Mon Sep 17 00:00:00 2001
From: jfh <jfh@imijmi.eu>
Date: Mon, 30 Aug 2021 17:44:51 +0300
Subject: [PATCH] Use inner dict methods during iteration.

---
 vm/src/builtins/set.rs | 66 +++++++++++++++++++-----------------------
 1 file changed, 29 insertions(+), 37 deletions(-)

diff --git a/vm/src/builtins/set.rs b/vm/src/builtins/set.rs
index b1c9109c2..abfbfb0e7 100644
--- a/vm/src/builtins/set.rs
+++ b/vm/src/builtins/set.rs
@@ -1,10 +1,10 @@
 /*
  * Builtin set type with a sequence of unique items.
  */
-use super::{pytype::PyTypeRef, PyDictRef};
+use super::{pytype::PyTypeRef, IterStatus, PyDictRef};
 use crate::common::hash::PyHash;
 use crate::common::rc::PyRc;
-use crate::dictdatatype;
+use crate::dictdatatype::{self, DictSize};
 use crate::function::{Args, FuncArgs, OptionalArg};
 use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable};
 use crate::vm::{ReprGuard, VirtualMachine};
@@ -186,15 +186,11 @@ impl PySetInner {
     }
 
     fn iter(&self) -> PySetIterator {
-        let set_size = SetSizeInfo {
-            position: 0,
-            size: Some(self.content.len()),
-        };
-
         PySetIterator {
             dict: PyRc::clone(&self.content),
-            elements: self.elements(),
-            size_info: AtomicCell::new(set_size),
+            size: self.content.size(),
+            position: AtomicCell::new(0),
+            status: AtomicCell::new(IterStatus::Active),
         }
     }
 
@@ -803,17 +799,12 @@ impl TryFromObject for SetIterable {
     }
 }
 
-#[derive(Copy, Clone, Default)]
-struct SetSizeInfo {
-    size: Option<usize>,
-    position: usize,
-}
-
 #[pyclass(module = false, name = "set_iterator")]
 pub(crate) struct PySetIterator {
     dict: PyRc<SetContentType>,
-    elements: Vec<PyObjectRef>,
-    size_info: AtomicCell<SetSizeInfo>,
+    size: DictSize,
+    position: AtomicCell<usize>,
+    status: AtomicCell<IterStatus>,
 }
 
 impl fmt::Debug for PySetIterator {
@@ -833,33 +824,34 @@ impl PyValue for PySetIterator {
 impl PySetIterator {
     #[pymethod(magic)]
     fn length_hint(&self) -> usize {
-        let set_len = self.dict.len();
-        let position = self.size_info.load().position;
-        set_len.saturating_sub(position)
+        if let IterStatus::Exhausted = self.status.load() {
+            0
+        } else {
+            self.dict.len_from_entry_index(self.position.load())
+        }
     }
 }
 
 impl PyIter for PySetIterator {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        let mut size_info = zelf.size_info.take();
-
-        if let Some(set_size) = size_info.size {
-            if set_size == zelf.dict.len() {
-                let index = size_info.position;
-                let item = zelf
-                    .elements
-                    .get(index)
-                    .ok_or_else(|| vm.new_stop_iteration())?;
-                size_info.position += 1;
-                zelf.size_info.store(size_info);
-                return Ok(item.clone());
-            } else {
-                size_info.size = None;
-                zelf.size_info.store(size_info);
+        match zelf.status.load() {
+            IterStatus::Exhausted => Err(vm.new_stop_iteration()),
+            IterStatus::Active => {
+                if zelf.dict.has_changed_size(&zelf.size) {
+                    zelf.status.store(IterStatus::Exhausted);
+                    return Err(
+                        vm.new_runtime_error("set changed size during iteration".to_owned())
+                    );
+                }
+                match zelf.dict.next_entry_atomic(&zelf.position) {
+                    Some((key, _)) => Ok(key),
+                    None => {
+                        zelf.status.store(IterStatus::Exhausted);
+                        Err(vm.new_stop_iteration())
+                    }
+                }
             }
         }
-
-        Err(vm.new_runtime_error("set changed size during iteration".into()))
     }
 }