PyIter protocol

DimitrisJim · Sep 29, 2021 · Sep 22, 2021 · Sep 22, 2021 · Sep 29, 2021 · 3273be1a6b98ded2b5fd7f1571169748d9b824ce
commit 3273be1a6b98ded2b5fd7f1571169748d9b824ce
diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs
@@ -85,13 +85,13 @@ impl PyDict {
                     return Err(vm.new_runtime_error("dict mutated during update".to_owned()));
                 }
             } else if let Some(keys) = vm.get_method(dict_obj.clone(), "keys") {
-                let keys = iterator::get_iter(vm, vm.invoke(&keys?, ())?)?;
+                let keys = vm.invoke(&keys?, ())?.get_iter(vm)?;
                 while let Some(key) = iterator::get_next_object(vm, &keys)? {
                     let val = dict_obj.get_item(key.clone(), vm)?;
                     dict.insert(vm, key, val)?;
                 }
             } else {
-                let iter = iterator::get_iter(vm, dict_obj)?;
+                let iter = dict_obj.get_iter(vm)?;
                 loop {
                     fn err(vm: &VirtualMachine) -> PyBaseExceptionRef {
                         vm.new_value_error("Iterator must have exactly two elements".to_owned())
@@ -100,7 +100,7 @@ impl PyDict {
                         Some(obj) => obj,
                         None => break,
                     };
-                    let elem_iter = iterator::get_iter(vm, element)?;
+                    let elem_iter = element.get_iter(vm)?;
                     let key = iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?;
                     let value =
                         iterator::get_next_object(vm, &elem_iter)?.ok_or_else(|| err(vm))?;

diff --git a/vm/src/builtins/enumerate.rs b/vm/src/builtins/enumerate.rs
@@ -6,7 +6,7 @@ use super::{
 use crate::common::lock::PyRwLock;
 use crate::{
     function::OptionalArg,
-    iterator,
+    protocol::PyIter,
     slots::{IteratorIterable, SlotConstructor, SlotIterator},
     IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
     TypeProtocol, VirtualMachine,
@@ -19,7 +19,7 @@ use num_traits::Zero;
 #[derive(Debug)]
 pub struct PyEnumerate {
     counter: PyRwLock<BigInt>,
-    iterator: PyObjectRef,
+    iterator: PyIter,
 }
 
 impl PyValue for PyEnumerate {
@@ -30,19 +30,20 @@ impl PyValue for PyEnumerate {
 
 #[derive(FromArgs)]
 pub struct EnumerateArgs {
-    iterable: PyObjectRef,
+    iterator: PyIter,
     #[pyarg(any, optional)]
     start: OptionalArg<PyIntRef>,
 }
 
 impl SlotConstructor for PyEnumerate {
     type Args = EnumerateArgs;
 
-    fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
-        let counter = args
-            .start
-            .map_or_else(BigInt::zero, |start| start.as_bigint().clone());
-        let iterator = iterator::get_iter(vm, args.iterable)?;
+    fn py_new(
+        cls: PyTypeRef,
+        Self::Args { iterator, start }: Self::Args,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        let counter = start.map_or_else(BigInt::zero, |start| start.as_bigint().clone());
         PyEnumerate {
             counter: PyRwLock::new(counter),
             iterator,
@@ -57,7 +58,7 @@ impl PyEnumerate {}
 impl IteratorIterable for PyEnumerate {}
 impl SlotIterator for PyEnumerate {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
-        let next_obj = iterator::call_next(vm, &zelf.iterator)?;
+        let next_obj = zelf.iterator.next(vm)?;
         let mut counter = zelf.counter.write();
         let position = counter.clone();
         *counter += 1;

diff --git a/vm/src/builtins/filter.rs b/vm/src/builtins/filter.rs
@@ -1,6 +1,6 @@
 use super::PyTypeRef;
 use crate::{
-    iterator,
+    protocol::PyIter,
     slots::{IteratorIterable, SlotConstructor, SlotIterator},
     PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine,
 };
@@ -13,7 +13,7 @@ use crate::{
 #[derive(Debug)]
 pub struct PyFilter {
     predicate: PyObjectRef,
-    iterator: PyObjectRef,
+    iterator: PyIter,
 }
 
 impl PyValue for PyFilter {
@@ -22,24 +22,10 @@ impl PyValue for PyFilter {
     }
 }
 
-#[derive(FromArgs)]
-pub struct FilterArgs {
-    #[pyarg(positional)]
-    function: PyObjectRef,
-    #[pyarg(positional)]
-    iterable: PyObjectRef,
-}
-
 impl SlotConstructor for PyFilter {
-    type Args = FilterArgs;
-
-    fn py_new(
-        cls: PyTypeRef,
-        Self::Args { function, iterable }: Self::Args,
-        vm: &VirtualMachine,
-    ) -> PyResult {
-        let iterator = iterator::get_iter(vm, iterable)?;
+    type Args = (PyObjectRef, PyIter);
 
+    fn py_new(cls: PyTypeRef, (function, iterator): Self::Args, vm: &VirtualMachine) -> PyResult {
         Self {
             predicate: function,
             iterator,
@@ -55,9 +41,8 @@ impl IteratorIterable for PyFilter {}
 impl SlotIterator for PyFilter {
     fn next(zelf: &PyRef<Self>, vm: &VirtualMachine) -> PyResult {
         let predicate = &zelf.predicate;
-        let iterator = &zelf.iterator;
         loop {
-            let next_obj = iterator::call_next(vm, iterator)?;
+            let next_obj = zelf.iterator.next(vm)?;
             let predicate_value = if vm.is_none(predicate) {
                 next_obj.clone()
             } else {

diff --git a/vm/src/builtins/make_module.rs b/vm/src/builtins/make_module.rs
@@ -25,7 +25,8 @@ mod decl {
             ArgBytesLike, ArgCallable, ArgIterable, FuncArgs, KwArgs, OptionalArg, OptionalOption,
             PosArgs,
         },
-        iterator, py_io,
+        protocol::PyIter,
+        py_io,
         readline::{Readline, ReadlineResult},
         scope::Scope,
         slots::PyComparisonOp,
@@ -413,14 +414,15 @@ mod decl {
         iter_target: PyObjectRef,
         sentinel: OptionalArg<PyObjectRef>,
         vm: &VirtualMachine,
-    ) -> PyResult {
+    ) -> PyResult<PyIter> {
         if let OptionalArg::Present(sentinel) = sentinel {
             let callable = ArgCallable::try_from_object(vm, iter_target)?;
-            Ok(PyCallableIterator::new(callable, sentinel)
+            let iterator = PyCallableIterator::new(callable, sentinel)
                 .into_ref(vm)
-                .into_object())
+                .into_object();
+            Ok(PyIter::new(iterator))
         } else {
-            iterator::get_iter(vm, iter_target)
+            iter_target.get_iter(vm)
         }
     }
 
@@ -511,7 +513,7 @@ mod decl {
         default_value: OptionalArg<PyObjectRef>,
         vm: &VirtualMachine,
     ) -> PyResult {
-        iterator::call_next(vm, &iterator).or_else(|err| {
+        PyIter::new(iterator).next(vm).or_else(|err| {
             if err.isinstance(&vm.ctx.exceptions.stop_iteration) {
                 default_value.ok_or(err)
             } else {

diff --git a/vm/src/builtins/map.rs b/vm/src/builtins/map.rs
@@ -2,6 +2,7 @@ use super::PyTypeRef;
 use crate::{
     function::PosArgs,
     iterator,
+    protocol::PyIter,
     slots::{IteratorIterable, SlotConstructor, SlotIterator},
     PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine,
 };
@@ -14,7 +15,7 @@ use crate::{
 #[derive(Debug)]
 pub struct PyMap {
     mapper: PyObjectRef,
-    iterators: Vec<PyObjectRef>,
+    iterators: Vec<PyIter>,
 }
 
 impl PyValue for PyMap {
@@ -24,18 +25,11 @@ impl PyValue for PyMap {
 }
 
 impl SlotConstructor for PyMap {
-    type Args = (PyObjectRef, PosArgs<PyObjectRef>);
+    type Args = (PyObjectRef, PosArgs<PyIter>);
 
-    fn py_new(cls: PyTypeRef, (function, iterables): Self::Args, vm: &VirtualMachine) -> PyResult {
-        let iterators = iterables
-            .into_iter()
-            .map(|iterable| iterator::get_iter(vm, iterable))
-            .collect::<Result<Vec<_>, _>>()?;
-        PyMap {
-            mapper: function,
-            iterators,
-        }
-        .into_pyresult_with_type(vm, cls)
+    fn py_new(cls: PyTypeRef, (mapper, iterators): Self::Args, vm: &VirtualMachine) -> PyResult {
+        let iterators = iterators.into_vec();
+        PyMap { mapper, iterators }.into_pyresult_with_type(vm, cls)
     }
 }
 
@@ -44,7 +38,7 @@ impl PyMap {
     #[pymethod(magic)]
     fn length_hint(&self, vm: &VirtualMachine) -> PyResult<usize> {
         self.iterators.iter().try_fold(0, |prev, cur| {
-            let cur = iterator::length_hint(vm, cur.clone())?.unwrap_or(0);
+            let cur = iterator::length_hint(vm, cur.as_object().clone())?.unwrap_or(0);
             let max = std::cmp::max(prev, cur);
             Ok(max)
         })
@@ -57,7 +51,7 @@ impl SlotIterator for PyMap {
         let next_objs = zelf
             .iterators
             .iter()
-            .map(|iterator| iterator::call_next(vm, iterator))
+            .map(|iterator| iterator.next(vm))
             .collect::<Result<Vec<_>, _>>()?;
 
         // the mapper itself can raise StopIteration which does stop the map iteration

diff --git a/vm/src/builtins/mappingproxy.rs b/vm/src/builtins/mappingproxy.rs
@@ -1,7 +1,6 @@
 use super::{PyDict, PyStrRef, PyTypeRef};
 use crate::{
     function::OptionalArg,
-    iterator,
     slots::{Iterable, SlotConstructor},
     IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue,
     TryFromObject, VirtualMachine,
@@ -137,7 +136,8 @@ impl Iterable for PyMappingProxy {
                 PyDict::from_attributes(c.attributes.read().clone(), vm)?.into_pyobject(vm)
             }
         };
-        iterator::get_iter(vm, obj)
+        let iter = obj.get_iter(vm)?;
+        Ok(iter.into_object())
     }
 }
 

diff --git a/vm/src/builtins/range.rs b/vm/src/builtins/range.rs
@@ -30,7 +30,7 @@ fn iter_search(
     vm: &VirtualMachine,
 ) -> PyResult<usize> {
     let mut count = 0;
-    let iter = iterator::get_iter(vm, obj)?;
+    let iter = obj.get_iter(vm)?;
     while let Some(element) = iterator::get_next_object(vm, &iter)? {
         if vm.bool_eq(&item, &element)? {
             match flag {

diff --git a/vm/src/builtins/zip.rs b/vm/src/builtins/zip.rs
@@ -1,15 +1,15 @@
 use super::PyTypeRef;
 use crate::{
     function::PosArgs,
-    iterator,
+    protocol::PyIter,
     slots::{IteratorIterable, SlotConstructor, SlotIterator},
-    PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, VirtualMachine,
+    PyClassImpl, PyContext, PyRef, PyResult, PyValue, VirtualMachine,
 };
 
 #[pyclass(module = false, name = "zip")]
 #[derive(Debug)]
 pub struct PyZip {
-    iterators: Vec<PyObjectRef>,
+    iterators: Vec<PyIter>,
 }
 
 impl PyValue for PyZip {
@@ -19,13 +19,10 @@ impl PyValue for PyZip {
 }
 
 impl SlotConstructor for PyZip {
-    type Args = PosArgs;
+    type Args = PosArgs<PyIter>;
 
-    fn py_new(cls: PyTypeRef, iterables: Self::Args, vm: &VirtualMachine) -> PyResult {
-        let iterators = iterables
-            .into_iter()
-            .map(|iterable| iterator::get_iter(vm, iterable))
-            .collect::<Result<Vec<_>, _>>()?;
+    fn py_new(cls: PyTypeRef, iterators: Self::Args, vm: &VirtualMachine) -> PyResult {
+        let iterators = iterators.into_vec();
         PyZip { iterators }.into_pyresult_with_type(vm, cls)
     }
 }
@@ -42,7 +39,7 @@ impl SlotIterator for PyZip {
             let next_objs = zelf
                 .iterators
                 .iter()
-                .map(|iterator| iterator::call_next(vm, iterator))
+                .map(|iterator| iterator.next(vm))
                 .collect::<Result<Vec<_>, _>>()?;
 
             Ok(vm.ctx.new_tuple(next_objs))

diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -17,6 +17,7 @@ use crate::{
     exceptions::{self, ExceptionCtor},
     function::FuncArgs,
     iterator,
+    protocol::PyIter,
     scope::Scope,
     slots::PyComparisonOp,
     IdProtocol, ItemProtocol, PyMethod, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
@@ -857,8 +858,8 @@ impl ExecutingFrame<'_> {
             }
             bytecode::Instruction::GetIter => {
                 let iterated_obj = self.pop_value();
-                let iter_obj = iterator::get_iter(vm, iterated_obj)?;
-                self.push_value(iter_obj);
+                let iter_obj = iterated_obj.get_iter(vm)?;
+                self.push_value(iter_obj.into_object());
                 Ok(None)
             }
             bytecode::Instruction::GetAwaitable => {
@@ -1442,7 +1443,7 @@ impl ExecutingFrame<'_> {
     fn _send(&self, coro: &PyObjectRef, val: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         match self.builtin_coro(coro) {
             Some(coro) => coro.send(val, vm),
-            None if vm.is_none(&val) => iterator::call_next(vm, coro),
+            None if vm.is_none(&val) => PyIter::new(coro).next(vm),
             None => {
                 let meth = vm.get_attribute(coro.clone(), "send")?;
                 vm.invoke(&meth, (val,))
@@ -1512,7 +1513,7 @@ impl ExecutingFrame<'_> {
 
     /// The top of stack contains the iterator, lets push it forward
     fn execute_for_iter(&mut self, vm: &VirtualMachine, target: bytecode::Label) -> FrameResult {
-        let top_of_stack = self.last_value();
+        let top_of_stack = PyIter::new(self.last_value());
         let next_obj = iterator::get_next_object(vm, &top_of_stack);
 
         // Check the next object:

diff --git a/vm/src/function/argument.rs b/vm/src/function/argument.rs
@@ -1,7 +1,7 @@
 use super::IntoFuncArgs;
-use crate::builtins::iter::PySequenceIterator;
 use crate::{
-    iterator, PyObjectRef, PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine,
+    builtins::iter::PySequenceIterator, iterator, protocol::PyIter, PyObjectRef, PyResult, PyValue,
+    TryFromObject, TypeProtocol, VirtualMachine,
 };
 use std::marker::PhantomData;
 
@@ -102,7 +102,7 @@ where
     type Item = PyResult<T>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        iterator::get_next_object(self.vm, &self.obj)
+        iterator::get_next_object(self.vm, &PyIter::new(&self.obj))
             .transpose()
             .map(|x| x.and_then(|obj| T::try_from_object(self.vm, obj)))
     }