py: Make dict use a bit less RAM when iterating; properly del values.

Heap RAM was being allocated to print dicts and do some other types of
iterating.  Now these iterations use 1 word of state on the stack.

Deleting elements from a dict was not allowing the value to be reclaimed
by the GC.  This is now fixed.
This commit is contained in:
Damien George 2014-09-17 15:53:03 +01:00
parent 1d7fb82f0a
commit 8a9b999f1c

View File

@ -38,17 +38,32 @@
#include "runtime.h" #include "runtime.h"
#include "builtin.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); 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) { 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; mp_obj_dict_t *self = self_in;
bool first = true; bool first = true;
print(env, "{"); 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; 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) { if (!first) {
print(env, ", "); 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; return mp_const_false;
} }
mp_uint_t size = o->map.alloc; mp_uint_t cur = 0;
mp_map_t *map = &o->map; mp_map_elem_t *next = NULL;
while ((next = dict_iter_next(o, &cur)) != NULL) {
for (mp_uint_t i = 0; i < size; i++) { mp_map_elem_t *elem = mp_map_lookup(&rhs->map, next->key, MP_MAP_LOOKUP);
if (MP_MAP_SLOT_IS_FILLED(map, i)) { if (elem == NULL || !mp_obj_equal(next->value, elem->value)) {
mp_map_elem_t *elem = mp_map_lookup(&rhs->map, map->table[i].key, MP_MAP_LOOKUP); return mp_const_false;
if (elem == NULL || !mp_obj_equal(map->table[i].value, elem->value)) {
return mp_const_false;
}
} }
} }
return mp_const_true; return mp_const_true;
@ -158,28 +170,14 @@ typedef struct _mp_obj_dict_it_t {
mp_uint_t cur; mp_uint_t cur;
} mp_obj_dict_it_t; } 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_obj_dict_it_t *self = self_in;
mp_uint_t max = self->dict->map.alloc; mp_map_elem_t *next = dict_iter_next(self->dict, &self->cur);
mp_map_t *map = &self->dict->map;
for (mp_uint_t i = self->cur; i < max; i++) { if (next == NULL) {
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 {
return MP_OBJ_STOP_ITERATION; 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, .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); mp_obj_dict_it_t *o = m_new_obj(mp_obj_dict_it_t);
o->base.type = &mp_type_dict_it; o->base.type = &mp_type_dict_it;
o->dict = dict; o->dict = o_in;
o->cur = cur; o->cur = 0;
return o; return o;
} }
STATIC mp_obj_t dict_getiter(mp_obj_t o_in) {
return mp_obj_new_dict_iterator(o_in, 0);
}
/******************************************************************************/ /******************************************************************************/
/* dict methods */ /* 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) { 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_map_elem_t *elem = mp_map_lookup(self, key, lookup_kind);
mp_obj_t value; mp_obj_t value;
if (elem == NULL || elem->value == NULL) { if (elem == NULL || elem->value == MP_OBJ_NULL) {
if (deflt == NULL) { if (deflt == MP_OBJ_NULL) {
if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) { if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
nlr_raise(mp_obj_new_exception_msg(&mp_type_KeyError, "<value>")); nlr_raise(mp_obj_new_exception_msg(&mp_type_KeyError, "<value>"));
} else { } 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 { } else {
value = deflt; value = deflt;
} }
if (lookup_kind == MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
elem->value = value;
}
} else { } else {
value = elem->value; value = elem->value;
} if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
if (lookup_kind == MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) { elem->value = MP_OBJ_NULL; // so that GC can collect the deleted value
elem->value = value; }
} }
return 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, return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
args[1], args[1],
n_args == 3 ? args[2] : NULL, n_args == 3 ? args[2] : MP_OBJ_NULL,
MP_MAP_LOOKUP); MP_MAP_LOOKUP);
} }
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get); 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, return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
args[1], args[1],
n_args == 3 ? args[2] : NULL, n_args == 3 ? args[2] : MP_OBJ_NULL,
MP_MAP_LOOKUP_REMOVE_IF_FOUND); MP_MAP_LOOKUP_REMOVE_IF_FOUND);
} }
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop); 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, return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
args[1], args[1],
n_args == 3 ? args[2] : NULL, n_args == 3 ? args[2] : MP_OBJ_NULL,
MP_MAP_LOOKUP_ADD_IF_NOT_FOUND); MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
} }
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_setdefault_obj, 2, 3, dict_setdefault); 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) { STATIC mp_obj_t dict_popitem(mp_obj_t self_in) {
assert(MP_OBJ_IS_TYPE(self_in, &mp_type_dict)); assert(MP_OBJ_IS_TYPE(self_in, &mp_type_dict));
mp_obj_dict_t *self = self_in; 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")); 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--; self->map.used--;
mp_obj_t items[] = {next->key, next->value}; mp_obj_t items[] = {next->key, next->value};
next->key = NULL; next->key = MP_OBJ_SENTINEL; // must mark key as sentinel to indicate that it was deleted
next->value = NULL; next->value = MP_OBJ_NULL;
mp_obj_t tuple = mp_obj_new_tuple(2, items); mp_obj_t tuple = mp_obj_new_tuple(2, items);
return tuple; 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)) { if (MP_OBJ_IS_TYPE(args[1], &mp_type_dict)) {
// update from other dictionary (make sure other is not self) // update from other dictionary (make sure other is not self)
if (args[1] != self) { if (args[1] != self) {
// TODO don't allocate heap object for this iterator mp_uint_t cur = 0;
mp_obj_t *dict_iter = mp_obj_new_dict_iterator(args[1], 0);
mp_map_elem_t *elem = NULL; 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; 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 { typedef struct _mp_obj_dict_view_it_t {
mp_obj_base_t base; mp_obj_base_t base;
mp_dict_view_kind_t kind; mp_dict_view_kind_t kind;
mp_obj_dict_it_t *iter; mp_obj_dict_t *dict;
mp_uint_t cur; mp_uint_t cur;
} mp_obj_dict_view_it_t; } 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) { STATIC mp_obj_t dict_view_it_iternext(mp_obj_t self_in) {
assert(MP_OBJ_IS_TYPE(self_in, &dict_view_it_type)); assert(MP_OBJ_IS_TYPE(self_in, &dict_view_it_type));
mp_obj_dict_view_it_t *self = self_in; 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) { switch (self->kind) {
case MP_DICT_VIEW_ITEMS: 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 */ assert(0); /* can't happen */
return mp_const_none; 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); mp_obj_dict_view_it_t *o = m_new_obj(mp_obj_dict_view_it_t);
o->base.type = &dict_view_it_type; o->base.type = &dict_view_it_type;
o->kind = view->kind; o->kind = view->kind;
o->iter = mp_obj_new_dict_iterator(view->dict, 0); o->dict = view->dict;
o->cur = 0;
return o; 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) { 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)); assert(MP_OBJ_IS_TYPE(self_in, &mp_type_dict));
mp_obj_dict_t *self = self_in; 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; return self_in;
} }