Merge pull request #2769 from DimitrisJim/fast_range_iter

Fast range iterator

Lib/test/test_range.py

@@ -484,10 +484,6 @@ class RangeTest(unittest.TestCase):
         self.assertNotIn(-1, r)
         self.assertNotIn(1, r)
 
-    # TODO: RUSTPYTHON.
-    # NOTE: Test passes, fast iterators are required for this to not
-    #       take ~ 1m.
-    @unittest.skip
     def test_range_iterators(self):
         # exercise 'fast' iterators, that use a rangeiterobject internally.
         # see issue 7298

vm/src/builtins/range.rs

@@ -1,7 +1,7 @@
 use crossbeam_utils::atomic::AtomicCell;
 use num_bigint::{BigInt, Sign};
 use num_integer::Integer;
-use num_traits::{One, Signed, Zero};
+use num_traits::{One, Signed, ToPrimitive, Zero};
 
 use super::int::{PyInt, PyIntRef};
 use super::pytype::PyTypeRef;
@@ -171,11 +171,10 @@ impl PyRange {
 
 pub fn init(context: &PyContext) {
     PyRange::extend_class(context, &context.types.range_type);
+    PyLongRangeIterator::extend_class(context, &context.types.longrange_iterator_type);
     PyRangeIterator::extend_class(context, &context.types.range_iterator_type);
 }
 
-type PyRangeRef = PyRef<PyRange>;
-
 #[pyimpl(with(Hashable, Comparable, Iterable))]
 impl PyRange {
     fn new(cls: PyTypeRef, stop: PyIntRef, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
@@ -217,23 +216,39 @@ impl PyRange {
     }
 
     #[pymethod(name = "__reversed__")]
-    fn reversed(&self, vm: &VirtualMachine) -> PyRangeIterator {
+    fn reversed(&self, vm: &VirtualMachine) -> PyResult {
         let start = self.start.as_bigint();
         let step = self.step.as_bigint();
 
         // Use CPython calculation for this:
+        let length = self.len();
         let new_stop = start - step;
-        let new_start = &new_stop + (self.len() * step);
-        let reversed = PyRange {
-            start: new_start.into_pyref(vm),
-            stop: new_stop.into_pyref(vm),
-            step: (-step).into_pyref(vm),
-        };
+        let start = &new_stop + length.clone() * step;
+        let step = -step;
 
-        PyRangeIterator {
-            position: AtomicCell::new(0),
-            range: reversed.into_ref(vm),
-        }
+        Ok(
+            if let (Some(start), Some(step), Some(_)) =
+                (start.to_isize(), step.to_isize(), new_stop.to_isize())
+            {
+                PyRangeIterator {
+                    index: AtomicCell::new(0),
+                    start,
+                    step,
+                    // Cannot fail. If start, stop and step all successfully convert to isize, then result of zelf.len will
+                    // always fit in a usize.
+                    length: length.to_usize().unwrap_or(0),
+                }
+                .into_object(vm)
+            } else {
+                PyLongRangeIterator {
+                    index: AtomicCell::new(0),
+                    start,
+                    step,
+                    length,
+                }
+                .into_object(vm)
+            },
+        )
     }
 
     #[pymethod(name = "__len__")]
@@ -412,19 +427,84 @@ impl Comparable for PyRange {
 
 impl Iterable for PyRange {
     fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        Ok(PyRangeIterator {
-            position: AtomicCell::new(0),
-            range: zelf,
+        let (start, stop, step, length) = (
+            zelf.start.as_bigint(),
+            zelf.stop.as_bigint(),
+            zelf.step.as_bigint(),
+            zelf.len(),
+        );
+        if let (Some(start), Some(step), Some(_)) =
+            (start.to_isize(), step.to_isize(), stop.to_isize())
+        {
+            Ok(PyRangeIterator {
+                index: AtomicCell::new(0),
+                start,
+                step,
+                // Cannot fail. If start, stop and step all successfully convert to isize, then result of zelf.len will
+                // always fit in a usize.
+                length: length.to_usize().unwrap_or(0),
+            }
+            .into_object(vm))
+        } else {
+            Ok(PyLongRangeIterator {
+                index: AtomicCell::new(0),
+                start: start.clone(),
+                step: step.clone(),
+                length,
+            }
+            .into_object(vm))
         }
-        .into_object(vm))
     }
 }
 
+// Semantically, this is the same as the previous representation.
+//
+// Unfortunately, since AtomicCell requires a Copy type, no BigInt implementations can
+// generally be used. As such, usize::MAX is the upper bound on number of elements (length)
+// the range can contain in RustPython.
+//
+// This doesn't preclude the range from containing large values, since start and step
+// can be BigInts, we can store any arbitrary range of values.
+#[pyclass(module = false, name = "longrange_iterator")]
+#[derive(Debug)]
+pub struct PyLongRangeIterator {
+    index: AtomicCell<usize>,
+    start: BigInt,
+    step: BigInt,
+    length: BigInt,
+}
+
+impl PyValue for PyLongRangeIterator {
+    fn class(vm: &VirtualMachine) -> &PyTypeRef {
+        &vm.ctx.types.longrange_iterator_type
+    }
+}
+
+#[pyimpl(with(PyIter))]
+impl PyLongRangeIterator {}
+
+impl PyIter for PyLongRangeIterator {
+    fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
+        let index = BigInt::from(zelf.index.fetch_add(1));
+        if index < zelf.length {
+            Ok(vm
+                .ctx
+                .new_int(zelf.start.clone() + index * zelf.step.clone()))
+        } else {
+            Err(vm.new_stop_iteration())
+        }
+    }
+}
+
+// When start, stop, step are isizes, we can use a faster more compact representation
+// that only operates using isizes to track values.
 #[pyclass(module = false, name = "range_iterator")]
 #[derive(Debug)]
 pub struct PyRangeIterator {
-    position: AtomicCell<usize>,
-    range: PyRangeRef,
+    index: AtomicCell<usize>,
+    start: isize,
+    step: isize,
+    length: usize,
 }
 
 impl PyValue for PyRangeIterator {
@@ -438,9 +518,10 @@ impl PyRangeIterator {}
 
 impl PyIter for PyRangeIterator {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        let position = BigInt::from(zelf.position.fetch_add(1));
-        if let Some(int) = zelf.range.get(&position) {
-            Ok(vm.ctx.new_int(int))
+        if zelf.index.load() < zelf.length {
+            Ok(vm
+                .ctx
+                .new_int(zelf.start + (zelf.index.fetch_add(1) as isize) * zelf.step))
         } else {
             Err(vm.new_stop_iteration())
         }

vm/src/types.rs

@@ -94,6 +94,7 @@ pub struct TypeZoo {
     pub str_type: PyTypeRef,
     pub range_type: PyTypeRef,
     pub range_iterator_type: PyTypeRef,
+    pub longrange_iterator_type: PyTypeRef,
     pub slice_type: PyTypeRef,
     pub type_type: PyTypeRef,
     pub zip_type: PyTypeRef,
@@ -187,6 +188,7 @@ impl TypeZoo {
             module_type: module::PyModule::init_bare_type().clone(),
             namespace_type: namespace::PyNamespace::init_bare_type().clone(),
             range_iterator_type: range::PyRangeIterator::init_bare_type().clone(),
+            longrange_iterator_type: range::PyLongRangeIterator::init_bare_type().clone(),
             set_iterator_type: set::PySetIterator::init_bare_type().clone(),
             str_iterator_type: pystr::PyStrIterator::init_bare_type().clone(),
             str_reverseiterator_type: pystr::PyStrReverseIterator::init_bare_type().clone(),