py: Fix bug in type_store_attr, trying to store to ROM.

2 files changed, 14 insertions, 5 deletions

diff --git a/py/map.h b/py/map.h
index 692ae92e38..0614af3c42 100644
--- a/py/map.h
+++ b/py/map.h
@@ -3,6 +3,11 @@ typedef struct _mp_map_elem_t {
     mp_obj_t value;
 } mp_map_elem_t;
 
+// TODO maybe have a truncated mp_map_t for fixed tables, since alloc=used
+// put alloc last in the structure, so the truncated version does not need it
+// this would save 1 ROM word for all ROM objects that have a locals_dict
+// would also need a trucated dict structure
+
 typedef struct _mp_map_t {
     machine_uint_t all_keys_are_qstrs : 1;
     machine_uint_t table_is_fixed_array : 1;
diff --git a/py/objtype.c b/py/objtype.c
index 4c709c7ab4..51bc9ca3de 100644
--- a/py/objtype.c
+++ b/py/objtype.c
@@ -325,12 +325,16 @@ STATIC bool type_store_attr(mp_obj_t self_in, qstr attr, mp_obj_t value) {
 
     if (self->locals_dict != NULL) {
         assert(MP_OBJ_IS_TYPE(self->locals_dict, &dict_type)); // Micro Python restriction, for now
-        mp_map_t *locals_map = ((void*)self->locals_dict + sizeof(mp_obj_base_t)); // XXX hack to get map object from dict object
-        mp_map_lookup(locals_map, MP_OBJ_NEW_QSTR(attr), MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = value;
-        return true;
-    } else {
-        return false;
+        mp_map_t *locals_map = mp_obj_dict_get_map(self->locals_dict);
+        mp_map_elem_t *elem = mp_map_lookup(locals_map, MP_OBJ_NEW_QSTR(attr), MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
+        // note that locals_map may be in ROM, so add will fail in that case
+        if (elem != NULL) {
+            elem->value = value;
+            return true;
+        }
     }
+
+    return false;
 }
 
 const mp_obj_type_t mp_type_type = {