map

RustPython · cthulahoops · Mar 9, 2019 · Mar 9, 2019 · Mar 9, 2019 · Mar 9, 2019
commit 303e19d7139c58586a330b72921d556f0b986833
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::<PyMap>()
     {
         let next_objs = iterators
             .iter()

diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -1512,10 +1512,6 @@ pub enum PyObjectPayload {
         predicate: PyObjectRef,
         iterator: PyObjectRef,
     },
-    MapIterator {
-        mapper: PyObjectRef,
-        iterators: Vec<PyObjectRef>,
-    },
     ZipIterator {
         iterators: Vec<PyObjectRef>,
     },
@@ -1573,7 +1569,6 @@ impl fmt::Debug for PyObjectPayload {
             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"),