From 8978e8ccdafbfebd2ee50f1efc0d98b32627b1ca Mon Sep 17 00:00:00 2001 From: jfh Date: Sat, 14 Aug 2021 22:07:41 +0300 Subject: [PATCH] Add reverse deque iterator. --- Lib/test/test_deque.py | 4 - Lib/test/test_iterlen.py | 8 -- vm/src/stdlib/collections.rs | 140 ++++++++++++++++++++++++++++++++--- 3 files changed, 131 insertions(+), 21 deletions(-) diff --git a/Lib/test/test_deque.py b/Lib/test/test_deque.py index 191b3a238..ad768222c 100644 --- a/Lib/test/test_deque.py +++ b/Lib/test/test_deque.py @@ -820,16 +820,12 @@ class TestVariousIteratorArgs(unittest.TestCase): self.assertRaises(TypeError, deque, seq_tests.IterNoNext(s)) self.assertRaises(ZeroDivisionError, deque, seq_tests.IterGenExc(s)) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_iter_with_altered_data(self): d = deque('abcdefg') it = iter(d) d.pop() self.assertRaises(RuntimeError, next, it) - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_runtime_error_on_empty_deque(self): d = deque() it = iter(d) diff --git a/Lib/test/test_iterlen.py b/Lib/test/test_iterlen.py index db01b5487..40bb39018 100644 --- a/Lib/test/test_iterlen.py +++ b/Lib/test/test_iterlen.py @@ -114,13 +114,9 @@ class TestDeque(TestTemporarilyImmutable, unittest.TestCase): self.it = iter(d) self.mutate = d.pop - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_immutable_during_iteration(self): super().test_immutable_during_iteration() - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_invariant(self): super().test_invariant() @@ -131,13 +127,9 @@ class TestDequeReversed(TestTemporarilyImmutable, unittest.TestCase): self.it = reversed(d) self.mutate = d.pop - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_immutable_during_iteration(self): super().test_immutable_during_iteration() - # TODO: RUSTPYTHON - @unittest.expectedFailure def test_invariant(self): super().test_invariant() diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs index 3bee31e8c..5773a68a5 100644 --- a/vm/src/stdlib/collections.rs +++ b/vm/src/stdlib/collections.rs @@ -2,12 +2,18 @@ pub(crate) use _collections::make_module; #[pymodule] mod _collections { - use crate::builtins::{PyInt, PyTypeRef}; use crate::common::lock::{PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard}; use crate::function::{FuncArgs, OptionalArg}; use crate::slots::{Comparable, Hashable, Iterable, PyComparisonOp, PyIter, Unhashable}; use crate::vm::ReprGuard; use crate::VirtualMachine; + use crate::{ + builtins::{ + IterStatus::{self, Active, Exhausted}, + PyInt, PyTypeRef, + }, + TypeProtocol, + }; use crate::{sequence, sliceable}; use crate::{PyComparisonValue, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, StaticType}; use crossbeam_utils::atomic::AtomicCell; @@ -288,6 +294,18 @@ mod _collections { *self.borrow_deque_mut() = rev; } + #[pymethod(magic)] + fn reversed(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + let length = zelf.len(); + Ok(PyReverseDequeIterator { + position: AtomicCell::new(length), + status: AtomicCell::new(if length > 0 { Active } else { Exhausted }), + length, + deque: zelf, + } + .into_object(vm)) + } + #[pymethod] fn rotate(&self, mid: OptionalArg) { let mut deque = self.borrow_deque_mut(); @@ -415,6 +433,8 @@ mod _collections { fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult { Ok(PyDequeIterator { position: AtomicCell::new(0), + status: AtomicCell::new(IterStatus::Active), + length: zelf.len(), deque: zelf, } .into_object(vm)) @@ -426,6 +446,8 @@ mod _collections { #[derive(Debug)] struct PyDequeIterator { position: AtomicCell, + status: AtomicCell, + length: usize, // To track length immutability. deque: PyDequeRef, } @@ -436,17 +458,117 @@ mod _collections { } #[pyimpl(with(PyIter))] - impl PyDequeIterator {} + impl PyDequeIterator { + #[pymethod(magic)] + fn length_hint(&self) -> usize { + match self.status.load() { + Active => self.deque.len().saturating_sub(self.position.load()), + Exhausted => 0, + } + } + + #[pymethod(magic)] + fn reduce( + zelf: PyRef, + vm: &VirtualMachine, + ) -> PyResult<(PyTypeRef, (PyDequeRef, PyObjectRef))> { + Ok(( + zelf.clone_class(), + (zelf.deque.clone(), vm.ctx.new_int(zelf.position.load())), + )) + } + } impl PyIter for PyDequeIterator { fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - let pos = zelf.position.fetch_add(1); - let deque = zelf.deque.borrow_deque(); - if pos < deque.len() { - let ret = deque[pos].clone(); - Ok(ret) - } else { - Err(vm.new_stop_iteration()) + match zelf.status.load() { + Exhausted => Err(vm.new_stop_iteration()), + Active => { + if zelf.length != zelf.deque.len() { + // Deque was changed while we iterated. + zelf.status.store(Exhausted); + Err(vm.new_runtime_error("Deque mutated during iteration".to_owned())) + } else { + let pos = zelf.position.fetch_add(1); + let deque = zelf.deque.borrow_deque(); + if pos < deque.len() { + let ret = deque[pos].clone(); + Ok(ret) + } else { + zelf.status.store(Exhausted); + Err(vm.new_stop_iteration()) + } + } + } + } + } + } + + #[pyattr] + #[pyclass(name = "_deque_reverse_iterator")] + #[derive(Debug)] + struct PyReverseDequeIterator { + position: AtomicCell, + status: AtomicCell, + length: usize, // To track length immutability. + deque: PyDequeRef, + } + + impl PyValue for PyReverseDequeIterator { + fn class(_vm: &VirtualMachine) -> &PyTypeRef { + Self::static_type() + } + } + + #[pyimpl(with(PyIter))] + impl PyReverseDequeIterator { + #[pymethod(magic)] + fn length_hint(&self) -> usize { + match self.status.load() { + Active => self.position.load(), + Exhausted => 0, + } + } + + #[pymethod(magic)] + fn reduce( + zelf: PyRef, + vm: &VirtualMachine, + ) -> PyResult<(PyTypeRef, (PyDequeRef, PyObjectRef))> { + Ok(( + zelf.clone_class(), + (zelf.deque.clone(), vm.ctx.new_int(zelf.position.load())), + )) + } + } + + impl PyIter for PyReverseDequeIterator { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + match zelf.status.load() { + Exhausted => Err(vm.new_stop_iteration()), + Active => { + // If length changes while we iterate, set to Exhausted and bail. + if zelf.length != zelf.deque.len() { + zelf.status.store(Exhausted); + Err(vm.new_runtime_error("Deque mutated during iteration".to_owned())) + } else { + let pos = zelf.position.fetch_sub(1) - 1; + let deque = zelf.deque.borrow_deque(); + if pos > 0 { + if let Some(obj) = deque.get(pos) { + return Ok(obj.clone()); + } + } + // We either are == 0 or deque.get returned None. Either way, set status + // to exhausted and return last item if pos == 0. + zelf.status.store(Exhausted); + if pos == 0 { + // Can safely index directly. + return Ok(deque[pos].clone()); + } + Err(vm.new_stop_iteration()) + } + } } } }