From ef4bea53c6d824c067b916ce3af654d93f280582 Mon Sep 17 00:00:00 2001
From: Gyubong <jopemachine@naver.com>
Date: Mon, 29 Aug 2022 19:04:39 +0900
Subject: [PATCH] Add `memory_iterator` (#4138)

---
 vm/src/builtins/memory.rs | 72 +++++++++++++++++++++++++++++++++++----
 vm/src/types/zoo.rs       |  2 ++
 whats_left.py             |  1 +
 3 files changed, 69 insertions(+), 6 deletions(-)

diff --git a/vm/src/builtins/memory.rs b/vm/src/builtins/memory.rs
index 6d0e95cb3..a59a99af8 100644
--- a/vm/src/builtins/memory.rs
+++ b/vm/src/builtins/memory.rs
@@ -1,6 +1,6 @@
 use super::{
-    PyBytes, PyBytesRef, PyInt, PyListRef, PySlice, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType,
-    PyTypeRef,
+    PositionIterInternal, PyBytes, PyBytesRef, PyInt, PyListRef, PySlice, PyStr, PyStrRef, PyTuple,
+    PyTupleRef, PyType, PyTypeRef,
 };
 use crate::{
     buffer::FormatSpec,
@@ -15,15 +15,20 @@ use crate::{
     function::Either,
     function::{FuncArgs, OptionalArg, PyComparisonValue},
     protocol::{
-        BufferDescriptor, BufferMethods, PyBuffer, PyMappingMethods, PySequenceMethods, VecBuffer,
+        BufferDescriptor, BufferMethods, PyBuffer, PyIterReturn, PyMappingMethods,
+        PySequenceMethods, VecBuffer,
     },
     sliceable::SequenceIndexOp,
-    types::{AsBuffer, AsMapping, AsSequence, Comparable, Constructor, Hashable, PyComparisonOp},
+    types::{
+        AsBuffer, AsMapping, AsSequence, Comparable, Constructor, Hashable, IterNext,
+        IterNextIterable, Iterable, PyComparisonOp, Unconstructible,
+    },
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
     TryFromBorrowedObject, TryFromObject, VirtualMachine,
 };
 use crossbeam_utils::atomic::AtomicCell;
 use itertools::Itertools;
+use rustpython_common::lock::PyMutex;
 use std::{cmp::Ordering, fmt::Debug, mem::ManuallyDrop, ops::Range};
 
 #[derive(FromArgs)]
@@ -61,7 +66,15 @@ impl Constructor for PyMemoryView {
     }
 }
 
-#[pyclass(with(Hashable, Comparable, AsBuffer, AsMapping, AsSequence, Constructor))]
+#[pyclass(with(
+    Hashable,
+    Comparable,
+    AsBuffer,
+    AsMapping,
+    AsSequence,
+    Constructor,
+    Iterable
+))]
 impl PyMemoryView {
     fn parse_format(format: &str, vm: &VirtualMachine) -> PyResult<FormatSpec> {
         FormatSpec::parse(format.as_bytes(), vm)
@@ -1038,7 +1051,8 @@ impl PyPayload for PyMemoryView {
 }
 
 pub(crate) fn init(ctx: &Context) {
-    PyMemoryView::extend_class(ctx, ctx.types.memoryview_type)
+    PyMemoryView::extend_class(ctx, ctx.types.memoryview_type);
+    PyMemoryViewIterator::extend_class(ctx, ctx.types.memoryviewiterator_type);
 }
 
 fn format_unpack(
@@ -1082,3 +1096,49 @@ fn is_equiv_format(a: &BufferDescriptor, b: &BufferDescriptor) -> bool {
 fn is_equiv_structure(a: &BufferDescriptor, b: &BufferDescriptor) -> bool {
     is_equiv_format(a, b) && is_equiv_shape(a, b)
 }
+
+impl Iterable for PyMemoryView {
+    fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
+        Ok(PyMemoryViewIterator {
+            internal: PyMutex::new(PositionIterInternal::new(zelf, 0)),
+        }
+        .into_pyobject(vm))
+    }
+}
+
+#[pyclass(module = false, name = "memory_iterator")]
+#[derive(Debug)]
+pub struct PyMemoryViewIterator {
+    internal: PyMutex<PositionIterInternal<PyRef<PyMemoryView>>>,
+}
+
+impl PyPayload for PyMemoryViewIterator {
+    fn class(vm: &VirtualMachine) -> &'static Py<PyType> {
+        vm.ctx.types.memoryviewiterator_type
+    }
+}
+
+#[pyclass(with(Constructor, IterNext))]
+impl PyMemoryViewIterator {
+    #[pymethod(magic)]
+    fn reduce(&self, vm: &VirtualMachine) -> PyTupleRef {
+        self.internal
+            .lock()
+            .builtins_iter_reduce(|x| x.clone().into(), vm)
+    }
+}
+impl Unconstructible for PyMemoryViewIterator {}
+
+impl IterNextIterable for PyMemoryViewIterator {}
+impl IterNext for PyMemoryViewIterator {
+    fn next(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
+        zelf.internal.lock().next(|mv, pos| {
+            let len = mv.len(vm)?;
+            Ok(if pos >= len {
+                PyIterReturn::StopIteration(None)
+            } else {
+                PyIterReturn::Return(mv.getitem_by_idx(pos.try_into().unwrap(), vm)?)
+            })
+        })
+    }
+}
diff --git a/vm/src/types/zoo.rs b/vm/src/types/zoo.rs
index e7d631810..d7e4eedb7 100644
--- a/vm/src/types/zoo.rs
+++ b/vm/src/types/zoo.rs
@@ -58,6 +58,7 @@ pub struct TypeZoo {
     pub dict_items_type: &'static Py<PyType>,
     pub map_type: &'static Py<PyType>,
     pub memoryview_type: &'static Py<PyType>,
+    pub memoryviewiterator_type: &'static Py<PyType>,
     pub tuple_type: &'static Py<PyType>,
     pub tuple_iterator_type: &'static Py<PyType>,
     pub set_type: &'static Py<PyType>,
@@ -160,6 +161,7 @@ impl TypeZoo {
             list_iterator_type: list::PyListIterator::init_bare_type(),
             list_reverseiterator_type: list::PyListReverseIterator::init_bare_type(),
             mappingproxy_type: mappingproxy::PyMappingProxy::init_bare_type(),
+            memoryviewiterator_type: memory::PyMemoryViewIterator::init_bare_type(),
             module_type: module::PyModule::init_bare_type(),
             namespace_type: namespace::PyNamespace::init_bare_type(),
             range_iterator_type: range::PyRangeIterator::init_bare_type(),
diff --git a/whats_left.py b/whats_left.py
index ccf55927e..ef97a602b 100755
--- a/whats_left.py
+++ b/whats_left.py
@@ -200,6 +200,7 @@ def gen_methods():
         "type(range(0).__iter__())",
         "type(str().__iter__())",
         "type(tuple().__iter__())",
+        "type(memoryview(bytearray(b'0')).__iter__())",
     ]
 
     methods = {}