diff --git a/crates/vm/src/stdlib/itertools.rs b/crates/vm/src/stdlib/itertools.rs index f09e556eb..d314779cf 100644 --- a/crates/vm/src/stdlib/itertools.rs +++ b/crates/vm/src/stdlib/itertools.rs @@ -7,7 +7,7 @@ mod decl { VirtualMachine, builtins::{PyGenericAlias, PyInt, PyList, PyTuple, PyTupleRef, PyType, PyTypeRef, int}, class::PyClassDef, - common::lock::{PyMutex, PyRwLock, PyRwLockWriteGuard}, + common::lock::{PyMutex, PyRwLock}, convert::ToPyObject, function::{ArgIndex, FuncArgs, NameIterables, PosArgs}, protocol::{PyIter, PyIterReturn, PyNumber}, @@ -1216,12 +1216,10 @@ mod decl { #[pyclass(name = "product", traverse)] #[derive(Debug, PyPayload)] struct PyItertoolsProduct { - pools: Vec>, + pools: Vec, #[pytraverse(skip)] idxs: PyRwLock>, #[pytraverse(skip)] - cur: AtomicCell, - #[pytraverse(skip)] stop: AtomicCell, } @@ -1256,12 +1254,19 @@ mod decl { .filter(|n| *n <= isize::MAX as usize / size_of::()) .ok_or_else(|| vm.new_overflow_error("repeat argument too large"))?; - let mut single: Vec> = Vec::new(); - for arg in iterables.iter() { - single.push(arg.try_to_value(vm)?); + let mut single = Vec::new(); + if repeat != 0 { + for arg in iterables.iter() { + let pool = if let Some(tuple) = arg.downcast_ref_if_exact::(vm) { + tuple.to_owned() + } else { + PyTuple::new_ref(arg.try_to_value(vm)?, &vm.ctx) + }; + single.push(pool); + } } - let mut pools: Vec> = Vec::new(); + let mut pools: Vec = Vec::new(); pools .try_reserve_exact(npools) .map_err(|_| vm.no_memory_error())?; @@ -1276,46 +1281,34 @@ mod decl { .map_err(|_| vm.no_memory_error())?; idxs.resize(npools, 0); - let l = pools.len(); + let empty = pools.iter().any(|pool| pool.is_empty()); Ok(Self { pools, idxs: PyRwLock::new(idxs), - cur: AtomicCell::new(l.wrapping_sub(1)), - stop: AtomicCell::new(false), + stop: AtomicCell::new(empty), }) } } #[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))] impl PyItertoolsProduct { - fn update_idxs(&self, mut idxs: PyRwLockWriteGuard<'_, Vec>) { - if idxs.is_empty() { - self.stop.store(true); - return; - } - - let cur = self.cur.load(); - let lst_idx = &self.pools[cur].len() - 1; - - if idxs[cur] == lst_idx { - if cur == 0 { - self.stop.store(true); + fn update_idxs(&self, idxs: &mut [usize]) { + for (index, pool) in idxs.iter_mut().zip(&self.pools).rev() { + *index += 1; + if *index < pool.len() { return; } - idxs[cur] = 0; - self.cur.fetch_sub(1); - self.update_idxs(idxs); - } else { - idxs[cur] += 1; - self.cur.store(idxs.len() - 1); + *index = 0; } + self.stop.store(true); } } impl SelfIter for PyItertoolsProduct {} impl IterNext for PyItertoolsProduct { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let mut idxs = zelf.idxs.write(); // stop signal if zelf.stop.load() { return Ok(PyIterReturn::StopIteration(None)); @@ -1323,13 +1316,6 @@ mod decl { let pools = &zelf.pools; - for p in pools { - if p.is_empty() { - return Ok(PyIterReturn::StopIteration(None)); - } - } - - let idxs = zelf.idxs.write(); let res = vm.ctx.new_tuple( pools .iter() @@ -1338,7 +1324,7 @@ mod decl { .collect(), ); - zelf.update_idxs(idxs); + zelf.update_idxs(&mut idxs); Ok(PyIterReturn::Return(res.into())) } diff --git a/extra_tests/snippets/stdlib_itertools.py b/extra_tests/snippets/stdlib_itertools.py index 2eb7ca7bb..16ea42f93 100644 --- a/extra_tests/snippets/stdlib_itertools.py +++ b/extra_tests/snippets/stdlib_itertools.py @@ -368,6 +368,11 @@ assert list(t[0]) == [1, 2, 3] assert list(t[0]) == [] # itertools.product + +product_source = iter([1]) +assert list(itertools.product(product_source, repeat=0)) == [()] +assert next(product_source) == 1 + it = itertools.product([1, 2], [3, 4]) assert (1, 3) == next(it) assert (1, 4) == next(it)