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.

1 files changed, 59 insertions, 63 deletions

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, "<value>"));
             } 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;
 }