py: allow presizing dicts

Author: romanz

py/map.c

@@ -129,9 +129,16 @@ void mp_map_clear(mp_map_t *map) {
 }
 
 STATIC void mp_map_rehash(mp_map_t *map) {
-    size_t old_alloc = map->alloc;
     size_t new_alloc = get_hash_alloc_greater_or_equal_to(map->alloc + 1);
-    DEBUG_printf("mp_map_rehash(%p): " UINT_FMT " -> " UINT_FMT "\n", map, old_alloc, new_alloc);
+    DEBUG_printf("mp_map_rehash(%p): " UINT_FMT " -> " UINT_FMT "\n", map, map->alloc, new_alloc);
+    mp_map_presize(map, new_alloc);
+}
+
+void mp_map_presize(mp_map_t *map, size_t new_alloc) {
+    size_t old_alloc = map->alloc;
+    if (new_alloc < old_alloc) {
+        mp_raise_msg_varg(&mp_type_ValueError, "Map capacity (" UINT_FMT ") must not decrease: " UINT_FMT, old_alloc, new_alloc);
+    }
     mp_map_elem_t *old_table = map->table;
     mp_map_elem_t *new_table = m_new0(mp_map_elem_t, new_alloc);
     // If we reach this point, table resizing succeeded, now we can edit the old map.

py/obj.h

@@ -466,6 +466,7 @@ void mp_map_free(mp_map_t *map);
 mp_map_elem_t *mp_map_lookup(mp_map_t *map, mp_obj_t index, mp_map_lookup_kind_t lookup_kind);
 void mp_map_clear(mp_map_t *map);
 void mp_map_dump(mp_map_t *map);
+void mp_map_presize(mp_map_t *map, size_t n);
 
 // Underlying set implementation (not set object)
 
@@ -946,6 +947,7 @@ typedef struct _mp_obj_dict_t {
 } mp_obj_dict_t;
 mp_obj_t mp_obj_dict_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args);
 void mp_obj_dict_init(mp_obj_dict_t *dict, size_t n_args);
+void mp_obj_dict_presize(mp_obj_dict_t *dict, size_t n_args);
 size_t mp_obj_dict_len(mp_obj_t self_in);
 mp_obj_t mp_obj_dict_get(mp_obj_t self_in, mp_obj_t index);
 mp_obj_t mp_obj_dict_store(mp_obj_t self_in, mp_obj_t key, mp_obj_t value);

py/objdict.c

@@ -617,6 +617,10 @@ void mp_obj_dict_init(mp_obj_dict_t *dict, size_t n_args) {
     mp_map_init(&dict->map, n_args);
 }
 
+void mp_obj_dict_presize(mp_obj_dict_t *dict, size_t n_args) {
+    mp_map_presize(&dict->map, n_args);
+}
+
 mp_obj_t mp_obj_new_dict(size_t n_args) {
     mp_obj_dict_t *o = m_new_obj(mp_obj_dict_t);
     mp_obj_dict_init(o, n_args);