diff --git a/crates/monty-proto/src/python/convert.rs b/crates/monty-proto/src/python/convert.rs index aeb604d39..587f61f14 100644 --- a/crates/monty-proto/src/python/convert.rs +++ b/crates/monty-proto/src/python/convert.rs @@ -237,6 +237,10 @@ fn round_trip_type_table(py: Python<'_>) -> PyResult<&'static Vec<(Py, Mo MontyType::ItertoolsIslice, MontyType::ItertoolsChain, MontyType::ItertoolsCycle, + MontyType::ItertoolsTakeWhile, + MontyType::ItertoolsDropWhile, + MontyType::ItertoolsFilterFalse, + MontyType::ItertoolsStarMap, MontyType::Tuple, MontyType::Dict, MontyType::Set, @@ -459,6 +463,10 @@ fn type_object_to_py(py: Python<'_>, t: MontyType) -> PyResult> { MontyType::ItertoolsIslice => cached!("itertools", "islice"), MontyType::ItertoolsChain => cached!("itertools", "chain"), MontyType::ItertoolsCycle => cached!("itertools", "cycle"), + MontyType::ItertoolsTakeWhile => cached!("itertools", "takewhile"), + MontyType::ItertoolsDropWhile => cached!("itertools", "dropwhile"), + MontyType::ItertoolsFilterFalse => cached!("itertools", "filterfalse"), + MontyType::ItertoolsStarMap => cached!("itertools", "starmap"), // Consistent with the Path *instance* arm, which marshals as PurePosixPath // and is instantiable on every host OS (unlike PosixPath on Windows). MontyType::Path => get_pure_posix_path(py).map(|b| b.clone().unbind()), diff --git a/crates/monty-types/src/object.rs b/crates/monty-types/src/object.rs index 9c9af8fbf..a9b60ead8 100644 --- a/crates/monty-types/src/object.rs +++ b/crates/monty-types/src/object.rs @@ -830,6 +830,14 @@ pub enum MontyType { ItertoolsCycle, #[strum(serialize = "NotImplementedType")] NotImplementedType, + #[strum(serialize = "itertools.takewhile")] + ItertoolsTakeWhile, + #[strum(serialize = "itertools.dropwhile")] + ItertoolsDropWhile, + #[strum(serialize = "itertools.filterfalse")] + ItertoolsFilterFalse, + #[strum(serialize = "itertools.starmap")] + ItertoolsStarMap, } impl fmt::Display for MontyType { diff --git a/crates/monty-typeshed/custom/itertools.pyi b/crates/monty-typeshed/custom/itertools.pyi index 3dea42066..c0b9b7005 100644 --- a/crates/monty-typeshed/custom/itertools.pyi +++ b/crates/monty-typeshed/custom/itertools.pyi @@ -10,8 +10,8 @@ # because Monty's `count()` rejects anything else, where CPython accepts any # number protocol. -from collections.abc import Iterable -from typing import Generic, TypeAlias, TypeVar, overload +from collections.abc import Callable, Iterable +from typing import Any, Generic, TypeAlias, TypeVar, overload from typing_extensions import Self @@ -65,3 +65,25 @@ class cycle(Generic[_T]): def __new__(cls, iterable: Iterable[_T], /) -> Self: ... def __next__(self) -> _T: ... def __iter__(self) -> Self: ... + +class takewhile(Generic[_T]): + def __new__(cls, predicate: Callable[[_T], object], iterable: Iterable[_T], /) -> Self: ... + def __next__(self) -> _T: ... + def __iter__(self) -> Self: ... + +class dropwhile(Generic[_T]): + def __new__(cls, predicate: Callable[[_T], object], iterable: Iterable[_T], /) -> Self: ... + def __next__(self) -> _T: ... + def __iter__(self) -> Self: ... + +class filterfalse(Generic[_T]): + def __new__(cls, predicate: Callable[[_T], object] | None, iterable: Iterable[_T], /) -> Self: ... + def __next__(self) -> _T: ... + def __iter__(self) -> Self: ... + +# Upstream typeshed overloads `starmap` per arity; Monty spreads each item +# dynamically, so one `Callable[..., _T]` covers every shape it accepts. +class starmap(Generic[_T]): + def __new__(cls, function: Callable[..., _T], iterable: Iterable[Iterable[Any]], /) -> Self: ... + def __next__(self) -> _T: ... + def __iter__(self) -> Self: ... diff --git a/crates/monty-typeshed/vendor/typeshed/stdlib/itertools.pyi b/crates/monty-typeshed/vendor/typeshed/stdlib/itertools.pyi index 3dea42066..c0b9b7005 100644 --- a/crates/monty-typeshed/vendor/typeshed/stdlib/itertools.pyi +++ b/crates/monty-typeshed/vendor/typeshed/stdlib/itertools.pyi @@ -10,8 +10,8 @@ # because Monty's `count()` rejects anything else, where CPython accepts any # number protocol. -from collections.abc import Iterable -from typing import Generic, TypeAlias, TypeVar, overload +from collections.abc import Callable, Iterable +from typing import Any, Generic, TypeAlias, TypeVar, overload from typing_extensions import Self @@ -65,3 +65,25 @@ class cycle(Generic[_T]): def __new__(cls, iterable: Iterable[_T], /) -> Self: ... def __next__(self) -> _T: ... def __iter__(self) -> Self: ... + +class takewhile(Generic[_T]): + def __new__(cls, predicate: Callable[[_T], object], iterable: Iterable[_T], /) -> Self: ... + def __next__(self) -> _T: ... + def __iter__(self) -> Self: ... + +class dropwhile(Generic[_T]): + def __new__(cls, predicate: Callable[[_T], object], iterable: Iterable[_T], /) -> Self: ... + def __next__(self) -> _T: ... + def __iter__(self) -> Self: ... + +class filterfalse(Generic[_T]): + def __new__(cls, predicate: Callable[[_T], object] | None, iterable: Iterable[_T], /) -> Self: ... + def __next__(self) -> _T: ... + def __iter__(self) -> Self: ... + +# Upstream typeshed overloads `starmap` per arity; Monty spreads each item +# dynamically, so one `Callable[..., _T]` covers every shape it accepts. +class starmap(Generic[_T]): + def __new__(cls, function: Callable[..., _T], iterable: Iterable[Iterable[Any]], /) -> Self: ... + def __next__(self) -> _T: ... + def __iter__(self) -> Self: ... diff --git a/crates/monty/src/builtins/filter.rs b/crates/monty/src/builtins/filter.rs index 580809205..16d9ddbaf 100644 --- a/crates/monty/src/builtins/filter.rs +++ b/crates/monty/src/builtins/filter.rs @@ -5,7 +5,7 @@ //! - `None` as predicate (filters falsy values) //! - Builtin functions (len, abs, etc.) //! - Type constructors (int, str, float, etc.) -//! - User-defined functions (via `vm.evaluate_function`) +//! - User-defined functions (via `call_predicate`) use crate::{ args::ArgValues, @@ -13,6 +13,7 @@ use crate::{ defer_drop, exception_private::RunResult, heap::{DropGuard, HeapData}, + predicate::call_predicate, types::{List, PyTrait}, value::Value, }; @@ -48,12 +49,7 @@ pub fn builtin_filter(vm: &mut VM<'_>, args: ArgValues) -> RunResult { // No predicate - use truthiness of element item.py_bool(vm)? } else { - // Clone for predicate call - the clone is consumed by evaluate_function - let item_for_predicate = item.clone_with_heap(vm); - let result = vm.evaluate_function("filter()", function, ArgValues::One(item_for_predicate))?; - let is_truthy = result.py_bool(vm); - result.drop_with(vm); - is_truthy? + call_predicate(function, item, "filter()", vm)? }; if should_include { diff --git a/crates/monty/src/dump_format.rs b/crates/monty/src/dump_format.rs index b88e9a110..cbb57ef18 100644 --- a/crates/monty/src/dump_format.rs +++ b/crates/monty/src/dump_format.rs @@ -196,7 +196,7 @@ mod tests { ); assert_eq!( static_strings_fingerprint(), - 0x239c_c721_30be_eba3, + 0xd864_f832_cad7_dea1, "static strings changed for dump version {DUMP_VERSION}" ); assert_eq!( diff --git a/crates/monty/src/intern.rs b/crates/monty/src/intern.rs index 4f133925b..cb2b13480 100644 --- a/crates/monty/src/intern.rs +++ b/crates/monty/src/intern.rs @@ -935,6 +935,14 @@ pub enum StaticStrings { /// Python's `NotImplemented` singleton representation. #[strum(serialize = "NotImplemented")] NotImplementedRepr, + /// `itertools.takewhile()` function. + Takewhile, + /// `itertools.dropwhile()` function. + Dropwhile, + /// `itertools.filterfalse()` function. + Filterfalse, + /// `itertools.starmap()` function. + Starmap, } /// Computes an FNV-1a hash over static-string identities and serialization. diff --git a/crates/monty/src/lib.rs b/crates/monty/src/lib.rs index 81afec931..1187b03b9 100644 --- a/crates/monty/src/lib.rs +++ b/crates/monty/src/lib.rs @@ -23,6 +23,7 @@ mod namespace; mod object_bridge; mod os_dispatch; mod parse; +mod predicate; mod prepare; mod repl; mod resource_checks; diff --git a/crates/monty/src/modules/itertools.rs b/crates/monty/src/modules/itertools.rs index 897662c60..16f3d9bfe 100644 --- a/crates/monty/src/modules/itertools.rs +++ b/crates/monty/src/modules/itertools.rs @@ -14,7 +14,9 @@ use crate::{ modules::ModuleFunctions, types::{ ItertoolsIter, Module, Type, - itertools::{Chain, Compress, Count, Cycle, Islice, Pairwise, Repeat}, + itertools::{ + Chain, Compress, Count, Cycle, DropWhile, FilterFalse, Islice, Pairwise, Repeat, StarMap, TakeWhile, + }, }, value::Value, }; @@ -30,6 +32,10 @@ pub(crate) enum ItertoolsFunctions { Islice, Chain, Cycle, + Takewhile, + Dropwhile, + Filterfalse, + Starmap, } /// Static mapping of attribute names to functions for module creation. @@ -41,6 +47,10 @@ const ITERTOOLS_FUNCTIONS: &[(StaticStrings, ItertoolsFunctions)] = &[ (StaticStrings::Islice, ItertoolsFunctions::Islice), (StaticStrings::Chain, ItertoolsFunctions::Chain), (StaticStrings::Cycle, ItertoolsFunctions::Cycle), + (StaticStrings::Takewhile, ItertoolsFunctions::Takewhile), + (StaticStrings::Dropwhile, ItertoolsFunctions::Dropwhile), + (StaticStrings::Filterfalse, ItertoolsFunctions::Filterfalse), + (StaticStrings::Starmap, ItertoolsFunctions::Starmap), ]; /// Creates the `itertools` module on the heap. @@ -67,6 +77,10 @@ pub(super) fn call(vm: &mut VM<'_>, function: ItertoolsFunctions, args: ArgValue ItertoolsFunctions::Islice => call_islice(vm, args), ItertoolsFunctions::Chain => call_chain(vm, args), ItertoolsFunctions::Cycle => call_cycle(vm, args), + ItertoolsFunctions::Takewhile => call_takewhile(vm, args), + ItertoolsFunctions::Dropwhile => call_dropwhile(vm, args), + ItertoolsFunctions::Filterfalse => call_filterfalse(vm, args), + ItertoolsFunctions::Starmap => call_starmap(vm, args), } } @@ -359,3 +373,72 @@ fn call_cycle(vm: &mut VM<'_>, args: ArgValues) -> RunResult { let iter = ItertoolsIter::Cycle(Cycle::new(source)); Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter)))) } + +/// Argument shape shared by `takewhile`, `dropwhile`, `filterfalse` and +/// `starmap`. +/// +/// All four are `PyArg_UnpackTuple(args, name, 2, 2, ...)` in CPython, so both +/// slots are positional-only, arity reads `takewhile expected 2 arguments, got +/// 1`, and keywords are rejected wholesale. The macro embeds the name, hence +/// one struct per callable rather than one shared struct. +macro_rules! callable_and_iterable_args { + ($struct_name:ident, $py_name:literal) => { + #[derive(FromArgs)] + #[from_args(name = $py_name, style = unpack)] + struct $struct_name { + #[from_args(pos_only)] + callable: Value, + #[from_args(pos_only)] + iterable: Value, + } + }; +} + +callable_and_iterable_args!(TakeWhileArgs, "takewhile"); +callable_and_iterable_args!(DropWhileArgs, "dropwhile"); +callable_and_iterable_args!(FilterFalseArgs, "filterfalse"); +callable_and_iterable_args!(StarMapArgs, "starmap"); + +/// `itertools.takewhile(predicate, iterable)` — the leading passing run. +fn call_takewhile(vm: &mut VM<'_>, args: ArgValues) -> RunResult { + let TakeWhileArgs { callable, iterable } = TakeWhileArgs::from_args(args, vm)?; + let (predicate, source) = resolve_source(callable, iterable, vm)?; + let iter = ItertoolsIter::TakeWhile(TakeWhile::new(predicate, source)); + Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter)))) +} + +/// `itertools.dropwhile(predicate, iterable)` — everything past that run. +fn call_dropwhile(vm: &mut VM<'_>, args: ArgValues) -> RunResult { + let DropWhileArgs { callable, iterable } = DropWhileArgs::from_args(args, vm)?; + let (predicate, source) = resolve_source(callable, iterable, vm)?; + let iter = ItertoolsIter::DropWhile(DropWhile::new(predicate, source)); + Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter)))) +} + +/// `itertools.filterfalse(predicate, iterable)` — the items it rejects. +fn call_filterfalse(vm: &mut VM<'_>, args: ArgValues) -> RunResult { + let FilterFalseArgs { callable, iterable } = FilterFalseArgs::from_args(args, vm)?; + let (predicate, source) = resolve_source(callable, iterable, vm)?; + let iter = ItertoolsIter::FilterFalse(FilterFalse::new(predicate, source)); + Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter)))) +} + +/// `itertools.starmap(function, iterable)` — each item spread as arguments. +fn call_starmap(vm: &mut VM<'_>, args: ArgValues) -> RunResult { + let StarMapArgs { callable, iterable } = StarMapArgs::from_args(args, vm)?; + let (function, source) = resolve_source(callable, iterable, vm)?; + let iter = ItertoolsIter::StarMap(StarMap::new(function, source)); + Ok(Value::Ref(vm.heap.allocate(HeapData::Itertools(iter)))) +} + +/// Resolves the iterable while keeping the callable safe from the error path. +/// +/// CPython resolves eagerly for all four, so a non-iterable raises here rather +/// than on the first `next()`. The callable itself is never type-checked: a +/// non-callable is only discovered when the adaptor first applies it. +fn resolve_source(callable: Value, iterable: Value, vm: &mut VM<'_>) -> RunResult<(Value, Value)> { + let mut guard = DropGuard::new(callable, vm); + let source = iterable.into_py_iter(guard.ctx())?; + let (callable, _) = guard.into_parts(); + Ok((callable, source)) +} diff --git a/crates/monty/src/object_bridge.rs b/crates/monty/src/object_bridge.rs index f13f4e864..fc564c91e 100644 --- a/crates/monty/src/object_bridge.rs +++ b/crates/monty/src/object_bridge.rs @@ -570,6 +570,10 @@ impl MontyTypeExt for MontyType { Self::ItertoolsIslice => Some(Type::ItertoolsIslice), Self::ItertoolsChain => Some(Type::ItertoolsChain), Self::ItertoolsCycle => Some(Type::ItertoolsCycle), + Self::ItertoolsTakeWhile => Some(Type::ItertoolsTakeWhile), + Self::ItertoolsDropWhile => Some(Type::ItertoolsDropWhile), + Self::ItertoolsFilterFalse => Some(Type::ItertoolsFilterFalse), + Self::ItertoolsStarMap => Some(Type::ItertoolsStarMap), Self::ItertoolsCount => Some(Type::ItertoolsCount), Self::ItertoolsRepeat => Some(Type::ItertoolsRepeat), Self::Tuple => Some(Type::Tuple), @@ -647,6 +651,10 @@ impl MontyTypeExt for MontyType { Type::ItertoolsIslice => Self::ItertoolsIslice, Type::ItertoolsChain => Self::ItertoolsChain, Type::ItertoolsCycle => Self::ItertoolsCycle, + Type::ItertoolsTakeWhile => Self::ItertoolsTakeWhile, + Type::ItertoolsDropWhile => Self::ItertoolsDropWhile, + Type::ItertoolsFilterFalse => Self::ItertoolsFilterFalse, + Type::ItertoolsStarMap => Self::ItertoolsStarMap, Type::ItertoolsCount => Self::ItertoolsCount, Type::ItertoolsRepeat => Self::ItertoolsRepeat, Type::Tuple => Self::Tuple, diff --git a/crates/monty/src/predicate.rs b/crates/monty/src/predicate.rs new file mode 100644 index 000000000..7e942966f --- /dev/null +++ b/crates/monty/src/predicate.rs @@ -0,0 +1,21 @@ +//! Calling a user-supplied predicate and taking its truthiness. +//! +//! Shared by `filter` and the `itertools` adaptors `takewhile`, `dropwhile` +//! and `filterfalse`, which differ only in what they do with the answer. + +use crate::{args::ArgValues, bytecode::VM, defer_drop, exception_private::RunResult, types::PyTrait, value::Value}; + +/// Applies `predicate` to `item` and returns its truthiness. +/// +/// `item` is only borrowed — the call gets its own reference, so the caller can +/// still yield the item afterwards. Goes through `evaluate_function`, which +/// runs a defined function's frame to completion, so a predicate that suspends +/// (an external function, an `os` call) is rejected rather than paused; `ctx` +/// names the caller in that error. +pub(crate) fn call_predicate(predicate: &Value, item: &Value, ctx: &'static str, vm: &mut VM<'_>) -> RunResult { + let arg = item.clone_with_heap(vm.heap); + let result = vm.evaluate_function(ctx, predicate, ArgValues::One(arg))?; + // Guarded because `py_bool` can raise through a user `__bool__`. + defer_drop!(result, vm); + result.py_bool(vm) +} diff --git a/crates/monty/src/types/iter.rs b/crates/monty/src/types/iter.rs index a31856f41..6409c0368 100644 --- a/crates/monty/src/types/iter.rs +++ b/crates/monty/src/types/iter.rs @@ -85,7 +85,12 @@ impl<'h, I: CollectIter<'h>> Iterator for HeapedIterator<'_, 'h, I> { Ok(Some(value)) => { self.yielded += 1; let estimated = self.yielded.saturating_mul(VALUE_SIZE); - match check_estimated_size(estimated, &self.vm.heap.tracker) { + // Size alone does not bound a source whose items are cheap or + // interned, and the drain reaches no VM dispatch checkpoint of + // its own, so `max_duration` needs its own poll here. + let checked = check_estimated_size(estimated, &self.vm.heap.tracker) + .and_then(|()| self.vm.heap.tracker.check_time_every(self.yielded)); + match checked { Ok(()) => Some(value), Err(error) => { *self.error = Some(error.into()); diff --git a/crates/monty/src/types/itertools/chain.rs b/crates/monty/src/types/itertools/chain.rs index c1a27b444..00a670515 100644 --- a/crates/monty/src/types/itertools/chain.rs +++ b/crates/monty/src/types/itertools/chain.rs @@ -65,7 +65,14 @@ impl Chain { /// Drains the current source, then resolves the next one, until all are spent. pub(super) fn next<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) -> RunResult> { + let mut steps = 0usize; loop { + // Native loop: the VM's dispatch checkpoint is per-`run()`, so a + // discarding pass over an infinite source reaches none. Poll the + // tracker so `max_duration` still bites (see `VM::run`'s + // `CHECK_INTERVAL`). + vm.heap.tracker.check_time_every(steps)?; + steps += 1; let ItertoolsIter::Chain(chain) = iter.get(vm.heap) else { unreachable!("dispatched on Kind::Chain") }; diff --git a/crates/monty/src/types/itertools/compress.rs b/crates/monty/src/types/itertools/compress.rs index 7eaa43ad0..715f97a10 100644 --- a/crates/monty/src/types/itertools/compress.rs +++ b/crates/monty/src/types/itertools/compress.rs @@ -52,7 +52,14 @@ impl Compress { /// Pulls one item from each side per step, yielding the datum when its selector /// is truthy and stopping as soon as *either* side runs out. pub(super) fn next<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) -> RunResult> { + let mut steps = 0usize; loop { + // Native loop: the VM's dispatch checkpoint is per-`run()`, so a + // discarding pass over an infinite source reaches none. Poll the + // tracker so `max_duration` still bites (see `VM::run`'s + // `CHECK_INTERVAL`). + vm.heap.tracker.check_time_every(steps)?; + steps += 1; let ItertoolsIter::Compress(compress) = iter.get(vm.heap) else { unreachable!("dispatched on Kind::Compress") }; diff --git a/crates/monty/src/types/itertools/dropwhile.rs b/crates/monty/src/types/itertools/dropwhile.rs new file mode 100644 index 000000000..b2f8ab006 --- /dev/null +++ b/crates/monty/src/types/itertools/dropwhile.rs @@ -0,0 +1,99 @@ +//! `itertools.dropwhile(predicate, iterable)` — everything after the leading run. + +use serde::{Deserialize, Serialize}; + +use crate::{ + bytecode::VM, + exception_private::RunResult, + heap::{HeapId, HeapRead}, + predicate::call_predicate, + types::itertools::{ + ItertoolsIter, + step::{next_item, next_tested}, + }, + value::Value, +}; + +/// Discards items while `predicate` holds, then yields the rest untested. +/// +/// `dropping` clears on the first item the predicate rejects, and that item is +/// the first one yielded; the predicate is never consulted again but stays +/// owned until destruction, as CPython holds `lz->func` for the iterator's +/// whole life. The source stays owned too: unlike `takewhile` this adaptor +/// never latches, so every later `next` drives it again. +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct DropWhile { + predicate: Value, + source: Value, + /// Whether the leading run is still being discarded. + dropping: bool, +} + +impl DropWhile { + /// Takes ownership of both, with `source` already resolved by `py_iter`. + pub(crate) fn new(predicate: Value, source: Value) -> Self { + Self { + predicate, + source, + dropping: true, + } + } + + /// Invokes `on_child` for each heap id this iterator owns (GC trace hook). + pub(crate) fn for_each_child_id(&self, mut on_child: impl FnMut(HeapId)) { + if let Value::Ref(id) = &self.predicate { + on_child(*id); + } + if let Value::Ref(id) = &self.source { + on_child(*id); + } + } + + /// Releases the refs this iterator owns (mirrors `for_each_child_id`). + pub(crate) fn py_dec_ref_ids(&mut self, stack: &mut Vec) { + self.predicate.py_dec_ref_ids(stack); + self.source.py_dec_ref_ids(stack); + } +} + +/// Skips items while the predicate holds; once it has failed, yields straight +/// through without consulting it again. +pub(super) fn next<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) -> RunResult> { + let mut steps = 0usize; + loop { + // Native loop: the VM's dispatch checkpoint is per-`run()`, so a + // discarding pass over an infinite source reaches none. Poll the + // tracker so `max_duration` still bites (see `VM::run`'s + // `CHECK_INTERVAL`). + vm.heap.tracker.check_time_every(steps)?; + steps += 1; + let ItertoolsIter::DropWhile(drop_while) = iter.get(vm.heap) else { + unreachable!("dispatched on Kind::DropWhile") + }; + let dropping = drop_while.dropping; + let source = drop_while.source.clone_with_heap(vm.heap); + // Past the leading run the predicate is never consulted again, so the + // item is yielded untested and the predicate is not even cloned. + if !dropping { + return next_item(source, vm); + } + let predicate = drop_while.predicate.clone_with_heap(vm.heap); + + let Some((item, accepted)) = next_tested(predicate, source, vm, |predicate, item, vm| { + call_predicate(predicate, item, "dropwhile()", vm) + })? + else { + return Ok(None); + }; + if !accepted { + let ItertoolsIter::DropWhile(drop_while) = iter.get_mut(vm.heap) else { + unreachable!("dispatched on Kind::DropWhile") + }; + // Retained, not released: only the flag says it is done with. + drop_while.dropping = false; + return Ok(Some(item)); + } + // Still in the leading run, so this item is discarded. + item.drop_with(vm); + } +} diff --git a/crates/monty/src/types/itertools/filterfalse.rs b/crates/monty/src/types/itertools/filterfalse.rs new file mode 100644 index 000000000..977246a26 --- /dev/null +++ b/crates/monty/src/types/itertools/filterfalse.rs @@ -0,0 +1,83 @@ +//! `itertools.filterfalse(predicate, iterable)` — the items a predicate rejects. + +use serde::{Deserialize, Serialize}; + +use crate::{ + bytecode::VM, + exception_private::RunResult, + heap::{HeapId, HeapRead}, + predicate::call_predicate, + types::{ + PyTrait, + itertools::{ItertoolsIter, step::next_tested}, + }, + value::Value, +}; + +/// Yields the items of `source` for which `predicate` is false. +/// +/// No spent flag: unlike the `while` adaptors this never stops early, so the +/// source's own exhaustion is the only end condition. +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct FilterFalse { + /// `Value::None` selects the truth test, as `filter(None, ...)` does. + predicate: Value, + source: Value, +} + +impl FilterFalse { + /// Takes ownership of both, with `source` already resolved by `py_iter`. + pub(crate) fn new(predicate: Value, source: Value) -> Self { + Self { predicate, source } + } + + /// Invokes `on_child` for each heap id this iterator owns (GC trace hook). + pub(crate) fn for_each_child_id(&self, mut on_child: impl FnMut(HeapId)) { + if let Value::Ref(id) = &self.predicate { + on_child(*id); + } + if let Value::Ref(id) = &self.source { + on_child(*id); + } + } + + /// Releases the refs this iterator owns (mirrors `for_each_child_id`). + pub(crate) fn py_dec_ref_ids(&mut self, stack: &mut Vec) { + self.predicate.py_dec_ref_ids(stack); + self.source.py_dec_ref_ids(stack); + } +} + +/// Pulls items until one the predicate rejects, which is the one yielded. +pub(super) fn next<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) -> RunResult> { + let mut steps = 0usize; + loop { + // Native loop: the VM's dispatch checkpoint is per-`run()`, so a + // discarding pass over an infinite source reaches none. Poll the + // tracker so `max_duration` still bites (see `VM::run`'s + // `CHECK_INTERVAL`). + vm.heap.tracker.check_time_every(steps)?; + steps += 1; + let ItertoolsIter::FilterFalse(filter) = iter.get(vm.heap) else { + unreachable!("dispatched on Kind::FilterFalse") + }; + let predicate = filter.predicate.clone_with_heap(vm.heap); + let source = filter.source.clone_with_heap(vm.heap); + + let Some((item, truthy)) = next_tested(predicate, source, vm, |predicate, item, vm| { + if matches!(predicate, Value::None) { + item.py_bool(vm) + } else { + call_predicate(predicate, item, "filterfalse()", vm) + } + })? + else { + return Ok(None); + }; + if !truthy { + return Ok(Some(item)); + } + // Accepted by the predicate, so this one is filtered out. + item.drop_with(vm); + } +} diff --git a/crates/monty/src/types/itertools/islice.rs b/crates/monty/src/types/itertools/islice.rs index 736f62500..bb5fd0c47 100644 --- a/crates/monty/src/types/itertools/islice.rs +++ b/crates/monty/src/types/itertools/islice.rs @@ -61,7 +61,14 @@ impl Islice { /// Discards items up to `next_index`, yields the one there, then advances by /// `step` — the shape of CPython's `islice_next`. pub(super) fn next<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) -> RunResult> { + let mut steps = 0usize; loop { + // Native loop: the VM's dispatch checkpoint is per-`run()`, so a + // discarding pass over an infinite source reaches none. Poll the + // tracker so `max_duration` still bites (see `VM::run`'s + // `CHECK_INTERVAL`). + vm.heap.tracker.check_time_every(steps)?; + steps += 1; let ItertoolsIter::Islice(islice) = iter.get(vm.heap) else { unreachable!("dispatched on Kind::Islice") }; diff --git a/crates/monty/src/types/itertools/mod.rs b/crates/monty/src/types/itertools/mod.rs index 9c369c467..07fd2b6cb 100644 --- a/crates/monty/src/types/itertools/mod.rs +++ b/crates/monty/src/types/itertools/mod.rs @@ -14,21 +14,33 @@ pub mod chain; pub mod compress; pub mod count; pub mod cycle; +pub mod dropwhile; +pub mod filterfalse; pub mod islice; pub mod pairwise; pub mod repeat; +pub mod starmap; +mod step; +pub mod takewhile; -use std::fmt::Write; +use std::{fmt::Write, mem}; pub(crate) use chain::Chain; pub(crate) use compress::Compress; pub(crate) use count::Count; pub(crate) use cycle::Cycle; +pub(crate) use dropwhile::DropWhile; +pub(crate) use filterfalse::FilterFalse; pub(crate) use islice::Islice; pub(crate) use pairwise::Pairwise; pub(crate) use repeat::Repeat; use serde::{Deserialize, Serialize}; +pub(crate) use starmap::StarMap; +pub(crate) use takewhile::TakeWhile; +// Only the 64-bit size budget below needs it. +#[cfg(target_pointer_width = "64")] +use crate::types::Dict; use crate::{ bytecode::VM, exception_private::RunResult, @@ -39,9 +51,10 @@ use crate::{ /// The state of one `itertools` iterator, whichever adaptor produced it. /// -/// Held inline: `HeapData` is 160 bytes and the widest adaptor needs 56. If one -/// ever exceeds that, box it in its variant here — never at the `HeapData` -/// boundary, where `heap_read_boxed` is only sound for reads. +/// Held inline, so this width is memcpy'd on every heap allocate and free along +/// with the rest of `HeapData` — which #636 shrank to 80 bytes, asserted in +/// `heap_data.rs`. The budget below keeps the family from becoming what sets +/// that size. #[derive(Debug, Serialize, Deserialize)] pub(crate) enum ItertoolsIter { Count(Count), @@ -51,8 +64,21 @@ pub(crate) enum ItertoolsIter { Islice(Islice), Chain(Chain), Cycle(Cycle), + TakeWhile(TakeWhile), + DropWhile(DropWhile), + FilterFalse(FilterFalse), + StarMap(StarMap), } +// `Dict` is the widest `HeapData` payload on 64-bit hosts, so it — not a +// literal — is the budget: staying under it keeps this family from setting +// `HeapData`'s size. Only there: on 32-bit (the wasm worker) `Dict` halves +// while the adaptors' `i64` fields do not, and other variants set the size. +// TODO: when this fails, box the offending variant (`GroupBy(Box)`), +// not the enum and not at the `HeapData` boundary. +#[cfg(target_pointer_width = "64")] +const _: () = assert!(mem::size_of::() <= mem::size_of::()); + /// Which adaptor an [`ItertoolsIter`] is, without borrowing it. /// /// `py_next` and friends need the variant to pick a per-type function, but they @@ -67,6 +93,10 @@ pub(crate) enum Kind { Islice, Chain, Cycle, + TakeWhile, + DropWhile, + FilterFalse, + StarMap, } impl ItertoolsIter { @@ -80,6 +110,10 @@ impl ItertoolsIter { Self::Islice(_) => Kind::Islice, Self::Chain(_) => Kind::Chain, Self::Cycle(_) => Kind::Cycle, + Self::TakeWhile(_) => Kind::TakeWhile, + Self::DropWhile(_) => Kind::DropWhile, + Self::FilterFalse(_) => Kind::FilterFalse, + Self::StarMap(_) => Kind::StarMap, } } @@ -93,6 +127,10 @@ impl ItertoolsIter { Self::Islice(_) => Type::ItertoolsIslice, Self::Chain(_) => Type::ItertoolsChain, Self::Cycle(_) => Type::ItertoolsCycle, + Self::TakeWhile(_) => Type::ItertoolsTakeWhile, + Self::DropWhile(_) => Type::ItertoolsDropWhile, + Self::FilterFalse(_) => Type::ItertoolsFilterFalse, + Self::StarMap(_) => Type::ItertoolsStarMap, } } @@ -108,7 +146,11 @@ impl ItertoolsIter { | Self::Compress(_) | Self::Islice(_) | Self::Chain(_) - | Self::Cycle(_) => true, + | Self::Cycle(_) + | Self::TakeWhile(_) + | Self::DropWhile(_) + | Self::FilterFalse(_) + | Self::StarMap(_) => true, } } @@ -123,7 +165,11 @@ impl ItertoolsIter { | Self::Compress(_) | Self::Islice(_) | Self::Chain(_) - | Self::Cycle(_) => 0, + | Self::Cycle(_) + | Self::TakeWhile(_) + | Self::DropWhile(_) + | Self::FilterFalse(_) + | Self::StarMap(_) => 0, Self::Repeat(repeat) => repeat.size_hint(), } } @@ -138,6 +184,10 @@ impl ItertoolsIter { Self::Islice(islice) => islice.for_each_child_id(on_child), Self::Chain(chain) => chain.for_each_child_id(on_child), Self::Cycle(cycle) => cycle.for_each_child_id(on_child), + Self::TakeWhile(take) => take.for_each_child_id(on_child), + Self::DropWhile(drop_while) => drop_while.for_each_child_id(on_child), + Self::FilterFalse(filter) => filter.for_each_child_id(on_child), + Self::StarMap(starmap) => starmap.for_each_child_id(on_child), } } } @@ -153,6 +203,10 @@ impl HeapItem for ItertoolsIter { Self::Islice(islice) => islice.py_dec_ref_ids(stack), Self::Chain(chain) => chain.py_dec_ref_ids(stack), Self::Cycle(cycle) => cycle.py_dec_ref_ids(stack), + Self::TakeWhile(take) => take.py_dec_ref_ids(stack), + Self::DropWhile(drop_while) => drop_while.py_dec_ref_ids(stack), + Self::FilterFalse(filter) => filter.py_dec_ref_ids(stack), + Self::StarMap(starmap) => starmap.py_dec_ref_ids(stack), } } } @@ -195,7 +249,15 @@ impl<'h> PyTrait<'h> for HeapRead<'h, ItertoolsIter> { // Source-driving adaptors re-enter `py_next` on their wrapped // iterator, recursing on the native Rust stack; charge a recursion // level so deep nesting raises `RecursionError`, not a stack overflow. - Kind::Pairwise | Kind::Compress | Kind::Islice | Kind::Chain | Kind::Cycle => { + Kind::Pairwise + | Kind::Compress + | Kind::Islice + | Kind::Chain + | Kind::Cycle + | Kind::TakeWhile + | Kind::DropWhile + | Kind::FilterFalse + | Kind::StarMap => { let mut guard = vm.recursion_guard()?; let vm = &mut *guard; match kind { @@ -204,6 +266,10 @@ impl<'h> PyTrait<'h> for HeapRead<'h, ItertoolsIter> { Kind::Islice => islice::next(self, vm), Kind::Chain => chain::next(self, vm), Kind::Cycle => cycle::next(self, vm), + Kind::TakeWhile => takewhile::next(self, vm), + Kind::DropWhile => dropwhile::next(self, vm), + Kind::FilterFalse => filterfalse::next(self, vm), + Kind::StarMap => starmap::next(self, vm), Kind::Count | Kind::Repeat => unreachable!("handled above"), } } @@ -217,7 +283,15 @@ impl<'h> PyTrait<'h> for HeapRead<'h, ItertoolsIter> { match self.get(vm.heap).kind() { Kind::Count => count::repr_fmt(self, f, vm, heap_ids), Kind::Repeat => repeat::repr_fmt(self, f, vm, heap_ids), - Kind::Pairwise | Kind::Compress | Kind::Islice | Kind::Chain | Kind::Cycle => { + Kind::Pairwise + | Kind::Compress + | Kind::Islice + | Kind::Chain + | Kind::Cycle + | Kind::TakeWhile + | Kind::DropWhile + | Kind::FilterFalse + | Kind::StarMap => { let type_name = self.py_type(vm).name(vm.heap, vm.interns); Ok(write!(f, "<{type_name} object>")?) } diff --git a/crates/monty/src/types/itertools/starmap.rs b/crates/monty/src/types/itertools/starmap.rs new file mode 100644 index 000000000..d3e9378af --- /dev/null +++ b/crates/monty/src/types/itertools/starmap.rs @@ -0,0 +1,89 @@ +//! `itertools.starmap(function, iterable)` — each item unpacked as the arguments. + +use serde::{Deserialize, Serialize}; + +use crate::{ + args::{ArgValues, KwargsValues}, + bytecode::VM, + defer_drop, + exception_private::RunResult, + heap::{HeapId, HeapRead}, + types::{iter::collect_owned_iterable, itertools::ItertoolsIter}, + value::Value, +}; + +/// Yields `function(*item)` for each item of `source`. +/// +/// Every item must itself be iterable — `starmap(pow, [5])` raises `TypeError: +/// 'int' object is not iterable` mid-iteration, exactly as in CPython. +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct StarMap { + function: Value, + source: Value, +} + +impl StarMap { + /// Takes ownership of both, with `source` already resolved by `py_iter`. + pub(crate) fn new(function: Value, source: Value) -> Self { + Self { function, source } + } + + /// Invokes `on_child` for each heap id this iterator owns (GC trace hook). + pub(crate) fn for_each_child_id(&self, mut on_child: impl FnMut(HeapId)) { + if let Value::Ref(id) = &self.function { + on_child(*id); + } + if let Value::Ref(id) = &self.source { + on_child(*id); + } + } + + /// Releases the refs this iterator owns (mirrors `for_each_child_id`). + pub(crate) fn py_dec_ref_ids(&mut self, stack: &mut Vec) { + self.function.py_dec_ref_ids(stack); + self.source.py_dec_ref_ids(stack); + } +} + +/// Pulls one item, spreads it into an argument list, and calls the function. +pub(super) fn next<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) -> RunResult> { + let ItertoolsIter::StarMap(starmap) = iter.get(vm.heap) else { + unreachable!("dispatched on Kind::StarMap") + }; + let function = starmap.function.clone_with_heap(vm.heap); + let source = starmap.source.clone_with_heap(vm.heap); + defer_drop!(function, vm); + defer_drop!(source, vm); + + let item = { + let mut read = source.read(vm); + read.py_next(vm) + }; + let Some(item) = item? else { + return Ok(None); + }; + // `collect_owned_iterable` consumes `item` on both paths, raising here for + // a non-iterable one, so the arguments never outlive a failed spread. + let args: Vec = collect_owned_iterable(item, vm)?; + vm.evaluate_function("starmap()", function, pack_args(args)).map(Some) +} + +/// Packs the spread item into the arity-specific [`ArgValues`] shape. +/// +/// The small forms are not interchangeable with `ArgsKargs`: extractors such as +/// `ArgValues::get_one_arg` match `One` structurally, so a one-element +/// `ArgsKargs` is rejected as the wrong shape (`abs() takes exactly one +/// argument`) even though the arity is right. +fn pack_args(items: Vec) -> ArgValues { + match <[Value; 1]>::try_from(items) { + Ok([first]) => ArgValues::One(first), + Err(items) => match <[Value; 2]>::try_from(items) { + Ok([first, second]) => ArgValues::Two(first, second), + Err(items) if items.is_empty() => ArgValues::Empty, + Err(items) => ArgValues::ArgsKargs { + args: items, + kwargs: KwargsValues::Empty, + }, + }, + } +} diff --git a/crates/monty/src/types/itertools/step.rs b/crates/monty/src/types/itertools/step.rs new file mode 100644 index 000000000..250f99c09 --- /dev/null +++ b/crates/monty/src/types/itertools/step.rs @@ -0,0 +1,42 @@ +//! The per-item step the predicate-driven adaptors share. +//! +//! `takewhile`, `dropwhile` and `filterfalse` each pull one item and decide +//! what to do with it. Only the decision differs, so the fetch — and the guards +//! that keep a raising test from leaking the item — live here rather than in +//! three copies. + +use crate::{bytecode::VM, defer_drop, exception_private::RunResult, heap::DropGuard, value::Value}; + +/// Pulls one item from `source`, releasing the caller's clone before returning. +/// +/// `source` is taken by value because it must be a clone: `py_next` re-enters +/// the VM, and the adaptor's own reference is unreachable behind that borrow. +pub(super) fn next_item(source: Value, vm: &mut VM<'_>) -> RunResult> { + defer_drop!(source, vm); + let mut read = source.read(vm); + read.py_next(vm) +} + +/// Pulls one item and applies `test` to it, returning both the item and the +/// answer. +/// +/// The item is guarded across `test`, so a predicate that raises drops it +/// instead of leaking, and `predicate` is released on every path. Callers get +/// the item back owned and decide whether to yield it — which is the only part +/// that differs between the adaptors. +pub(super) fn next_tested<'h>( + predicate: Value, + source: Value, + vm: &mut VM<'h>, + test: impl FnOnce(&Value, &Value, &mut VM<'h>) -> RunResult, +) -> RunResult> { + defer_drop!(predicate, vm); + let Some(item) = next_item(source, vm)? else { + return Ok(None); + }; + let mut item_guard = DropGuard::new(item, vm); + let (item, vm) = item_guard.as_parts_mut(); + let answer = test(predicate, item, vm)?; + let (item, _) = item_guard.into_parts(); + Ok(Some((item, answer))) +} diff --git a/crates/monty/src/types/itertools/takewhile.rs b/crates/monty/src/types/itertools/takewhile.rs new file mode 100644 index 000000000..5e96bae70 --- /dev/null +++ b/crates/monty/src/types/itertools/takewhile.rs @@ -0,0 +1,104 @@ +//! `itertools.takewhile(predicate, iterable)` — the leading run that passes. + +use serde::{Deserialize, Serialize}; + +use crate::{ + bytecode::VM, + exception_private::RunResult, + heap::{DropWithContext, HeapId, HeapRead}, + predicate::call_predicate, + types::itertools::{ItertoolsIter, step::next_tested}, + value::Value, +}; + +/// Yields items from `source` until `predicate` first returns false. +/// +/// Both fields go `None` together at that point, which latches the adaptor as +/// spent and releases them there and then: CPython stops calling the predicate +/// — and stops touching the source — once it has failed, so the rejected item +/// is the last thing either value ever sees. A source that merely runs out is +/// not a rejection and does not latch, so both stay `Some`. +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct TakeWhile { + predicate: Option, + source: Option, +} + +impl TakeWhile { + /// Takes ownership of both, with `source` already resolved by `py_iter`. + pub(crate) fn new(predicate: Value, source: Value) -> Self { + Self { + predicate: Some(predicate), + source: Some(source), + } + } + + /// Invokes `on_child` for each heap id this iterator owns (GC trace hook). + pub(crate) fn for_each_child_id(&self, mut on_child: impl FnMut(HeapId)) { + if let Some(Value::Ref(id)) = &self.predicate { + on_child(*id); + } + if let Some(Value::Ref(id)) = &self.source { + on_child(*id); + } + } + + /// Releases the refs this iterator owns (mirrors `for_each_child_id`). + pub(crate) fn py_dec_ref_ids(&mut self, stack: &mut Vec) { + if let Some(predicate) = &mut self.predicate { + predicate.py_dec_ref_ids(stack); + } + if let Some(source) = &mut self.source { + source.py_dec_ref_ids(stack); + } + } +} + +/// Pulls one item and yields it if the predicate holds, latching on the first +/// rejection so nothing after it is examined. +pub(super) fn next<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) -> RunResult> { + let ItertoolsIter::TakeWhile(take) = iter.get(vm.heap) else { + unreachable!("dispatched on Kind::TakeWhile") + }; + let (Some(predicate), Some(source)) = (&take.predicate, &take.source) else { + return Ok(None); + }; + let predicate = predicate.clone_with_heap(vm.heap); + let source = source.clone_with_heap(vm.heap); + + // Exhaustion does NOT latch: CPython only sets `stop` when the predicate + // fails, so a source that raises `StopIteration` and later yields again is + // re-driven on the next call rather than treated as spent. + let Some((item, accepted)) = next_tested(predicate, source, vm, |predicate, item, vm| { + call_predicate(predicate, item, "takewhile()", vm) + })? + else { + return Ok(None); + }; + + if accepted { + Ok(Some(item)) + } else { + // The rejected item is discarded rather than yielded, and the latch + // means nothing after it is ever examined. + item.drop_with(vm); + finish(iter, vm); + Ok(None) + } +} + +/// Latches the adaptor as spent, so neither the source nor the predicate is +/// reached again. +/// +/// Releasing both here rather than at destruction lets whatever the source +/// holds (a generator's frame, a file) be reclaimed as soon as the run ends, +/// matching what `pairwise` and `islice` do when they stop. +fn finish<'h>(iter: &mut HeapRead<'h, ItertoolsIter>, vm: &mut VM<'h>) { + let ItertoolsIter::TakeWhile(take) = iter.get_mut(vm.heap) else { + unreachable!("dispatched on Kind::TakeWhile") + }; + // Both taken under the one borrow, then dropped. + let (predicate, source) = (take.predicate.take(), take.source.take()); + predicate.drop_with(vm); + source.drop_with(vm); +} diff --git a/crates/monty/src/types/type.rs b/crates/monty/src/types/type.rs index 5d63bef93..45f16421f 100644 --- a/crates/monty/src/types/type.rs +++ b/crates/monty/src/types/type.rs @@ -193,6 +193,14 @@ pub enum Type { ItertoolsChain, #[strum(serialize = "itertools.cycle")] ItertoolsCycle, + #[strum(serialize = "itertools.takewhile")] + ItertoolsTakeWhile, + #[strum(serialize = "itertools.dropwhile")] + ItertoolsDropWhile, + #[strum(serialize = "itertools.filterfalse")] + ItertoolsFilterFalse, + #[strum(serialize = "itertools.starmap")] + ItertoolsStarMap, } /// Writes the canonical static name of every non-[`Instance`](Type::Instance) @@ -331,6 +339,10 @@ impl Type { | Self::ItertoolsIslice | Self::ItertoolsChain | Self::ItertoolsCycle + | Self::ItertoolsTakeWhile + | Self::ItertoolsDropWhile + | Self::ItertoolsFilterFalse + | Self::ItertoolsStarMap ) } diff --git a/crates/monty/test_cases/itertools__adaptors.py b/crates/monty/test_cases/itertools__adaptors.py index b04032eba..f188f6e72 100644 --- a/crates/monty/test_cases/itertools__adaptors.py +++ b/crates/monty/test_cases/itertools__adaptors.py @@ -258,6 +258,32 @@ except TypeError as exc: assert str(exc) == 'cycle() takes no keyword arguments' + +# starmap has no spent flag either, so a source that stops and then yields again +# is re-driven rather than treated as finished. +class StutteringPairs: + def __init__(self): + self.calls = 0 + + def __iter__(self): + return self + + def __next__(self): + self.calls += 1 + if self.calls == 2: + raise StopIteration + return (self.calls, 2) + + +starred = itertools.starmap(pow, StutteringPairs()) +assert next(starred) == 1 +try: + next(starred) + assert False, 'expected StopIteration' +except StopIteration: + pass +assert next(starred) == 9 + # === Iterator protocol === # Every adaptor is its own iterator, and exhaustion raises StopIteration rather # than returning a sentinel. @@ -373,3 +399,244 @@ def __next__(self): # Membership consumes the adaptor until it matches. assert 3 in itertools.chain([1, 2], [3]) assert 'z' not in itertools.compress('abc', [1, 1, 1]) + + +# === takewhile === +assert list(itertools.takewhile(lambda x: x < 3, [1, 2, 3, 4, 1])) == [1, 2] +assert list(itertools.takewhile(lambda x: x < 3, [])) == [] +assert list(itertools.takewhile(lambda x: False, [1, 2])) == [] +assert list(itertools.takewhile(lambda x: True, 'ab')) == ['a', 'b'] +# A None predicate is only reached when there is an item to test. +assert list(itertools.takewhile(None, [])) == [] + +# The predicate stops being called at the first rejection, and the adaptor +# stays spent afterwards. +seen = [] + + +def under_two(x): + seen.append(x) + return x < 2 + + +spent = itertools.takewhile(under_two, [1, 2, 3]) +assert list(spent) == [1] +assert seen == [1, 2] +assert list(spent) == [] +assert seen == [1, 2] + + +# Only a rejected item latches the adaptor. A source that raises StopIteration +# and then yields again is re-driven, as CPython's takewhile does — it is +# `pairwise`/`islice` that release their source when it runs out, not this one. +class Stuttering: + def __init__(self): + self.calls = 0 + + def __iter__(self): + return self + + def __next__(self): + self.calls += 1 + if self.calls == 2: + raise StopIteration + return self.calls + + +stutter = itertools.takewhile(lambda x: x < 4, Stuttering()) +assert next(stutter) == 1 +try: + next(stutter) + assert False, 'expected StopIteration' +except StopIteration: + pass +assert next(stutter) == 3 +# The same source behaviour through the other two, which never latched. Each is +# driven *past* the StopIteration, since stopping at the first item would pass +# whether or not the adaptor wrongly treated exhaustion as terminal. +for adaptor in ( + itertools.dropwhile(lambda x: False, Stuttering()), + itertools.filterfalse(lambda x: False, Stuttering()), +): + assert next(adaptor) == 1 + try: + next(adaptor) + assert False, 'expected StopIteration' + except StopIteration: + pass + assert next(adaptor) == 3 + + +# Latching stops the source being touched, not only the predicate being called: +# a second drain must not reach it. `Counting` reports how often it was asked. +class Counting: + def __init__(self, items): + self.items = list(items) + self.reads = 0 + + def __iter__(self): + return self + + def __next__(self): + self.reads += 1 + if not self.items: + raise StopIteration + return self.items.pop(0) + + +counted = Counting([1, 5, 2]) +latched = itertools.takewhile(lambda x: x < 3, counted) +assert list(latched) == [1] +assert counted.reads == 2 +assert list(latched) == [] +assert counted.reads == 2 + +# === dropwhile === +assert list(itertools.dropwhile(lambda x: x < 3, [1, 2, 3, 4, 1])) == [3, 4, 1] +assert list(itertools.dropwhile(lambda x: x < 3, [])) == [] +assert list(itertools.dropwhile(lambda x: True, [1, 2])) == [] +assert list(itertools.dropwhile(lambda x: False, [1, 2])) == [1, 2] + +# Once the predicate has failed it is never consulted again, so later items +# are yielded even when they would have satisfied it. +dropped = [] + + +def small(x): + dropped.append(x) + return x < 2 + + +assert list(itertools.dropwhile(small, [1, 2, 3, 0])) == [2, 3, 0] +assert dropped == [1, 2] + +# === filterfalse === +assert list(itertools.filterfalse(lambda x: x % 2, range(6))) == [0, 2, 4] +assert list(itertools.filterfalse(lambda x: True, [1, 2])) == [] +assert list(itertools.filterfalse(lambda x: False, [1, 2])) == [1, 2] +assert list(itertools.filterfalse(lambda x: x % 2, [])) == [] +# A None predicate selects the truth test, keeping the falsy items. +assert list(itertools.filterfalse(None, [0, 1, '', 'a', [], None])) == [0, '', [], None] +assert list(itertools.filterfalse(None, [1, 'a', [2]])) == [] + +# === starmap === +assert list(itertools.starmap(pow, [(2, 5), (3, 2)])) == [32, 9] +assert list(itertools.starmap(lambda a, b: a + b, ['ab', 'cd'])) == ['ab', 'cd'] +assert list(itertools.starmap(max, [[1, 5, 3]])) == [5] +assert list(itertools.starmap(pow, [])) == [] +# Items are spread, so a single-element item calls a single-argument function. +assert list(itertools.starmap(abs, [(-2,), (3,)])) == [2, 3] + +# === Iterator protocol === +for adaptor in ( + itertools.takewhile(bool, [1]), + itertools.dropwhile(bool, [1]), + itertools.filterfalse(bool, [0]), + itertools.starmap(pow, [(2, 2)]), +): + assert iter(adaptor) is adaptor + +exhausted = itertools.takewhile(lambda x: True, [1]) +assert next(exhausted) == 1 +try: + next(exhausted) + assert False, 'expected StopIteration' +except StopIteration: + pass + + +# Returning one from a user `__iter__` works, which needs the adaptor to count +# as a concrete iterator type and not just as something iterable. +class Wrapped: + def __init__(self, adaptor): + self.adaptor = adaptor + + def __iter__(self): + return self.adaptor + + +assert list(Wrapped(itertools.takewhile(lambda x: x < 3, [1, 2, 3]))) == [1, 2] +assert list(Wrapped(itertools.dropwhile(lambda x: x < 3, [1, 2, 3]))) == [3] +assert list(Wrapped(itertools.filterfalse(None, [0, 1]))) == [0] +assert list(Wrapped(itertools.starmap(pow, [(2, 3)]))) == [8] + + +# === Signature errors === +for name, builder in ( + ('takewhile', itertools.takewhile), + ('dropwhile', itertools.dropwhile), + ('filterfalse', itertools.filterfalse), + ('starmap', itertools.starmap), +): + try: + builder() + assert False, 'expected TypeError' + except TypeError as exc: + assert str(exc) == name + ' expected 2 arguments, got 0' + try: + builder(bool) + assert False, 'expected TypeError' + except TypeError as exc: + assert str(exc) == name + ' expected 2 arguments, got 1' + try: + builder(bool, [], []) + assert False, 'expected TypeError' + except TypeError as exc: + assert str(exc) == name + ' expected 2 arguments, got 3' + try: + builder(bool, iterable=[]) + assert False, 'expected TypeError' + except TypeError as exc: + assert str(exc) == name + '() takes no keyword arguments' + # The iterable is resolved eagerly, so a non-iterable raises up front. + try: + builder(bool, 5) + assert False, 'expected TypeError' + except TypeError as exc: + assert str(exc) == "'int' object is not iterable" + +# A None predicate is a call-time failure, not a construction-time one. +for adaptor in (itertools.takewhile(None, [1]), itertools.dropwhile(None, [1])): + try: + next(adaptor) + assert False, 'expected TypeError' + except TypeError as exc: + assert str(exc) == "'NoneType' object is not callable" + +# starmap needs each item to be iterable, discovered as it reaches them. +bad_items = itertools.starmap(pow, [5]) +try: + next(bad_items) + assert False, 'expected TypeError' +except TypeError as exc: + assert str(exc) == "'int' object is not iterable" + + +# === Exceptions propagate out of the callable === +def explode(x): + raise ValueError('bang') + + +for adaptor in ( + itertools.takewhile(explode, [1]), + itertools.dropwhile(explode, [1]), + itertools.filterfalse(explode, [1]), + itertools.starmap(explode, [(1,)]), +): + try: + next(adaptor) + assert False, 'expected ValueError' + except ValueError as exc: + assert str(exc) == 'bang' + +# === Composition with the source-wrapping adaptors === +assert list(itertools.takewhile(lambda x: x < 4, itertools.count())) == [0, 1, 2, 3] +assert list(itertools.islice(itertools.dropwhile(lambda x: x < 3, itertools.count()), 3)) == [3, 4, 5] +assert list(itertools.filterfalse(None, itertools.islice(itertools.cycle([0, 1]), 4))) == [0, 0] +assert list(itertools.starmap(pow, itertools.pairwise([2, 3, 4]))) == [8, 81] +assert sorted(itertools.filterfalse(lambda x: x > 1, itertools.chain([0, 1], [2]))) == [0, 1] + +# Every argument-count shape, since each is packed differently internally. +assert list(itertools.starmap(lambda: 7, [()])) == [7] +assert list(itertools.starmap(lambda a, b, c: a + b + c, [(1, 2, 3)])) == [6] +assert list(itertools.starmap(lambda *a: len(a), [(1, 2, 3, 4)])) == [4] diff --git a/crates/monty/test_cases/refcount__itertools_adaptors.py b/crates/monty/test_cases/refcount__itertools_adaptors.py index 9f2a7fd89..b0008394c 100644 --- a/crates/monty/test_cases/refcount__itertools_adaptors.py +++ b/crates/monty/test_cases/refcount__itertools_adaptors.py @@ -120,5 +120,103 @@ def __next__(self): except TypeError as exc: assert str(exc) == "'int' object is not iterable" + +# The predicate-driven adaptors own a CALLABLE as well as a source, so each has +# a second trace edge. A closure is used deliberately: a plain `def` is an +# immediate `Value`, not a heap ref, so it would exercise no hook at all. +def make_shorter_than(limit): + bound = list(range(limit)) + + def shorter(x): + return len(x) < len(bound) + + return shorter + + +def make_adder(): + bound = [1] + + def add(a, b=0): + return a + b + len(bound) + + return add + + +def make_boom(): + bound = [1] + + def boom(*args): + raise ValueError('boom' + str(len(bound))) + + return boom + + +# Each closure is passed inline and never named, so the adaptor's callable +# field is its only referrer; the items are lists for the same reason. +take_live = itertools.takewhile(make_shorter_than(3), [[1], [2]]) +next(take_live) +drop_live = itertools.dropwhile(make_shorter_than(0), [[1], [2]]) +next(drop_live) +filter_live = itertools.filterfalse(make_shorter_than(0), [[1], [2]]) +next(filter_live) +star_live = itertools.starmap(make_adder(), [(1,), (2,)]) +next(star_live) + +# filterfalse with a None predicate leaves only the source edge, so a hook that +# traces the callable twice still fails to reach these. +filter_none = itertools.filterfalse(None, [[1], []]) + +# The freeing paths: `py_dec_ref_ids` runs only on release, so each of these +# must be dropped rather than merely held. +gone_take = itertools.takewhile(make_shorter_than(3), [[1], [2]]) +next(gone_take) +gone_take = None +gone_drop = itertools.dropwhile(make_shorter_than(0), [[1], [2]]) +next(gone_drop) +gone_drop = None +gone_filter = itertools.filterfalse(make_shorter_than(0), [[1], [2]]) +next(gone_filter) +gone_filter = None +gone_star = itertools.starmap(make_adder(), [(1,)]) +next(gone_star) +gone_star = None + +# A rejected item is dropped rather than yielded — the guard path inside `next`. +rejected = itertools.takewhile(make_shorter_than(0), [[1], [2]]) +assert list(rejected) == [] + +# A callable that raises leaves `next` through a `?` while the guard still +# holds the item being tested, and for starmap the arguments already collected. +pred_erroring = itertools.takewhile(make_boom(), [[1], [2]]) +try: + next(pred_erroring) +except ValueError: + pass + +star_erroring = itertools.starmap(make_boom(), [(1, 2)]) +try: + next(star_erroring) +except ValueError: + pass + + +# Spending an adaptor releases what it can no longer reach, THERE AND THEN +# rather than at destruction — as `pairwise` and `islice` do above. Each source +# and callable is named separately, so a count of 1 means the spent adaptor let +# go of it and 2 means it is still held. The adaptors stay bound so it is the +# release being measured, not their destruction. +take_pred = make_shorter_than(0) +take_source = iter([[1], [2]]) +latched_take = itertools.takewhile(take_pred, take_source) +assert list(latched_take) == [] + +# `dropwhile` releases neither: the predicate goes uncalled after the first +# rejection but stays owned to destruction, as CPython holds `lz->func`, and +# it never latches, so every later `next` drives the source again. +drop_pred = make_shorter_than(1) +drop_source = iter([[], [1]]) +past_drop = itertools.dropwhile(drop_pred, drop_source) +assert next(past_drop) == [1] + len('done') -# ref-counts={'itertools': 1, 'live': 1, 'primed': 1, 'cyclic': 2, 'paired': 1, 'sliced': 1, 'chained': 1, 'cycled': 1, 'replaying': 1, 'Boom': 2, 'erroring': 1, 'spent_source': 1, 'spent_pairwise': 1, 'stopped_source': 1, 'stopped_islice': 1, 'drained_source': 1, 'drained_islice': 1, 'chain_drained_source': 1, 'chain_drained': 1, 'chain_unreached_source': 1, 'chain_failed': 1} +# ref-counts={'itertools': 1, 'live': 1, 'primed': 1, 'cyclic': 2, 'paired': 1, 'sliced': 1, 'chained': 1, 'cycled': 1, 'replaying': 1, 'Boom': 2, 'erroring': 1, 'spent_source': 1, 'spent_pairwise': 1, 'stopped_source': 1, 'stopped_islice': 1, 'drained_source': 1, 'drained_islice': 1, 'chain_drained_source': 1, 'chain_drained': 1, 'chain_unreached_source': 1, 'chain_failed': 1, 'take_live': 1, 'drop_live': 1, 'filter_live': 1, 'star_live': 1, 'filter_none': 1, 'rejected': 1, 'pred_erroring': 1, 'star_erroring': 1, 'take_pred': 1, 'take_source': 1, 'latched_take': 1, 'drop_pred': 2, 'drop_source': 2, 'past_drop': 1} diff --git a/crates/monty/tests/main.rs b/crates/monty/tests/main.rs index 111335c9e..cf35f96ff 100644 --- a/crates/monty/tests/main.rs +++ b/crates/monty/tests/main.rs @@ -155,6 +155,40 @@ fn external_function_in_next_raises_not_implemented() { ); } +/// The `itertools` adaptors that apply a callable drive it through +/// `evaluate_function`, so one reaching an external function cannot suspend and +/// raises `NotImplementedError` (see `limitations/itertools.md`). Rust-side for +/// the same reason as the tests above: on CPython the external is an ordinary +/// function and the call would succeed. +/// +/// Both call sites are covered — the predicate helper shared by `takewhile`, +/// `dropwhile` and `filterfalse`, and `starmap`, which calls its function +/// itself and so names itself in the error separately. +#[test] +fn external_function_as_itertools_callable_raises_not_implemented() { + for (call, adaptor) in [ + ("itertools.takewhile(ext_fn, [1])", "takewhile"), + ("itertools.starmap(ext_fn, [(1,)])", "starmap"), + ] { + let expr = format!("list({call})"); + let code = format!("import itertools\n\n{expr}"); + let ex = MontyRun::new(code, "test.py", vec!["ext_fn".to_owned()], CompileOptions::default()).unwrap(); + let err = ex + .run_no_limits(vec![MontyObject::Function { + name: "ext_fn".to_owned(), + docstring: None, + }]) + .unwrap_err(); + let carets = "~".repeat(expr.len()); + assert_eq!( + err.to_string(), + format!( + "Traceback (most recent call last):\n File \"test.py\", line 3, in \n {expr}\n {carets}\nNotImplementedError: {adaptor}(): external function 'ext_fn' is not yet supported in this context" + ) + ); + } +} + /// The 3-arg `type()` form rejects non-empty bases because Monty classes /// cannot inherit (documented in `limitations/classes.md`). Kept as a /// Rust-side test because CPython accepts bases, so the comparative diff --git a/crates/monty/tests/resource_limits.rs b/crates/monty/tests/resource_limits.rs index f3977939c..350951bd1 100644 --- a/crates/monty/tests/resource_limits.rs +++ b/crates/monty/tests/resource_limits.rs @@ -1008,3 +1008,31 @@ a < b assert_eq!(exc.exc_type(), ExcType::RecursionError, "build: {build}"); } } + +/// Every `itertools` adaptor whose `next` can loop natively without yielding. +/// +/// Each pairs a discarding or draining adaptor with an infinite source, so the +/// loop never returns to the VM. `dropwhile` appears twice because a builtin +/// predicate and a short user-defined one fail the same way: the dispatch +/// checkpoint is per-`run()`, so a callback under `CHECK_INTERVAL` +/// instructions restarts the countdown instead of reaching it. +const ITERTOOLS_INFINITE_LOOPS: &[&str] = &[ + "next(itertools.dropwhile(bool, itertools.count(1)))", + "def always(x):\n return True\nnext(itertools.dropwhile(always, itertools.count(1)))", + "next(itertools.filterfalse(bool, itertools.count(1)))", + "next(itertools.compress(itertools.count(1), itertools.repeat(0)))", + "next(itertools.islice(itertools.count(1), 10**18, None))", + "next(itertools.starmap(max, itertools.repeat(itertools.count(1))))", +]; + +/// Test that adaptors discarding items from an infinite source still time out. +/// +/// These loops sit inside one bytecode instruction and drive native sources, so +/// nothing returns to the dispatch checkpoint; each must poll the tracker +/// itself or `max_duration` is unenforceable. +#[test] +fn timeout_in_itertools_adaptor_loops() { + for expr in ITERTOOLS_INFINITE_LOOPS { + assert_timeout_in_builtin(&format!("import itertools\n{expr}"), expr); + } +} diff --git a/limitations/itertools.md b/limitations/itertools.md index 0d1a6eb03..2a01a03f0 100644 --- a/limitations/itertools.md +++ b/limitations/itertools.md @@ -8,13 +8,15 @@ notes below. `count(start=0, step=1)`, `repeat(object, times=?)`, `pairwise(iterable)`, `compress(data, selectors)`, `islice(iterable, [start,] stop[, step])`, -`chain(*iterables)`, `cycle(iterable)`. +`chain(*iterables)`, `cycle(iterable)`, `takewhile(predicate, iterable)`, +`dropwhile(predicate, iterable)`, `filterfalse(predicate, iterable)`, +`starmap(function, iterable)`. ## Not implemented Everything else: `accumulate`, `batched`, `combinations`, -`combinations_with_replacement`, `dropwhile`, `filterfalse`, `groupby`, -`permutations`, `product`, `starmap`, `takewhile`, `tee`, `zip_longest`. +`combinations_with_replacement`, `groupby`, `permutations`, `product`, `tee`, +`zip_longest`. `chain.from_iterable` is also absent, even though `chain` itself is implemented: it is a classmethod reached through an attribute on the `chain` @@ -53,6 +55,16 @@ raise `AttributeError` at runtime. is ``, where CPython appends ` at 0x...`. This is Monty's general iterator treatment (see ./iter.md), not specific to `itertools`. +- **A callable that suspends is rejected, not paused.** `takewhile`, + `dropwhile`, `filterfalse` and `starmap` apply their callable through the + synchronous `evaluate_function` path, which runs a frame to completion and + cannot yield to the host. A callable that reaches an external function, an + `os` operation, or a host method call therefore raises + `NotImplementedError: takewhile(): external function 'f' is not yet supported + in this context` where CPython would simply call it. This is the same + restriction that applies to `__init__`, `__next__` and `__repr__` (see + `limitations/classes.md`); ordinary sandbox-defined functions and lambdas are + unaffected. - **Crossing the host boundary loses the repr.** A `count` / `repeat` object returned to the host arrives as `` / `` rather than its in-sandbox `repr()` @@ -87,6 +99,14 @@ or duration limit, and then raises `MemoryError` rather than exhausting. Under runs until the host itself runs out of memory. This is the same exposure as a `while True:` loop, not something specific to `itertools`. +The adaptors that discard items without yielding — `dropwhile` and +`filterfalse` before their first accepted item, `compress` past a falsy run, +`islice` skipping to `start`, `chain` crossing an exhausted source — poll +`max_duration` themselves while looping, so a discarding pass over an infinite +source raises `TimeoutError` instead of spinning. The poll is amortized (once +per 64 items), so the limit can be overshot by up to that much work. CPython +has no duration limit at all and would loop forever. + `cycle(iterable)` must buffer every item it has seen so far in order to replay them, and that buffer is charged against `max_memory` as it grows, so cycling over a very long source raises `MemoryError` at the limit rather than at the