diff --git a/py/map.c b/py/map.c index 01209c9b74..b88181989d 100644 --- a/py/map.c +++ b/py/map.c @@ -9,7 +9,8 @@ #include "map.h" // approximatelly doubling primes; made with Mathematica command: Table[Prime[Floor[(1.7)^n]], {n, 3, 24}] -static int doubling_primes[] = {7, 19, 43, 89, 179, 347, 647, 1229, 2297, 4243, 7829, 14347, 26017, 47149, 84947, 152443, 273253, 488399, 869927, 1547173, 2745121, 4861607}; +// prefixed with zero for the empty case. +static int doubling_primes[] = {0, 7, 19, 43, 89, 179, 347, 647, 1229, 2297, 4243, 7829, 14347, 26017, 47149, 84947, 152443, 273253, 488399, 869927, 1547173, 2745121, 4861607}; int get_doubling_prime_greater_or_equal_to(int x) { for (int i = 0; i < sizeof(doubling_primes) / sizeof(int); i++) { @@ -38,7 +39,32 @@ mp_map_t *mp_map_new(mp_map_kind_t kind, int n) { return map; } -mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found) { +void mp_map_clear(mp_map_t *map) { + map->used = 0; + machine_uint_t a = map->alloc; + map->alloc = 0; + map->table = m_renew(mp_map_elem_t, map->table, a, map->alloc); + mp_map_elem_t nul = {NULL, NULL}; + for (uint i=0; ialloc; i++) { + map->table[i] = nul; + } +} + +static void mp_map_rehash (mp_map_t *map) { + int old_alloc = map->alloc; + mp_map_elem_t *old_table = map->table; + map->alloc = get_doubling_prime_greater_or_equal_to(map->alloc + 1); + map->used = 0; + map->table = m_new0(mp_map_elem_t, map->alloc); + for (int i = 0; i < old_alloc; i++) { + if (old_table[i].key != NULL) { + mp_map_lookup_helper(map, old_table[i].key, true, false)->value = old_table[i].value; + } + } + m_del(mp_map_elem_t, old_table, old_alloc); +} + +mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found, bool remove_if_found) { bool is_map_mp_obj = (map->kind == MP_MAP_OBJ); machine_uint_t hash; if (is_map_mp_obj) { @@ -46,6 +72,13 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n } else { hash = (machine_uint_t)index; } + if (map->alloc == 0) { + if (add_if_not_found) { + mp_map_rehash(map); + } else { + return NULL; + } + } uint pos = hash % map->alloc; for (;;) { mp_map_elem_t *elem = &map->table[pos]; @@ -54,17 +87,7 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n if (add_if_not_found) { if (map->used + 1 >= map->alloc) { // not enough room in table, rehash it - int old_alloc = map->alloc; - mp_map_elem_t *old_table = map->table; - map->alloc = get_doubling_prime_greater_or_equal_to(map->alloc + 1); - map->used = 0; - map->table = m_new0(mp_map_elem_t, map->alloc); - for (int i = 0; i < old_alloc; i++) { - if (old_table[i].key != NULL) { - mp_map_lookup_helper(map, old_table[i].key, true)->value = old_table[i].value; - } - } - m_del(mp_map_elem_t, old_table, old_alloc); + mp_map_rehash(map); // restart the search for the new element pos = hash % map->alloc; } else { @@ -82,6 +105,16 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n elem->key = index; } */ + if (remove_if_found) { + map->used--; + /* this leaks this memory (but see dict_get_helper) */ + mp_map_elem_t *retval = m_new(mp_map_elem_t, 1); + retval->key = elem->key; + retval->value = elem->value; + elem->key = NULL; + elem->value = NULL; + return retval; + } return elem; } else { // not yet found, keep searching in this table @@ -92,7 +125,7 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n mp_map_elem_t* mp_qstr_map_lookup(mp_map_t *map, qstr index, bool add_if_not_found) { mp_obj_t o = (mp_obj_t)(machine_uint_t)index; - return mp_map_lookup_helper(map, o, add_if_not_found); + return mp_map_lookup_helper(map, o, add_if_not_found, false); } /******************************************************************************/ @@ -106,6 +139,7 @@ void mp_set_init(mp_set_t *set, int n) { mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, bool add_if_not_found) { int hash = mp_obj_hash(index); + assert(set->alloc); /* FIXME: if alloc is ever 0 when doing a lookup, this'll fail: */ int pos = hash % set->alloc; for (;;) { mp_obj_t elem = set->table[pos]; diff --git a/py/map.h b/py/map.h index f8ca886aa4..f0a57758c5 100644 --- a/py/map.h +++ b/py/map.h @@ -26,8 +26,9 @@ typedef struct _mp_set_t { int get_doubling_prime_greater_or_equal_to(int x); void mp_map_init(mp_map_t *map, mp_map_kind_t kind, int n); mp_map_t *mp_map_new(mp_map_kind_t kind, int n); -mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found); +mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found, bool remove_if_found); mp_map_elem_t* mp_qstr_map_lookup(mp_map_t *map, qstr index, bool add_if_not_found); +void mp_map_clear(mp_map_t *map); void mp_set_init(mp_set_t *set, int n); mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, bool add_if_not_found); diff --git a/py/objdict.c b/py/objdict.c index 66a442fbaa..c36d4ccb03 100644 --- a/py/objdict.c +++ b/py/objdict.c @@ -17,20 +17,23 @@ typedef struct _mp_obj_dict_t { mp_map_t map; } mp_obj_dict_t; +static mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, int cur); +static mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in); + static void dict_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in) { mp_obj_dict_t *self = self_in; bool first = true; print(env, "{"); - for (int i = 0; i < self->map.alloc; i++) { - if (self->map.table[i].key != NULL) { - if (!first) { - print(env, ", "); - } - first = false; - mp_obj_print_helper(print, env, self->map.table[i].key); - print(env, ": "); - mp_obj_print_helper(print, env, self->map.table[i].value); + mp_obj_t *dict_iter = mp_obj_new_dict_iterator(self, 0); + mp_map_elem_t *next = NULL; + while ((next = dict_it_iternext_elem(dict_iter)) != NULL) { + if (!first) { + print(env, ", "); } + first = false; + mp_obj_print_helper(print, env, next->key); + print(env, ": "); + mp_obj_print_helper(print, env, next->value); } print(env, "}"); } @@ -47,7 +50,7 @@ static mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { case RT_BINARY_OP_SUBSCR: { // dict load - mp_map_elem_t *elem = mp_map_lookup_helper(&o->map, rhs_in, false); + mp_map_elem_t *elem = mp_map_lookup_helper(&o->map, rhs_in, false, false); if (elem == NULL) { nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "")); } else { @@ -60,12 +63,211 @@ static mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { } } + +/******************************************************************************/ +/* dict iterator */ + +typedef struct _mp_obj_dict_it_t { + mp_obj_base_t base; + mp_obj_dict_t *dict; + machine_uint_t cur; +} mp_obj_dict_it_t; + +static mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in) { + mp_obj_dict_it_t *self = self_in; + machine_uint_t max = self->dict->map.alloc; + mp_map_elem_t *table = self->dict->map.table; + + for (int i = self->cur; i < max; i++) { + if (table[i].key != NULL) { + self->cur = i + 1; + return &(table[i]); + } + } + + return NULL; +} + +mp_obj_t dict_it_iternext(mp_obj_t self_in) { + mp_map_elem_t *next = dict_it_iternext_elem(self_in); + + if (next != NULL) { + return next->key; + } else { + return mp_const_stop_iteration; + } +} + +static const mp_obj_type_t dict_it_type = { + { &mp_const_type }, + "dict_iterator", + .iternext = dict_it_iternext, +}; + +static mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, int cur) { + mp_obj_dict_it_t *o = m_new_obj(mp_obj_dict_it_t); + o->base.type = &dict_it_type; + o->dict = dict; + o->cur = cur; + return o; +} + +static mp_obj_t dict_getiter(mp_obj_t o_in) { + return mp_obj_new_dict_iterator(o_in, 0); +} + +/******************************************************************************/ +/* dict methods */ + +static mp_obj_t dict_clear(mp_obj_t self_in) { + assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); + mp_obj_dict_t *self = self_in; + + mp_map_clear(&self->map); + + return mp_const_none; +} +static MP_DEFINE_CONST_FUN_OBJ_1(dict_clear_obj, dict_clear); + +static mp_obj_t dict_copy(mp_obj_t self_in) { + assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); + mp_obj_dict_t *self = self_in; + mp_obj_dict_t *other = mp_obj_new_dict(self->map.alloc); + other->map.used = self->map.used; + memcpy(other->map.table, self->map.table, self->map.alloc * sizeof(mp_map_elem_t)); + return other; +} +static MP_DEFINE_CONST_FUN_OBJ_1(dict_copy_obj, dict_copy); + +static mp_obj_t dict_get_helper(mp_map_t *self, mp_obj_t key, mp_obj_t deflt, bool pop, bool set) { + mp_map_elem_t *elem = mp_map_lookup_helper(self, key, set, pop); + mp_obj_t value; + if (elem == NULL || elem->value == NULL) { + if (deflt == NULL) { + if (pop) { + nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "")); + } else { + value = mp_const_none; + } + } else { + value = deflt; + } + } else { + value = elem->value; + if (pop) { + /* catch the leak (from mp_map_lookup_helper) */ + m_free(elem, sizeof(mp_map_elem_t)); + } + } + if (set) { + elem->value = value; + } + return value; +} + +static mp_obj_t dict_get(int n_args, const mp_obj_t *args) { + assert(2 <= n_args && n_args <= 3); + assert(MP_OBJ_IS_TYPE(args[0], &dict_type)); + + return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, + args[1], + n_args == 3 ? args[2] : NULL, + false, false); +} +static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get); + +static mp_obj_t dict_pop(int n_args, const mp_obj_t *args) { + assert(2 <= n_args && n_args <= 3); + assert(MP_OBJ_IS_TYPE(args[0], &dict_type)); + + return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, + args[1], + n_args == 3 ? args[2] : NULL, + true, false); +} +static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop); + + +static mp_obj_t dict_setdefault(int n_args, const mp_obj_t *args) { + assert(2 <= n_args && n_args <= 3); + assert(MP_OBJ_IS_TYPE(args[0], &dict_type)); + + return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, + args[1], + n_args == 3 ? args[2] : NULL, + false, true); +} +static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_setdefault_obj, 2, 3, dict_setdefault); + + +static mp_obj_t dict_popitem(mp_obj_t self_in) { + assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); + mp_obj_dict_t *self = self_in; + if (self->map.used == 0) { + nlr_jump(mp_obj_new_exception_msg(MP_QSTR_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; + mp_obj_t tuple = mp_obj_new_tuple(2, items); + + return tuple; +} +static MP_DEFINE_CONST_FUN_OBJ_1(dict_popitem_obj, dict_popitem); + +static mp_obj_t dict_update(mp_obj_t self_in, mp_obj_t iterable) { + assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); + mp_obj_dict_t *self = self_in; + /* TODO: check for the "keys" method */ + mp_obj_t iter = rt_getiter(iterable); + mp_obj_t next = NULL; + while ((next = rt_iternext(iter)) != mp_const_stop_iteration) { + mp_obj_t inneriter = rt_getiter(next); + mp_obj_t key = rt_iternext(inneriter); + mp_obj_t value = rt_iternext(inneriter); + mp_obj_t stop = rt_iternext(inneriter); + if (key == mp_const_stop_iteration + || value == mp_const_stop_iteration + || stop != mp_const_stop_iteration) { + nlr_jump(mp_obj_new_exception_msg( + MP_QSTR_ValueError, + "dictionary update sequence has the wrong length")); + } else { + mp_map_lookup_helper(&self->map, key, true, false)->value = value; + } + } + + return mp_const_none; +} +static MP_DEFINE_CONST_FUN_OBJ_2(dict_update_obj, dict_update); + + +/******************************************************************************/ +/* dict constructors & etc */ + +static const mp_method_t dict_type_methods[] = { + { "clear", &dict_clear_obj }, + { "copy", &dict_copy_obj }, + { "get", &dict_get_obj }, + { "pop", &dict_pop_obj }, + { "popitem", &dict_popitem_obj }, + { "setdefault", &dict_setdefault_obj }, + { "update", &dict_update_obj }, + { NULL, NULL }, // end-of-list sentinel +}; + const mp_obj_type_t dict_type = { { &mp_const_type }, "dict", .print = dict_print, .make_new = dict_make_new, .binary_op = dict_binary_op, + .getiter = dict_getiter, + .methods = dict_type_methods, }; mp_obj_t mp_obj_new_dict(int n_args) { @@ -76,19 +278,12 @@ mp_obj_t mp_obj_new_dict(int n_args) { } uint mp_obj_dict_len(mp_obj_t self_in) { - mp_obj_dict_t *self = self_in; - uint len = 0; - for (int i = 0; i < self->map.alloc; i++) { - if (self->map.table[i].key != NULL) { - len += 1; - } - } - return len; + return ((mp_obj_dict_t *)self_in)->map.used; } mp_obj_t mp_obj_dict_store(mp_obj_t self_in, mp_obj_t key, mp_obj_t value) { assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); mp_obj_dict_t *self = self_in; - mp_map_lookup_helper(&self->map, key, true)->value = value; + mp_map_lookup_helper(&self->map, key, true, false)->value = value; return self_in; } diff --git a/py/runtime.c b/py/runtime.c index 962d539c04..5c476d9ecf 100644 --- a/py/runtime.c +++ b/py/runtime.c @@ -879,13 +879,13 @@ mp_obj_t rt_getiter(mp_obj_t o_in) { mp_obj_t rt_iternext(mp_obj_t o_in) { if (MP_OBJ_IS_SMALL_INT(o_in)) { - nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "? 'int' object is not iterable")); + nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "'int' object is not an iterator")); } else { mp_obj_base_t *o = o_in; if (o->type->iternext != NULL) { return o->type->iternext(o_in); } else { - nlr_jump(mp_obj_new_exception_msg_1_arg(MP_QSTR_TypeError, "? '%s' object is not iterable", o->type->name)); + nlr_jump(mp_obj_new_exception_msg_1_arg(MP_QSTR_TypeError, "'%s' object is not an iterator", o->type->name)); } } } diff --git a/tests/basics/tests/dict_clear.py b/tests/basics/tests/dict_clear.py new file mode 100644 index 0000000000..6be2778bea --- /dev/null +++ b/tests/basics/tests/dict_clear.py @@ -0,0 +1,6 @@ +d = {1: 2, 3: 4} +print(d) +d.clear() +print(d) +d[2] = 42 +print(d) diff --git a/tests/basics/tests/dict_copy.py b/tests/basics/tests/dict_copy.py new file mode 100644 index 0000000000..c3eb7ffc18 --- /dev/null +++ b/tests/basics/tests/dict_copy.py @@ -0,0 +1,5 @@ +a = {i: 2*i for i in range(1000)} +b = a.copy() +for i in range(1000): + print(i, b[i]) +print(len(b)) diff --git a/tests/basics/tests/dict_get.py b/tests/basics/tests/dict_get.py new file mode 100644 index 0000000000..fb43a45eab --- /dev/null +++ b/tests/basics/tests/dict_get.py @@ -0,0 +1,3 @@ +for d in {}, {42:2}: + print(d.get(42)) + print(d.get(42,2)) diff --git a/tests/basics/tests/dict_iterator.py b/tests/basics/tests/dict_iterator.py new file mode 100644 index 0000000000..f190e32ffd --- /dev/null +++ b/tests/basics/tests/dict_iterator.py @@ -0,0 +1,3 @@ +d = {1: 2, 3: 4} +for i in d: + print(i, d[i]) diff --git a/tests/basics/tests/dict_pop.py b/tests/basics/tests/dict_pop.py new file mode 100644 index 0000000000..602560ce9d --- /dev/null +++ b/tests/basics/tests/dict_pop.py @@ -0,0 +1,12 @@ +d = {1: 2, 3: 4} +print(d.pop(3), d) +print(d) +print(d.pop(1, 42), d) +print(d.pop(1, 42), d) +print(d.pop(1, None), d) +try: + print(d.pop(1), "!!!",) +except KeyError: + print("Raised KeyError") +else: + print("Did not rise KeyError!") diff --git a/tests/basics/tests/dict_popitem.py b/tests/basics/tests/dict_popitem.py new file mode 100644 index 0000000000..184735cde6 --- /dev/null +++ b/tests/basics/tests/dict_popitem.py @@ -0,0 +1,11 @@ +d={1:2,3:4} +print(d.popitem()) +print(d) +print(d.popitem()) +print(d) +try: + print(d.popitem(), "!!!",) +except KeyError: + print("Raised KeyError") +else: + print("Did not raise KeyError") diff --git a/tests/basics/tests/dict_setdefault.py b/tests/basics/tests/dict_setdefault.py new file mode 100644 index 0000000000..57d0ba4518 --- /dev/null +++ b/tests/basics/tests/dict_setdefault.py @@ -0,0 +1,13 @@ +d = {} +print(d.setdefault(1)) +print(d.setdefault(1)) +print(d.setdefault(5, 42)) +print(d.setdefault(5, 1)) +print(d[1]) +print(d[5]) +d.pop(5) +print(d.setdefault(5, 1)) +print(d[1]) +print(d[5]) + + diff --git a/tests/basics/tests/dict_update.py b/tests/basics/tests/dict_update.py new file mode 100644 index 0000000000..e7ae0bd965 --- /dev/null +++ b/tests/basics/tests/dict_update.py @@ -0,0 +1,10 @@ +d = {1:2, 3:4} +print(d) +d.update(["ab"]) +print(d[1]) +print(d[3]) +print(d["a"]) +print(len(d)) +d.update([(1,4)]) +print(d[1]) +print(len(d))