From 39b174e00a05dfaf4ac6f2e17cee4892a60e92c3 Mon Sep 17 00:00:00 2001
From: "John R. Lenton" <jlenton@gmail.com>
Date: Wed, 15 Jan 2014 01:10:09 +0000
Subject: [PATCH] Added map

---
 py/mpqstrraw.h            |  1 +
 py/obj.h                  |  3 +++
 py/objmap.c               | 55 +++++++++++++++++++++++++++++++++++++++
 py/py.mk                  |  1 +
 py/runtime.c              |  1 +
 tests/basics/tests/map.py |  4 +++
 6 files changed, 65 insertions(+)
 create mode 100644 py/objmap.c
 create mode 100644 tests/basics/tests/map.py

diff --git a/py/mpqstrraw.h b/py/mpqstrraw.h
index 8d681c9a49..bbca1bb43e 100644
--- a/py/mpqstrraw.h
+++ b/py/mpqstrraw.h
@@ -50,6 +50,7 @@ Q(issubclass)
 Q(iter)
 Q(len)
 Q(list)
+Q(map)
 Q(max)
 Q(min)
 Q(next)
diff --git a/py/obj.h b/py/obj.h
index 00d152ec4b..7779c47101 100644
--- a/py/obj.h
+++ b/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;
 
diff --git a/py/objmap.c b/py/objmap.c
new file mode 100644
index 0000000000..2e07d76b7b
--- /dev/null
+++ b/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,
+};
diff --git a/py/py.mk b/py/py.mk
index 275c92d126..ba26fc0297 100644
--- a/py/py.mk
+++ b/py/py.mk
@@ -84,6 +84,7 @@ PY_O_BASENAME = \
 	objgenerator.o \
 	objint.o \
 	objlist.o \
+	objmap.o \
 	objmodule.o \
 	objnone.o \
 	objrange.o \
diff --git a/py/runtime.c b/py/runtime.c
index a76af3db7c..c5292dc263 100644
--- a/py/runtime.c
+++ b/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);
diff --git a/tests/basics/tests/map.py b/tests/basics/tests/map.py
new file mode 100644
index 0000000000..62dca44ede
--- /dev/null
+++ b/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))))