From a9d16e21b367e1606fe4937c9641cc0b01728c60 Mon Sep 17 00:00:00 2001 From: Alexander Scharinger Date: Sat, 11 Jun 2022 00:01:28 +0200 Subject: [PATCH 1/6] test: itertools.chain evaluate lazliy --- extra_tests/snippets/stdlib_itertools.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/extra_tests/snippets/stdlib_itertools.py b/extra_tests/snippets/stdlib_itertools.py index 05376d4b05a..bd5d201cf83 100644 --- a/extra_tests/snippets/stdlib_itertools.py +++ b/extra_tests/snippets/stdlib_itertools.py @@ -50,6 +50,12 @@ with assert_raises(TypeError): next(x) +# iterables are lazily evaluted +x = chain.from_iterable(itertools.repeat(range(2))) +assert next(x) == 0 +assert next(x) == 1 +assert next(x) == 0 +assert next(x) == 1 # itertools.count tests From 4c11720af09ee5f0abb90b30da93cfecf7db4d71 Mon Sep 17 00:00:00 2001 From: Alexander Scharinger Date: Sun, 12 Jun 2022 20:03:03 +0200 Subject: [PATCH 2/6] feat: itertools.chain evaluate lazily --- vm/src/stdlib/itertools.rs | 86 ++++++++++++++++++++------------------ 1 file changed, 46 insertions(+), 40 deletions(-) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index cc4eeb2d00c..f9999bd0262 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -7,7 +7,7 @@ mod decl { rc::PyRc, }; use crate::{ - builtins::{int, PyGenericAlias, PyInt, PyIntRef, PyTuple, PyTupleRef, PyTypeRef}, + builtins::{int, PyGenericAlias, PyInt, PyIntRef, PyList, PyTuple, PyTupleRef, PyTypeRef}, convert::ToPyObject, function::{ArgCallable, FuncArgs, OptionalArg, OptionalOption, PosArgs}, identifier, @@ -25,19 +25,18 @@ mod decl { #[pyclass(name = "chain")] #[derive(Debug, PyPayload)] struct PyItertoolsChain { - iterables: Vec, - cur_idx: AtomicCell, - cached_iter: PyRwLock>, + source: PyRwLock>, + active: PyRwLock>, } #[pyimpl(with(IterNext))] impl PyItertoolsChain { #[pyslot] fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult { + let args_list = PyList::from(args.args); PyItertoolsChain { - iterables: args.args, - cur_idx: AtomicCell::new(0), - cached_iter: PyRwLock::new(None), + source: PyRwLock::new(Some(args_list.to_pyobject(vm).get_iter(vm)?)), + active: PyRwLock::new(None), } .into_ref_with_type(vm, cls) .map(Into::into) @@ -46,13 +45,12 @@ mod decl { #[pyclassmethod] fn from_iterable( cls: PyTypeRef, - iterable: PyObjectRef, + source: PyObjectRef, vm: &VirtualMachine, ) -> PyResult> { PyItertoolsChain { - iterables: iterable.try_to_value(vm)?, - cur_idx: AtomicCell::new(0), - cached_iter: PyRwLock::new(None), + source: PyRwLock::new(Some(source.get_iter(vm)?)), + active: PyRwLock::new(None), } .into_ref_with_type(vm, cls) } @@ -65,37 +63,45 @@ mod decl { impl IterNextIterable for PyItertoolsChain {} impl IterNext for PyItertoolsChain { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - loop { - let pos = zelf.cur_idx.load(); - if pos >= zelf.iterables.len() { - break; - } - let cur_iter = if zelf.cached_iter.read().is_none() { - // We need to call "get_iter" outside of the lock. - let iter = zelf.iterables[pos].clone().get_iter(vm)?; - *zelf.cached_iter.write() = Some(iter.clone()); - iter - } else if let Some(cached_iter) = (*zelf.cached_iter.read()).clone() { - cached_iter - } else { - // Someone changed cached iter to None since we checked. - continue; - }; - - // We need to call "next" outside of the lock. - match cur_iter.next(vm) { - Ok(PyIterReturn::Return(ok)) => return Ok(PyIterReturn::Return(ok)), - Ok(PyIterReturn::StopIteration(_)) => { - zelf.cur_idx.fetch_add(1); - *zelf.cached_iter.write() = None; - } - Err(err) => { - return Err(err); + let next = || { + let source = zelf.source.read().clone(); + match source { + None => { + return Ok(PyIterReturn::StopIteration(None)); } + Some(source) => loop { + let active = zelf.active.read().clone(); + match active { + None => match source.next(vm) { + Ok(PyIterReturn::Return(ok)) => { + *zelf.active.write() = Some(ok.get_iter(vm)?); + } + Ok(PyIterReturn::StopIteration(_)) => { + return Ok(PyIterReturn::StopIteration(None)); + } + Err(err) => { + return Err(err); + } + }, + Some(active) => match active.next(vm) { + Ok(PyIterReturn::Return(ok)) => { + return Ok(PyIterReturn::Return(ok)); + } + Ok(PyIterReturn::StopIteration(_)) => { + *zelf.active.write() = None; + } + Err(err) => { + return Err(err); + } + }, + } + }, } - } - - Ok(PyIterReturn::StopIteration(None)) + }; + next().map_err(|err| { + *zelf.source.write() = None; + err + }) } } From 33bcc89e9447b3b58a72259ddb46f8fdb1fb78d1 Mon Sep 17 00:00:00 2001 From: Alexander Scharinger Date: Tue, 14 Jun 2022 08:19:43 +0200 Subject: [PATCH 3/6] test: itertools.chain stop iteration on error --- extra_tests/snippets/stdlib_itertools.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/extra_tests/snippets/stdlib_itertools.py b/extra_tests/snippets/stdlib_itertools.py index bd5d201cf83..a5b91d0cde9 100644 --- a/extra_tests/snippets/stdlib_itertools.py +++ b/extra_tests/snippets/stdlib_itertools.py @@ -57,6 +57,12 @@ assert next(x) == 0 assert next(x) == 1 +x = chain(1, [2]) +with assert_raises(TypeError): + next(x) +with assert_raises(StopIteration): + next(x) + # itertools.count tests # default arguments From ec18a12faf2283d718e1ca724999ef40b6ec86b1 Mon Sep 17 00:00:00 2001 From: Alexander Scharinger Date: Tue, 14 Jun 2022 18:57:00 +0200 Subject: [PATCH 4/6] refactor --- vm/src/stdlib/itertools.rs | 62 ++++++++++++++++++-------------------- 1 file changed, 30 insertions(+), 32 deletions(-) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index f9999bd0262..818d4116c6d 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -63,42 +63,40 @@ mod decl { impl IterNextIterable for PyItertoolsChain {} impl IterNext for PyItertoolsChain { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let next = || { - let source = zelf.source.read().clone(); - match source { - None => { - return Ok(PyIterReturn::StopIteration(None)); + let source = if let Some(source) = zelf.source.read().clone() { + source + } else { + return Ok(PyIterReturn::StopIteration(None)); + }; + let next = loop { + let option_active = zelf.active.read().clone(); + if let Some(active) = option_active { + match active.next(vm) { + Ok(PyIterReturn::Return(ok)) => { + break Ok(PyIterReturn::Return(ok)); + } + Ok(PyIterReturn::StopIteration(_)) => { + *zelf.active.write() = None; + } + Err(err) => { + break Err(err); + } } - Some(source) => loop { - let active = zelf.active.read().clone(); - match active { - None => match source.next(vm) { - Ok(PyIterReturn::Return(ok)) => { - *zelf.active.write() = Some(ok.get_iter(vm)?); - } - Ok(PyIterReturn::StopIteration(_)) => { - return Ok(PyIterReturn::StopIteration(None)); - } - Err(err) => { - return Err(err); - } - }, - Some(active) => match active.next(vm) { - Ok(PyIterReturn::Return(ok)) => { - return Ok(PyIterReturn::Return(ok)); - } - Ok(PyIterReturn::StopIteration(_)) => { - *zelf.active.write() = None; - } - Err(err) => { - return Err(err); - } - }, + } else { + match source.next(vm) { + Ok(PyIterReturn::Return(ok)) => { + *zelf.active.write() = Some(ok.get_iter(vm)?); } - }, + Ok(PyIterReturn::StopIteration(_)) => { + break Ok(PyIterReturn::StopIteration(None)); + } + Err(err) => { + break Err(err); + } + } } }; - next().map_err(|err| { + next.map_err(|err| { *zelf.source.write() = None; err }) From 76ef074b5a13b987d6ffe7f4d1921cabeb434825 Mon Sep 17 00:00:00 2001 From: Alexander Scharinger Date: Tue, 14 Jun 2022 18:59:27 +0200 Subject: [PATCH 5/6] Release source when PyIterReturn::StopIteration --- vm/src/stdlib/itertools.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 818d4116c6d..ecc1893fb98 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -96,10 +96,13 @@ mod decl { } } }; - next.map_err(|err| { - *zelf.source.write() = None; - err - }) + match next { + Err(_) | Ok(PyIterReturn::StopIteration(_)) => { + *zelf.source.write() = None; + } + _ => {} + }; + next } } From 480f5f1d259b6a25dd3f82f70a088ff3d6c5766b Mon Sep 17 00:00:00 2001 From: Alexander Scharinger Date: Tue, 14 Jun 2022 19:53:41 +0200 Subject: [PATCH 6/6] fix tests --- vm/src/stdlib/itertools.rs | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index ecc1893fb98..26f2710920f 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -69,8 +69,8 @@ mod decl { return Ok(PyIterReturn::StopIteration(None)); }; let next = loop { - let option_active = zelf.active.read().clone(); - if let Some(active) = option_active { + let maybe_active = zelf.active.read().clone(); + if let Some(active) = maybe_active { match active.next(vm) { Ok(PyIterReturn::Return(ok)) => { break Ok(PyIterReturn::Return(ok)); @@ -84,9 +84,14 @@ mod decl { } } else { match source.next(vm) { - Ok(PyIterReturn::Return(ok)) => { - *zelf.active.write() = Some(ok.get_iter(vm)?); - } + Ok(PyIterReturn::Return(ok)) => match ok.get_iter(vm) { + Ok(iter) => { + *zelf.active.write() = Some(iter); + } + Err(err) => { + break Err(err); + } + }, Ok(PyIterReturn::StopIteration(_)) => { break Ok(PyIterReturn::StopIteration(None)); }