diff --git a/py/objdict.c b/py/objdict.c index d25cb4b377..a378bf6db8 100644 --- a/py/objdict.c +++ b/py/objdict.c @@ -38,17 +38,32 @@ #include "runtime.h" #include "builtin.h" -STATIC mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, mp_uint_t cur); -STATIC mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in); STATIC mp_obj_t dict_update(mp_uint_t n_args, const mp_obj_t *args, mp_map_t *kwargs); +// This is a helper function to iterate through a dictionary. The state of +// the iteration is held in *cur and should be initialised with zero for the +// first call. Will return NULL when no more elements are available. +STATIC mp_map_elem_t *dict_iter_next(mp_obj_dict_t *dict, mp_uint_t *cur) { + mp_uint_t max = dict->map.alloc; + mp_map_t *map = &dict->map; + + for (mp_uint_t i = *cur; i < max; i++) { + if (MP_MAP_SLOT_IS_FILLED(map, i)) { + *cur = i + 1; + return &(map->table[i]); + } + } + + return NULL; +} + STATIC void dict_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) { mp_obj_dict_t *self = self_in; bool first = true; print(env, "{"); - mp_obj_t *dict_iter = mp_obj_new_dict_iterator(self, 0); + mp_uint_t cur = 0; mp_map_elem_t *next = NULL; - while ((next = dict_it_iternext_elem(dict_iter)) != MP_OBJ_STOP_ITERATION) { + while ((next = dict_iter_next(self, &cur)) != NULL) { if (!first) { print(env, ", "); } @@ -94,15 +109,12 @@ STATIC mp_obj_t dict_binary_op(mp_uint_t op, mp_obj_t lhs_in, mp_obj_t rhs_in) { return mp_const_false; } - mp_uint_t size = o->map.alloc; - mp_map_t *map = &o->map; - - for (mp_uint_t i = 0; i < size; i++) { - if (MP_MAP_SLOT_IS_FILLED(map, i)) { - mp_map_elem_t *elem = mp_map_lookup(&rhs->map, map->table[i].key, MP_MAP_LOOKUP); - if (elem == NULL || !mp_obj_equal(map->table[i].value, elem->value)) { - return mp_const_false; - } + mp_uint_t cur = 0; + mp_map_elem_t *next = NULL; + while ((next = dict_iter_next(o, &cur)) != NULL) { + mp_map_elem_t *elem = mp_map_lookup(&rhs->map, next->key, MP_MAP_LOOKUP); + if (elem == NULL || !mp_obj_equal(next->value, elem->value)) { + return mp_const_false; } } return mp_const_true; @@ -158,28 +170,14 @@ typedef struct _mp_obj_dict_it_t { mp_uint_t cur; } mp_obj_dict_it_t; -STATIC mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in) { +STATIC mp_obj_t dict_it_iternext(mp_obj_t self_in) { mp_obj_dict_it_t *self = self_in; - mp_uint_t max = self->dict->map.alloc; - mp_map_t *map = &self->dict->map; + mp_map_elem_t *next = dict_iter_next(self->dict, &self->cur); - for (mp_uint_t i = self->cur; i < max; i++) { - if (MP_MAP_SLOT_IS_FILLED(map, i)) { - self->cur = i + 1; - return &(map->table[i]); - } - } - - return MP_OBJ_STOP_ITERATION; -} - -mp_obj_t dict_it_iternext(mp_obj_t self_in) { - mp_map_elem_t *next = dict_it_iternext_elem(self_in); - - if (next != MP_OBJ_STOP_ITERATION) { - return next->key; - } else { + if (next == NULL) { return MP_OBJ_STOP_ITERATION; + } else { + return next->key; } } @@ -190,18 +188,14 @@ STATIC const mp_obj_type_t mp_type_dict_it = { .iternext = dict_it_iternext, }; -STATIC mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, mp_uint_t cur) { +STATIC mp_obj_t dict_getiter(mp_obj_t o_in) { mp_obj_dict_it_t *o = m_new_obj(mp_obj_dict_it_t); o->base.type = &mp_type_dict_it; - o->dict = dict; - o->cur = cur; + o->dict = o_in; + o->cur = 0; return o; } -STATIC mp_obj_t dict_getiter(mp_obj_t o_in) { - return mp_obj_new_dict_iterator(o_in, 0); -} - /******************************************************************************/ /* dict methods */ @@ -259,8 +253,8 @@ STATIC MP_DEFINE_CONST_CLASSMETHOD_OBJ(dict_fromkeys_obj, (const mp_obj_t)&dict_ STATIC mp_obj_t dict_get_helper(mp_map_t *self, mp_obj_t key, mp_obj_t deflt, mp_map_lookup_kind_t lookup_kind) { mp_map_elem_t *elem = mp_map_lookup(self, key, lookup_kind); mp_obj_t value; - if (elem == NULL || elem->value == NULL) { - if (deflt == NULL) { + if (elem == NULL || elem->value == MP_OBJ_NULL) { + if (deflt == MP_OBJ_NULL) { if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) { nlr_raise(mp_obj_new_exception_msg(&mp_type_KeyError, "")); } else { @@ -269,11 +263,14 @@ STATIC mp_obj_t dict_get_helper(mp_map_t *self, mp_obj_t key, mp_obj_t deflt, mp } else { value = deflt; } + if (lookup_kind == MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) { + elem->value = value; + } } else { value = elem->value; - } - if (lookup_kind == MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) { - elem->value = value; + if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) { + elem->value = MP_OBJ_NULL; // so that GC can collect the deleted value + } } return value; } @@ -284,7 +281,7 @@ STATIC mp_obj_t dict_get(mp_uint_t n_args, const mp_obj_t *args) { return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, args[1], - n_args == 3 ? args[2] : NULL, + n_args == 3 ? args[2] : MP_OBJ_NULL, MP_MAP_LOOKUP); } STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get); @@ -295,7 +292,7 @@ STATIC mp_obj_t dict_pop(mp_uint_t n_args, const mp_obj_t *args) { return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, args[1], - n_args == 3 ? args[2] : NULL, + n_args == 3 ? args[2] : MP_OBJ_NULL, MP_MAP_LOOKUP_REMOVE_IF_FOUND); } STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop); @@ -307,7 +304,7 @@ STATIC mp_obj_t dict_setdefault(mp_uint_t n_args, const mp_obj_t *args) { return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, args[1], - n_args == 3 ? args[2] : NULL, + n_args == 3 ? args[2] : MP_OBJ_NULL, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND); } STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_setdefault_obj, 2, 3, dict_setdefault); @@ -316,16 +313,15 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_setdefault_obj, 2, 3, dict_setde STATIC mp_obj_t dict_popitem(mp_obj_t self_in) { assert(MP_OBJ_IS_TYPE(self_in, &mp_type_dict)); mp_obj_dict_t *self = self_in; - if (self->map.used == 0) { + mp_uint_t cur = 0; + mp_map_elem_t *next = dict_iter_next(self, &cur); + if (next == NULL) { nlr_raise(mp_obj_new_exception_msg(&mp_type_KeyError, "popitem(): dictionary is empty")); } - mp_obj_dict_it_t *iter = mp_obj_new_dict_iterator(self, 0); - - mp_map_elem_t *next = dict_it_iternext_elem(iter); self->map.used--; mp_obj_t items[] = {next->key, next->value}; - next->key = NULL; - next->value = NULL; + next->key = MP_OBJ_SENTINEL; // must mark key as sentinel to indicate that it was deleted + next->value = MP_OBJ_NULL; mp_obj_t tuple = mp_obj_new_tuple(2, items); return tuple; @@ -344,10 +340,9 @@ STATIC mp_obj_t dict_update(mp_uint_t n_args, const mp_obj_t *args, mp_map_t *kw if (MP_OBJ_IS_TYPE(args[1], &mp_type_dict)) { // update from other dictionary (make sure other is not self) if (args[1] != self) { - // TODO don't allocate heap object for this iterator - mp_obj_t *dict_iter = mp_obj_new_dict_iterator(args[1], 0); + mp_uint_t cur = 0; mp_map_elem_t *elem = NULL; - while ((elem = dict_it_iternext_elem(dict_iter)) != MP_OBJ_STOP_ITERATION) { + while ((elem = dict_iter_next((mp_obj_dict_t*)args[1], &cur)) != NULL) { mp_map_lookup(&self->map, elem->key, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = elem->value; } } @@ -402,7 +397,7 @@ STATIC char *mp_dict_view_names[] = {"dict_items", "dict_keys", "dict_values"}; typedef struct _mp_obj_dict_view_it_t { mp_obj_base_t base; mp_dict_view_kind_t kind; - mp_obj_dict_it_t *iter; + mp_obj_dict_t *dict; mp_uint_t cur; } mp_obj_dict_view_it_t; @@ -415,9 +410,11 @@ typedef struct _mp_obj_dict_view_t { STATIC mp_obj_t dict_view_it_iternext(mp_obj_t self_in) { assert(MP_OBJ_IS_TYPE(self_in, &dict_view_it_type)); mp_obj_dict_view_it_t *self = self_in; - mp_map_elem_t *next = dict_it_iternext_elem(self->iter); + mp_map_elem_t *next = dict_iter_next(self->dict, &self->cur); - if (next != MP_OBJ_STOP_ITERATION) { + if (next == NULL) { + return MP_OBJ_STOP_ITERATION; + } else { switch (self->kind) { case MP_DICT_VIEW_ITEMS: { @@ -432,8 +429,6 @@ STATIC mp_obj_t dict_view_it_iternext(mp_obj_t self_in) { assert(0); /* can't happen */ return mp_const_none; } - } else { - return MP_OBJ_STOP_ITERATION; } } @@ -450,7 +445,8 @@ STATIC mp_obj_t dict_view_getiter(mp_obj_t view_in) { mp_obj_dict_view_it_t *o = m_new_obj(mp_obj_dict_view_it_t); o->base.type = &dict_view_it_type; o->kind = view->kind; - o->iter = mp_obj_new_dict_iterator(view->dict, 0); + o->dict = view->dict; + o->cur = 0; return o; } @@ -580,7 +576,7 @@ mp_obj_t mp_obj_dict_store(mp_obj_t self_in, mp_obj_t key, mp_obj_t value) { mp_obj_t mp_obj_dict_delete(mp_obj_t self_in, mp_obj_t key) { assert(MP_OBJ_IS_TYPE(self_in, &mp_type_dict)); mp_obj_dict_t *self = self_in; - dict_get_helper(&self->map, key, NULL, MP_MAP_LOOKUP_REMOVE_IF_FOUND); + dict_get_helper(&self->map, key, MP_OBJ_NULL, MP_MAP_LOOKUP_REMOVE_IF_FOUND); return self_in; }