Added map

Author: chipaca

py/mpqstrraw.h

@@ -50,6 +50,7 @@ Q(issubclass)
 Q(iter)
 Q(len)
 Q(list)
+Q(map)
 Q(max)
 Q(min)
 Q(next)

py/obj.h

@@ -294,6 +294,9 @@ void mp_obj_list_get(mp_obj_t self_in, uint *len, mp_obj_t **items);
 void mp_obj_list_store(mp_obj_t self_in, mp_obj_t index, mp_obj_t value);
 mp_obj_t mp_obj_list_sort(mp_obj_t args, struct _mp_map_t *kwargs);
 
+// map (the python builtin, not the dict implementation detail)
+extern const mp_obj_type_t map_type;
+
 // enumerate
 extern const mp_obj_type_t enumerate_type;
 

py/objmap.c

@@ -0,0 +1,55 @@
+#include <stdlib.h>
+#include <assert.h>
+
+#include "misc.h"
+#include "mpconfig.h"
+#include "obj.h"
+#include "runtime.h"
+
+typedef struct _mp_obj_map_t {
+    mp_obj_base_t base;
+    machine_uint_t n_iters;
+    mp_obj_t fun;
+    mp_obj_t iters[];
+} mp_obj_map_t;
+
+static mp_obj_t map_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args) {
+    /* NOTE: args are backwards */
+    mp_obj_map_t *o = m_new_obj_var(mp_obj_map_t, mp_obj_t, n_args - 1);
+    assert(n_args >= 2);
+    o->base.type = &map_type;
+    o->n_iters = n_args - 1;
+    o->fun = args[n_args - 1];
+    for (int i = 0; i < n_args - 1; i++) {
+        o->iters[i] = rt_getiter(args[n_args-i-2]);
+    }
+    return o;
+}
+
+static mp_obj_t map_getiter(mp_obj_t self_in) {
+    return self_in;
+}
+
+static mp_obj_t map_iternext(mp_obj_t self_in) {
+    assert(MP_OBJ_IS_TYPE(self_in, &map_type));
+    mp_obj_map_t *self = self_in;
+    mp_obj_t *nextses = m_new(mp_obj_t, self->n_iters);
+
+    for (int i = 0; i < self->n_iters; i++) {
+        mp_obj_t next = rt_iternext(self->iters[i]);
+        if (next == mp_const_stop_iteration) {
+            m_del(mp_obj_t, nextses, self->n_iters);
+            return mp_const_stop_iteration;
+        }
+        nextses[i] = next;
+    }
+    return rt_call_function_n(self->fun, self->n_iters, nextses);
+}
+
+const mp_obj_type_t map_type = {
+    { &mp_const_type },
+    "map",
+    .make_new = map_make_new,
+    .getiter = map_getiter,
+    .iternext = map_iternext,
+};

py/py.mk

@@ -84,6 +84,7 @@ PY_O_BASENAME = \
 	objgenerator.o \
 	objint.o \
 	objlist.o \
+	objmap.o \
 	objmodule.o \
 	objnone.o \
 	objrange.o \

py/runtime.c

@@ -111,6 +111,7 @@ void rt_init(void) {
 #endif
     mp_map_add_qstr(&map_builtins, MP_QSTR_int, (mp_obj_t)&int_type);
     mp_map_add_qstr(&map_builtins, MP_QSTR_list, (mp_obj_t)&list_type);
+    mp_map_add_qstr(&map_builtins, MP_QSTR_map, (mp_obj_t)&map_type);
     mp_map_add_qstr(&map_builtins, MP_QSTR_set, (mp_obj_t)&set_type);
     mp_map_add_qstr(&map_builtins, MP_QSTR_tuple, (mp_obj_t)&tuple_type);
     mp_map_add_qstr(&map_builtins, MP_QSTR_type, (mp_obj_t)&mp_const_type);

tests/basics/tests/map.py

@@ -0,0 +1,4 @@
+print(list(map(lambda x: x & 1, range(-3, 4))))
+print(list(map(abs, range(-3, 4))))
+print(list(map(set, [[i] for i in range(-3, 4)])))
+print(list(map(pow, range(4), range(4))))