From 350c4de4e586c94c82379f95b514fb441b46aa26 Mon Sep 17 00:00:00 2001 From: Kangzhi Shi Date: Fri, 1 Oct 2021 20:46:17 +0200 Subject: [PATCH] rebase to use PyIterReturn --- vm/src/builtins/bytearray.rs | 20 ++++--- vm/src/builtins/bytes.rs | 20 +++---- vm/src/builtins/dict.rs | 16 +++--- vm/src/builtins/enumerate.rs | 17 ++++-- vm/src/builtins/iter.rs | 102 +++++++++++++++-------------------- vm/src/builtins/list.rs | 40 +++++++------- vm/src/builtins/pystr.rs | 12 ++--- vm/src/builtins/set.rs | 8 +-- vm/src/builtins/tuple.rs | 21 ++++---- vm/src/stdlib/collections.rs | 60 +++++++++------------ 10 files changed, 145 insertions(+), 171 deletions(-) diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs index 2a365a22ce..e79fb91a58 100644 --- a/vm/src/builtins/bytearray.rs +++ b/vm/src/builtins/bytearray.rs @@ -751,7 +751,7 @@ impl PyByteArrayIterator { fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef { self.internal .lock() - .builtin_iter_reduce(|x| x.clone().into_object(), vm) + .builtins_iter_reduce(|x| x.clone().into_object(), vm) } #[pymethod(magic)] @@ -763,15 +763,13 @@ impl PyByteArrayIterator { } impl IteratorIterable for PyByteArrayIterator {} impl SlotIterator for PyByteArrayIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.internal.lock().next( - |bytearray, pos| { - let buf = bytearray.borrow_buf(); - buf.get(pos) - .ok_or_else(|| vm.new_stop_iteration()) - .map(|&x| vm.ctx.new_int(x)) - }, - vm, - ) + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + zelf.internal.lock().next(|bytearray, pos| { + let buf = bytearray.borrow_buf(); + Ok(match buf.get(pos) { + Some(&x) => PyIterReturn::Return(vm.ctx.new_int(x)), + None => PyIterReturn::StopIteration(None), + }) + }) } } diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs index 316ce8b086..38b7185780 100644 --- a/vm/src/builtins/bytes.rs +++ b/vm/src/builtins/bytes.rs @@ -605,7 +605,7 @@ impl PyBytesIterator { fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef { self.internal .lock() - .builtin_iter_reduce(|x| x.clone().into_object(), vm) + .builtins_iter_reduce(|x| x.clone().into_object(), vm) } #[pymethod(magic)] @@ -617,17 +617,13 @@ impl PyBytesIterator { } impl IteratorIterable for PyBytesIterator {} impl SlotIterator for PyBytesIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.internal.lock().next( - |bytes, pos| { - bytes - .as_bytes() - .get(pos) - .ok_or_else(|| vm.new_stop_iteration()) - .map(|&x| vm.ctx.new_int(x)) - }, - vm, - ) + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + zelf.internal.lock().next(|bytes, pos| { + Ok(match bytes.as_bytes().get(pos) { + Some(&x) => PyIterReturn::Return(vm.ctx.new_int(x)), + None => PyIterReturn::StopIteration(None), + }) + }) } } diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs index 26d56867fd..30087c0fbd 100644 --- a/vm/src/builtins/dict.rs +++ b/vm/src/builtins/dict.rs @@ -733,7 +733,7 @@ macro_rules! dict_iterator { impl IteratorIterable for $iter_name {} impl SlotIterator for $iter_name { #[allow(clippy::redundant_closure_call)] - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let mut internal = zelf.internal.lock(); if let IterStatus::Active(dict) = &internal.status { if dict.entries.has_changed_size(&zelf.size) { @@ -745,15 +745,15 @@ macro_rules! dict_iterator { match dict.entries.next_entry(internal.position) { Some((position, key, value)) => { internal.position = position; - Ok(($result_fn)(vm, key, value)) + Ok(PyIterReturn::Return(($result_fn)(vm, key, value))) } None => { internal.status = IterStatus::Exhausted; - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } else { - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } @@ -793,7 +793,7 @@ macro_rules! dict_iterator { impl IteratorIterable for $reverse_iter_name {} impl SlotIterator for $reverse_iter_name { #[allow(clippy::redundant_closure_call)] - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let mut internal = zelf.internal.lock(); if let IterStatus::Active(dict) = &internal.status { if dict.entries.has_changed_size(&zelf.size) { @@ -809,15 +809,15 @@ macro_rules! dict_iterator { } else { internal.position = position; } - Ok(($result_fn)(vm, key, value)) + Ok(PyIterReturn::Return(($result_fn)(vm, key, value))) } None => { internal.status = IterStatus::Exhausted; - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } else { - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } diff --git a/vm/src/builtins/enumerate.rs b/vm/src/builtins/enumerate.rs index 596d228d01..4922e491d9 100644 --- a/vm/src/builtins/enumerate.rs +++ b/vm/src/builtins/enumerate.rs @@ -1,5 +1,6 @@ use super::{IterStatus, PositionIterInternal, PyIntRef, PyTypeRef}; use crate::common::lock::{PyMutex, PyRwLock}; +use crate::TypeProtocol; use crate::{ function::OptionalArg, protocol::{PyIter, PyIterReturn}, @@ -105,16 +106,26 @@ impl PyReverseSequenceIterator { fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef { self.internal .lock() - .builtin_reversed_reduce(|x| x.clone(), vm) + .builtins_reversed_reduce(|x| x.clone(), vm) } } impl IteratorIterable for PyReverseSequenceIterator {} impl SlotIterator for PyReverseSequenceIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { zelf.internal .lock() - .rev_next(|obj, pos| obj.get_item(pos, vm), vm) + .rev_next(|obj, pos| match obj.get_item(pos, vm) { + Ok(ret) => Ok(PyIterReturn::Return(ret)), + Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => { + Ok(PyIterReturn::StopIteration(None)) + } + Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => { + let args = e.get_arg(0); + Ok(PyIterReturn::StopIteration(args)) + } + Err(e) => Err(e), + }) } } diff --git a/vm/src/builtins/iter.rs b/vm/src/builtins/iter.rs index b285dbd820..de39b23e35 100644 --- a/vm/src/builtins/iter.rs +++ b/vm/src/builtins/iter.rs @@ -71,7 +71,7 @@ impl PositionIterInternal { } } - pub fn builtin_iter_reduce(&self, f: F, vm: &VirtualMachine) -> PyObjectRef + pub fn builtins_iter_reduce(&self, f: F, vm: &VirtualMachine) -> PyObjectRef where F: FnOnce(&T) -> PyObjectRef, { @@ -79,7 +79,7 @@ impl PositionIterInternal { self._reduce(iter, f, vm) } - pub fn builtin_reversed_reduce(&self, f: F, vm: &VirtualMachine) -> PyObjectRef + pub fn builtins_reversed_reduce(&self, f: F, vm: &VirtualMachine) -> PyObjectRef where F: FnOnce(&T) -> PyObjectRef, { @@ -87,66 +87,42 @@ impl PositionIterInternal { self._reduce(reversed, f, vm) } - pub fn next(&mut self, f: F, vm: &VirtualMachine) -> PyResult + fn _next(&mut self, f: F, op: OP) -> PyResult where - F: FnOnce(&T, usize) -> PyResult, + F: FnOnce(&T, usize) -> PyResult, + OP: FnOnce(&mut Self), { if let IterStatus::Active(obj) = &self.status { - match f(obj, self.position) { - Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => { - self.status = IterStatus::Exhausted; - Err(e) - } - Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => { - self.status = IterStatus::Exhausted; - Err(vm.new_stop_iteration()) - } - Err(e) if e.isinstance(&vm.ctx.exceptions.runtime_error) => { - self.status = IterStatus::Exhausted; - Err(e) - } - Err(e) => Err(e), - Ok(ret) => { - self.position += 1; - Ok(ret) - } + let ret = f(obj, self.position); + if let Ok(PyIterReturn::Return(_)) = ret { + op(self); + } else { + self.status = IterStatus::Exhausted; } + ret } else { - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } - pub fn rev_next(&mut self, f: F, vm: &VirtualMachine) -> PyResult + pub fn next(&mut self, f: F) -> PyResult where - F: FnOnce(&T, usize) -> PyResult, + F: FnOnce(&T, usize) -> PyResult, { - if let IterStatus::Active(obj) = &self.status { - match f(obj, self.position) { - Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => { - self.status = IterStatus::Exhausted; - Err(e) - } - Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => { - self.status = IterStatus::Exhausted; - Err(vm.new_stop_iteration()) - } - Err(e) if e.isinstance(&vm.ctx.exceptions.runtime_error) => { - self.status = IterStatus::Exhausted; - Err(e) - } - Err(e) => Err(e), - Ok(ret) => { - if self.position == 0 { - self.status = IterStatus::Exhausted; - } else { - self.position -= 1; - } - Ok(ret) - } + self._next(f, |zelf| zelf.position += 1) + } + + pub fn rev_next(&mut self, f: F) -> PyResult + where + F: FnOnce(&T, usize) -> PyResult, + { + self._next(f, |zelf| { + if zelf.position == 0 { + zelf.status = IterStatus::Exhausted; + } else { + zelf.position -= 1; } - } else { - Err(vm.new_stop_iteration()) - } + }) } pub fn length_hint(&self, f: F) -> usize @@ -221,7 +197,7 @@ impl PySequenceIterator { #[pymethod(magic)] fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef { - self.internal.lock().builtin_iter_reduce(|x| x.clone(), vm) + self.internal.lock().builtins_iter_reduce(|x| x.clone(), vm) } #[pymethod(magic)] @@ -232,10 +208,20 @@ impl PySequenceIterator { impl IteratorIterable for PySequenceIterator {} impl SlotIterator for PySequenceIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { zelf.internal .lock() - .next(|obj, pos| obj.get_item(pos, vm), vm) + .next(|obj, pos| match obj.get_item(pos, vm) { + Ok(ret) => Ok(PyIterReturn::Return(ret)), + Err(e) if e.isinstance(&vm.ctx.exceptions.index_error) => { + Ok(PyIterReturn::StopIteration(None)) + } + Err(e) if e.isinstance(&vm.ctx.exceptions.stop_iteration) => { + let args = e.get_arg(0); + Ok(PyIterReturn::StopIteration(args)) + } + Err(e) => Err(e), + }) } } @@ -264,18 +250,18 @@ impl PyCallableIterator { impl IteratorIterable for PyCallableIterator {} impl SlotIterator for PyCallableIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let status = zelf.status.upgradable_read(); if let IterStatus::Active(callable) = &*status { let ret = callable.invoke((), vm)?; if vm.bool_eq(&ret, &zelf.sentinel)? { *PyRwLockUpgradableReadGuard::upgrade(status) = IterStatus::Exhausted; - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } else { - Ok(ret) + Ok(PyIterReturn::Return(ret)) } } else { - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } diff --git a/vm/src/builtins/list.rs b/vm/src/builtins/list.rs index 8c7ae63593..05a99fdd60 100644 --- a/vm/src/builtins/list.rs +++ b/vm/src/builtins/list.rs @@ -489,22 +489,20 @@ impl PyListIterator { fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef { self.internal .lock() - .builtin_iter_reduce(|x| x.clone().into_object(), vm) + .builtins_iter_reduce(|x| x.clone().into_object(), vm) } } impl IteratorIterable for PyListIterator {} impl SlotIterator for PyListIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.internal.lock().next( - |list, pos| { - let vec = list.borrow_vec(); - vec.get(pos) - .ok_or_else(|| vm.new_stop_iteration()) - .map(|x| x.clone()) - }, - vm, - ) + fn next(zelf: &PyRef, _vm: &VirtualMachine) -> PyResult { + zelf.internal.lock().next(|list, pos| { + let vec = list.borrow_vec(); + Ok(match vec.get(pos) { + Some(x) => PyIterReturn::Return(x.clone()), + None => PyIterReturn::StopIteration(None), + }) + }) } } @@ -538,22 +536,20 @@ impl PyListReverseIterator { fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef { self.internal .lock() - .builtin_reversed_reduce(|x| x.clone().into_object(), vm) + .builtins_reversed_reduce(|x| x.clone().into_object(), vm) } } impl IteratorIterable for PyListReverseIterator {} impl SlotIterator for PyListReverseIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.internal.lock().rev_next( - |list, pos| { - let vec = list.borrow_vec(); - vec.get(pos) - .ok_or_else(|| vm.new_stop_iteration()) - .map(|x| x.clone()) - }, - vm, - ) + fn next(zelf: &PyRef, _vm: &VirtualMachine) -> PyResult { + zelf.internal.lock().rev_next(|list, pos| { + let vec = list.borrow_vec(); + Ok(match vec.get(pos) { + Some(x) => PyIterReturn::Return(x.clone()), + None => PyIterReturn::StopIteration(None), + }) + }) } } diff --git a/vm/src/builtins/pystr.rs b/vm/src/builtins/pystr.rs index a6bf250bc8..26c1d45f3c 100644 --- a/vm/src/builtins/pystr.rs +++ b/vm/src/builtins/pystr.rs @@ -200,13 +200,13 @@ impl PyStrIterator { self.internal .lock() .0 - .builtin_iter_reduce(|x| x.clone().into_object(), vm) + .builtins_iter_reduce(|x| x.clone().into_object(), vm) } } impl IteratorIterable for PyStrIterator {} impl SlotIterator for PyStrIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let mut internal = zelf.internal.lock(); if let IterStatus::Active(s) = &internal.0.status { @@ -216,20 +216,18 @@ impl SlotIterator for PyStrIterator { if let Some((offset, ch)) = value.char_indices().nth(internal.0.position) { internal.0.position += 1; internal.1 = offset + ch.len_utf8(); - return Ok(ch.into_pyobject(vm)); + return Ok(PyIterReturn::Return(ch.into_pyobject(vm))); } } else if let Some(value) = value.get(internal.1..) { if let Some(ch) = value.chars().next() { internal.0.position += 1; internal.1 += ch.len_utf8(); - return Ok(ch.into_pyobject(vm)); + return Ok(PyIterReturn::Return(ch.into_pyobject(vm))); } } internal.0.status = Exhausted; - Err(vm.new_stop_iteration()) - } else { - Err(vm.new_stop_iteration()) } + Ok(PyIterReturn::StopIteration(None)) } } diff --git a/vm/src/builtins/set.rs b/vm/src/builtins/set.rs index d4ffc0c338..e57dac8fab 100644 --- a/vm/src/builtins/set.rs +++ b/vm/src/builtins/set.rs @@ -853,7 +853,7 @@ impl PySetIterator { impl IteratorIterable for PySetIterator {} impl SlotIterator for PySetIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { let mut internal = zelf.internal.lock(); if let IterStatus::Active(dict) = &internal.status { if dict.has_changed_size(&zelf.size) { @@ -863,15 +863,15 @@ impl SlotIterator for PySetIterator { match dict.next_entry(internal.position) { Some((position, key, _)) => { internal.position = position; - Ok(key) + Ok(PyIterReturn::Return(key)) } None => { internal.status = IterStatus::Exhausted; - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } else { - Err(vm.new_stop_iteration()) + Ok(PyIterReturn::StopIteration(None)) } } } diff --git a/vm/src/builtins/tuple.rs b/vm/src/builtins/tuple.rs index e692d636c1..c99aabf7e0 100644 --- a/vm/src/builtins/tuple.rs +++ b/vm/src/builtins/tuple.rs @@ -342,23 +342,20 @@ impl PyTupleIterator { fn reduce(&self, vm: &VirtualMachine) -> PyObjectRef { self.internal .lock() - .builtin_iter_reduce(|x| x.clone().into_object(), vm) + .builtins_iter_reduce(|x| x.clone().into_object(), vm) } } impl IteratorIterable for PyTupleIterator {} impl SlotIterator for PyTupleIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.internal.lock().next( - |tuple, pos| { - tuple - .as_slice() - .get(pos) - .ok_or_else(|| vm.new_stop_iteration()) - .map(|x| x.clone()) - }, - vm, - ) + fn next(zelf: &PyRef, _vm: &VirtualMachine) -> PyResult { + zelf.internal.lock().next(|tuple, pos| { + Ok(if let Some(ret) = tuple.as_slice().get(pos) { + PyIterReturn::Return(ret.clone()) + } else { + PyIterReturn::StopIteration(None) + }) + }) } } diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs index 6444e8198b..d003c9c901 100644 --- a/vm/src/stdlib/collections.rs +++ b/vm/src/stdlib/collections.rs @@ -634,22 +634,17 @@ mod _collections { impl IteratorIterable for PyDequeIterator {} impl SlotIterator for PyDequeIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.internal.lock().next( - |deque, pos| { - if zelf.state != deque.state.load() { - return Err( - vm.new_runtime_error("Deque mutated during iteration".to_owned()) - ); - } - let deque = deque.borrow_deque(); - deque - .get(pos) - .ok_or_else(|| vm.new_stop_iteration()) - .map(|x| x.clone()) - }, - vm, - ) + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + zelf.internal.lock().next(|deque, pos| { + if zelf.state != deque.state.load() { + return Err(vm.new_runtime_error("Deque mutated during iteration".to_owned())); + } + let deque = deque.borrow_deque(); + Ok(match deque.get(pos) { + Some(x) => PyIterReturn::Return(x.clone()), + None => PyIterReturn::StopIteration(None), + }) + }) } } @@ -706,26 +701,23 @@ mod _collections { impl IteratorIterable for PyReverseDequeIterator {} impl SlotIterator for PyReverseDequeIterator { - fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { - zelf.internal.lock().next( - |deque, pos| { - if deque.state.load() != zelf.state { - return Err( - vm.new_runtime_error("Deque mutated during iteration".to_owned()) - ); - } - let deque = deque.borrow_deque(); - let pos = deque + fn next(zelf: &PyRef, vm: &VirtualMachine) -> PyResult { + zelf.internal.lock().next(|deque, pos| { + if deque.state.load() != zelf.state { + return Err(vm.new_runtime_error("Deque mutated during iteration".to_owned())); + } + let deque = deque.borrow_deque(); + Ok( + match deque .len() .checked_sub(pos + 1) - .ok_or_else(|| vm.new_stop_iteration())?; - deque - .get(pos) - .ok_or_else(|| vm.new_stop_iteration()) - .map(|x| x.clone()) - }, - vm, - ) + .and_then(|pos| deque.get(pos)) + { + Some(x) => PyIterReturn::Return(x.clone()), + None => PyIterReturn::StopIteration(None), + }, + ) + }) } } }