Make objiter use atomic types

coolreader18 · coolreader18 · commit a0cd62aa7460 · 2020-04-11T15:16:51.000-05:00
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -2,7 +2,6 @@
 //!
 //! Implements functions listed here: https://docs.python.org/3/library/builtins.html
 
-use std::cell::Cell;
 use std::char;
 use std::str;
 
@@ -672,12 +671,9 @@ fn builtin_reversed(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         vm.get_method_or_type_error(obj.clone(), "__getitem__", || {
             "argument to reversed() must be a sequence".to_owned()
         })?;
-        let len = vm.call_method(&obj.clone(), "__len__", PyFuncArgs::default())?;
-        let obj_iterator = objiter::PySequenceIterator {
-            position: Cell::new(objint::get_value(&len).to_isize().unwrap() - 1),
-            obj: obj.clone(),
-            reversed: true,
-        };
+        let len = vm.call_method(&obj, "__len__", PyFuncArgs::default())?;
+        let len = objint::get_value(&len).to_isize().unwrap();
+        let obj_iterator = objiter::PySequenceIterator::new_reversed(obj, len);
         Ok(obj_iterator.into_ref(vm).into_object())
     }
 }
diff --git a/vm/src/obj/objiter.rs b/vm/src/obj/objiter.rs
@@ -2,8 +2,8 @@
  * Various types to support iteration.
  */
 
+use crossbeam_utils::atomic::AtomicCell;
 use num_traits::{Signed, ToPrimitive};
-use std::cell::Cell;
 
 use super::objint::PyInt;
 use super::objsequence;
@@ -28,12 +28,9 @@ pub fn get_iter(vm: &VirtualMachine, iter_target: &PyObjectRef) -> PyResult {
         vm.get_method_or_type_error(iter_target.clone(), "__getitem__", || {
             format!("Cannot iterate over {}", iter_target.class().name)
         })?;
-        let obj_iterator = PySequenceIterator {
-            position: Cell::new(0),
-            obj: iter_target.clone(),
-            reversed: false,
-        };
-        Ok(obj_iterator.into_ref(vm).into_object())
+        Ok(PySequenceIterator::new_forward(iter_target.clone())
+            .into_ref(vm)
+            .into_object())
     }
 }
 
@@ -140,7 +137,7 @@ pub fn length_hint(vm: &VirtualMachine, iter: PyObjectRef) -> PyResult<Option<us
 #[pyclass]
 #[derive(Debug)]
 pub struct PySequenceIterator {
-    pub position: Cell<isize>,
+    pub position: AtomicCell<isize>,
     pub obj: PyObjectRef,
     pub reversed: bool,
 }
@@ -153,14 +150,31 @@ impl PyValue for PySequenceIterator {
 
 #[pyimpl]
 impl PySequenceIterator {
+    pub fn new_forward(obj: PyObjectRef) -> Self {
+        Self {
+            position: AtomicCell::new(0),
+            obj,
+            reversed: false,
+        }
+    }
+
+    pub fn new_reversed(obj: PyObjectRef, len: isize) -> Self {
+        Self {
+            position: AtomicCell::new(len - 1),
+            obj,
+            reversed: true,
+        }
+    }
+
     #[pymethod(name = "__next__")]
     fn next(&self, vm: &VirtualMachine) -> PyResult {
-        if self.position.get() >= 0 {
+        let pos = self.position.load();
+        if pos >= 0 {
             let step: isize = if self.reversed { -1 } else { 1 };
-            let number = vm.ctx.new_int(self.position.get());
+            let number = vm.ctx.new_int(pos);
             match vm.call_method(&self.obj, "__getitem__", vec![number]) {
                 Ok(val) => {
-                    self.position.set(self.position.get() + step);
+                    self.position.store(pos + step);
                     Ok(val)
                 }
                 Err(ref e) if objtype::isinstance(&e, &vm.ctx.exceptions.index_error) => {
@@ -181,7 +195,7 @@ impl PySequenceIterator {
 
     #[pymethod(name = "__length_hint__")]
     fn length_hint(&self, vm: &VirtualMachine) -> PyResult<isize> {
-        let pos = self.position.get();
+        let pos = self.position.load();
         let hint = if self.reversed {
             pos + 1
         } else {
@@ -195,11 +209,7 @@ impl PySequenceIterator {
 }
 
 pub fn seq_iter_method(obj: PyObjectRef) -> PySequenceIterator {
-    PySequenceIterator {
-        position: Cell::new(0),
-        obj,
-        reversed: false,
-    }
+    PySequenceIterator::new_forward(obj)
 }
 
 #[pyclass(name = "callable_iterator")]
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -1,5 +1,4 @@
 use std::any::Any;
-use std::cell::Cell;
 use std::collections::HashMap;
 use std::fmt;
 use std::marker::PhantomData;
@@ -1012,13 +1011,9 @@ where
             })?;
             Self::try_from_object(
                 vm,
-                objiter::PySequenceIterator {
-                    position: Cell::new(0),
-                    obj: obj.clone(),
-                    reversed: false,
-                }
-                .into_ref(vm)
-                .into_object(),
+                objiter::PySequenceIterator::new_forward(obj)
+                    .into_ref(vm)
+                    .into_object(),
             )
         }
     }
