Merge pull request #634 from RustPython/joey/map-zip-filter-enumerate-to-any

cthulahoops · web-flow · commit 56bb6d23cc40 · 2019-03-09T06:42:04.000Z
Convert map, zip, filter, and enumerate to Any payload
diff --git a/vm/src/obj/objenumerate.rs b/vm/src/obj/objenumerate.rs
@@ -3,11 +3,26 @@ use std::ops::AddAssign;
 
 use super::objint;
 use super::objiter;
-use crate::pyobject::{PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyResult, TypeProtocol};
+use crate::pyobject::{
+    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectPayload2, PyObjectRef, PyResult,
+    TypeProtocol,
+};
 use crate::vm::VirtualMachine;
 use num_bigint::BigInt;
 use num_traits::Zero;
 
+#[derive(Debug)]
+pub struct PyEnumerate {
+    counter: RefCell<BigInt>,
+    iterator: PyObjectRef,
+}
+
+impl PyObjectPayload2 for PyEnumerate {
+    fn required_type(ctx: &PyContext) -> PyObjectRef {
+        ctx.enumerate_type()
+    }
+}
+
 fn enumerate_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
@@ -22,9 +37,11 @@ fn enumerate_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     };
     let iterator = objiter::get_iter(vm, iterable)?;
     Ok(PyObject::new(
-        PyObjectPayload::EnumerateIterator {
-            counter: RefCell::new(counter),
-            iterator,
+        PyObjectPayload::AnyRustValue {
+            value: Box::new(PyEnumerate {
+                counter: RefCell::new(counter),
+                iterator,
+            }),
         },
         cls.clone(),
     ))
@@ -37,10 +54,10 @@ fn enumerate_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(enumerate, Some(vm.ctx.enumerate_type()))]
     );
 
-    if let PyObjectPayload::EnumerateIterator {
+    if let Some(PyEnumerate {
         ref counter,
         ref iterator,
-    } = enumerate.payload
+    }) = enumerate.payload()
     {
         let next_obj = objiter::call_next(vm, iterator)?;
         let result = vm
diff --git a/vm/src/obj/objfilter.rs b/vm/src/obj/objfilter.rs
@@ -1,10 +1,24 @@
-use super::objbool;
-use super::objiter;
 use crate::pyobject::{
-    IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyResult, TypeProtocol,
+    IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectPayload2, PyObjectRef,
+    PyResult, TypeProtocol,
 };
 use crate::vm::VirtualMachine; // Required for arg_check! to use isinstance
 
+use super::objbool;
+use super::objiter;
+
+#[derive(Debug)]
+pub struct PyFilter {
+    predicate: PyObjectRef,
+    iterator: PyObjectRef,
+}
+
+impl PyObjectPayload2 for PyFilter {
+    fn required_type(ctx: &PyContext) -> PyObjectRef {
+        ctx.filter_type()
+    }
+}
+
 fn filter_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
@@ -13,9 +27,11 @@ fn filter_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     );
     let iterator = objiter::get_iter(vm, iterable)?;
     Ok(PyObject::new(
-        PyObjectPayload::FilterIterator {
-            predicate: function.clone(),
-            iterator,
+        PyObjectPayload::AnyRustValue {
+            value: Box::new(PyFilter {
+                predicate: function.clone(),
+                iterator,
+            }),
         },
         cls.clone(),
     ))
@@ -24,10 +40,10 @@ fn filter_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 fn filter_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(filter, Some(vm.ctx.filter_type()))]);
 
-    if let PyObjectPayload::FilterIterator {
+    if let Some(PyFilter {
         ref predicate,
         ref iterator,
-    } = filter.payload
+    }) = filter.payload()
     {
         loop {
             let next_obj = objiter::call_next(vm, iterator)?;
diff --git a/vm/src/obj/objmap.rs b/vm/src/obj/objmap.rs
@@ -1,6 +1,22 @@
+use crate::pyobject::{
+    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectPayload2, PyObjectRef, PyResult,
+    TypeProtocol,
+};
+use crate::vm::VirtualMachine;
+
 use super::objiter;
-use crate::pyobject::{PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyResult, TypeProtocol};
-use crate::vm::VirtualMachine; // Required for arg_check! to use isinstance
+
+#[derive(Debug)]
+pub struct PyMap {
+    mapper: PyObjectRef,
+    iterators: Vec<PyObjectRef>,
+}
+
+impl PyObjectPayload2 for PyMap {
+    fn required_type(ctx: &PyContext) -> PyObjectRef {
+        ctx.map_type()
+    }
+}
 
 fn map_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     no_kwargs!(vm, args);
@@ -15,9 +31,11 @@ fn map_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
             .map(|iterable| objiter::get_iter(vm, iterable))
             .collect::<Result<Vec<_>, _>>()?;
         Ok(PyObject::new(
-            PyObjectPayload::MapIterator {
-                mapper: function.clone(),
-                iterators,
+            PyObjectPayload::AnyRustValue {
+                value: Box::new(PyMap {
+                    mapper: function.clone(),
+                    iterators,
+                }),
             },
             cls.clone(),
         ))
@@ -27,10 +45,10 @@ fn map_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 fn map_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(map, Some(vm.ctx.map_type()))]);
 
-    if let PyObjectPayload::MapIterator {
+    if let Some(PyMap {
         ref mapper,
         ref iterators,
-    } = map.payload
+    }) = map.payload()
     {
         let next_objs = iterators
             .iter()
diff --git a/vm/src/obj/objzip.rs b/vm/src/obj/objzip.rs
@@ -1,6 +1,21 @@
+use crate::pyobject::{
+    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectPayload2, PyObjectRef, PyResult,
+    TypeProtocol,
+};
+use crate::vm::VirtualMachine;
+
 use super::objiter;
-use crate::pyobject::{PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyResult, TypeProtocol};
-use crate::vm::VirtualMachine; // Required for arg_check! to use isinstance
+
+#[derive(Debug)]
+pub struct PyZip {
+    iterators: Vec<PyObjectRef>,
+}
+
+impl PyObjectPayload2 for PyZip {
+    fn required_type(ctx: &PyContext) -> PyObjectRef {
+        ctx.zip_type()
+    }
+}
 
 fn zip_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     no_kwargs!(vm, args);
@@ -11,15 +26,17 @@ fn zip_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         .map(|iterable| objiter::get_iter(vm, iterable))
         .collect::<Result<Vec<_>, _>>()?;
     Ok(PyObject::new(
-        PyObjectPayload::ZipIterator { iterators },
+        PyObjectPayload::AnyRustValue {
+            value: Box::new(PyZip { iterators }),
+        },
         cls.clone(),
     ))
 }
 
 fn zip_next(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(zip, Some(vm.ctx.zip_type()))]);
 
-    if let PyObjectPayload::ZipIterator { ref iterators } = zip.payload {
+    if let Some(PyZip { ref iterators }) = zip.payload() {
         if iterators.is_empty() {
             Err(objiter::new_stop_iteration(vm))
         } else {
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -1499,21 +1499,6 @@ pub enum PyObjectPayload {
         position: Cell<usize>,
         iterated_obj: PyObjectRef,
     },
-    EnumerateIterator {
-        counter: RefCell<BigInt>,
-        iterator: PyObjectRef,
-    },
-    FilterIterator {
-        predicate: PyObjectRef,
-        iterator: PyObjectRef,
-    },
-    MapIterator {
-        mapper: PyObjectRef,
-        iterators: Vec<PyObjectRef>,
-    },
-    ZipIterator {
-        iterators: Vec<PyObjectRef>,
-    },
     Slice {
         start: Option<BigInt>,
         stop: Option<BigInt>,
@@ -1562,10 +1547,6 @@ impl fmt::Debug for PyObjectPayload {
             PyObjectPayload::MemoryView { ref obj } => write!(f, "bytes/bytearray {:?}", obj),
             PyObjectPayload::WeakRef { .. } => write!(f, "weakref"),
             PyObjectPayload::Iterator { .. } => write!(f, "iterator"),
-            PyObjectPayload::EnumerateIterator { .. } => write!(f, "enumerate"),
-            PyObjectPayload::FilterIterator { .. } => write!(f, "filter"),
-            PyObjectPayload::MapIterator { .. } => write!(f, "map"),
-            PyObjectPayload::ZipIterator { .. } => write!(f, "zip"),
             PyObjectPayload::Slice { .. } => write!(f, "slice"),
             PyObjectPayload::Function { .. } => write!(f, "function"),
             PyObjectPayload::Generator { .. } => write!(f, "generator"),
