|
| 1 | +#include <stdlib.h> |
| 2 | +#include <assert.h> |
| 3 | + |
| 4 | +#include "nlr.h" |
| 5 | +#include "misc.h" |
| 6 | +#include "mpconfig.h" |
| 7 | +#include "mpqstr.h" |
| 8 | +#include "obj.h" |
| 9 | +#include "runtime.h" |
| 10 | + |
| 11 | +typedef struct _mp_obj_map_t { |
| 12 | + mp_obj_base_t base; |
| 13 | + machine_uint_t n_iters; |
| 14 | + mp_obj_t fun; |
| 15 | + mp_obj_t iters[]; |
| 16 | +} mp_obj_map_t; |
| 17 | + |
| 18 | +static mp_obj_t map_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args) { |
| 19 | + /* NOTE: args are backwards */ |
| 20 | + if (n_args < 2) { |
| 21 | + nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "map must have at least 2 arguments")); |
| 22 | + } |
| 23 | + assert(n_args >= 2); |
| 24 | + mp_obj_map_t *o = m_new_obj_var(mp_obj_map_t, mp_obj_t, n_args - 1); |
| 25 | + o->base.type = &map_type; |
| 26 | + o->n_iters = n_args - 1; |
| 27 | + o->fun = args[n_args - 1]; |
| 28 | + for (int i = 0; i < n_args - 1; i++) { |
| 29 | + o->iters[i] = rt_getiter(args[n_args-i-2]); |
| 30 | + } |
| 31 | + return o; |
| 32 | +} |
| 33 | + |
| 34 | +static mp_obj_t map_getiter(mp_obj_t self_in) { |
| 35 | + return self_in; |
| 36 | +} |
| 37 | + |
| 38 | +static mp_obj_t map_iternext(mp_obj_t self_in) { |
| 39 | + assert(MP_OBJ_IS_TYPE(self_in, &map_type)); |
| 40 | + mp_obj_map_t *self = self_in; |
| 41 | + mp_obj_t *nextses = m_new(mp_obj_t, self->n_iters); |
| 42 | + |
| 43 | + for (int i = 0; i < self->n_iters; i++) { |
| 44 | + mp_obj_t next = rt_iternext(self->iters[i]); |
| 45 | + if (next == mp_const_stop_iteration) { |
| 46 | + m_del(mp_obj_t, nextses, self->n_iters); |
| 47 | + return mp_const_stop_iteration; |
| 48 | + } |
| 49 | + nextses[i] = next; |
| 50 | + } |
| 51 | + return rt_call_function_n(self->fun, self->n_iters, nextses); |
| 52 | +} |
| 53 | + |
| 54 | +const mp_obj_type_t map_type = { |
| 55 | + { &mp_const_type }, |
| 56 | + "map", |
| 57 | + .make_new = map_make_new, |
| 58 | + .getiter = map_getiter, |
| 59 | + .iternext = map_iternext, |
| 60 | +}; |
0 commit comments