From 6b088a671ac010d78c566a86e70a460721485fa6 Mon Sep 17 00:00:00 2001 From: Paul Sokolovsky Date: Mon, 20 Jun 2016 15:50:31 +0300 Subject: [PATCH] extmod/modbtree: Implement keys(), values(), items() iterators. Each takes optional args of starting key, ending key, and flags (ending key inclusive, reverse order). --- extmod/modbtree.c | 58 +++++++++++++++++++++++++++++++++++++---------- 1 file changed, 46 insertions(+), 12 deletions(-) diff --git a/extmod/modbtree.c b/extmod/modbtree.c index 955dafa70e..b717172f46 100644 --- a/extmod/modbtree.c +++ b/extmod/modbtree.c @@ -44,6 +44,10 @@ typedef struct _mp_obj_btree_t { mp_obj_t end_key; #define FLAG_END_KEY_INCL 1 #define FLAG_DESC 2 + #define FLAG_ITER_TYPE_MASK 0xc0 + #define FLAG_ITER_KEYS 0x40 + #define FLAG_ITER_VALUES 0x80 + #define FLAG_ITER_ITEMS 0xc0 byte flags; } mp_obj_btree_t; @@ -126,17 +130,37 @@ STATIC mp_obj_t btree_seq(size_t n_args, const mp_obj_t *args) { } STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_seq_obj, 2, 4, btree_seq); -STATIC mp_obj_t btree_range(size_t n_args, const mp_obj_t *args) { +STATIC mp_obj_t btree_init_iter(size_t n_args, const mp_obj_t *args, byte type) { mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]); - self->start_key = args[1]; - self->end_key = args[2]; - self->flags = 0; - if (n_args > 3) { - self->flags = MP_OBJ_SMALL_INT_VALUE(args[3]); + self->flags = type; + self->start_key = mp_const_none; + self->end_key = mp_const_none; + if (n_args > 1) { + self->start_key = args[1]; + if (n_args > 2) { + self->end_key = args[2]; + if (n_args > 3) { + self->flags |= MP_OBJ_SMALL_INT_VALUE(args[3]); + } + } } return args[0]; } -STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_range_obj, 3, 4, btree_range); + +STATIC mp_obj_t btree_keys(size_t n_args, const mp_obj_t *args) { + return btree_init_iter(n_args, args, FLAG_ITER_KEYS); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_keys_obj, 1, 4, btree_keys); + +STATIC mp_obj_t btree_values(size_t n_args, const mp_obj_t *args) { + return btree_init_iter(n_args, args, FLAG_ITER_VALUES); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_values_obj, 1, 4, btree_values); + +STATIC mp_obj_t btree_items(size_t n_args, const mp_obj_t *args) { + return btree_init_iter(n_args, args, FLAG_ITER_ITEMS); +} +STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_items_obj, 1, 4, btree_items); STATIC mp_obj_t btree_iternext(mp_obj_t self_in) { mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in); @@ -179,10 +203,18 @@ STATIC mp_obj_t btree_iternext(mp_obj_t self_in) { } } - mp_obj_tuple_t *pair = mp_obj_new_tuple(2, NULL); - pair->items[0] = mp_obj_new_bytes(key.data, key.size); - pair->items[1] = mp_obj_new_bytes(val.data, val.size); - return pair; + switch (self->flags & FLAG_ITER_TYPE_MASK) { + case FLAG_ITER_KEYS: + return mp_obj_new_bytes(key.data, key.size); + case FLAG_ITER_VALUES: + return mp_obj_new_bytes(val.data, val.size); + default: { + mp_obj_tuple_t *pair = mp_obj_new_tuple(2, NULL); + pair->items[0] = mp_obj_new_bytes(key.data, key.size); + pair->items[1] = mp_obj_new_bytes(val.data, val.size); + return pair; + } + } } STATIC mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) { @@ -223,7 +255,9 @@ STATIC const mp_rom_map_elem_t btree_locals_dict_table[] = { { MP_ROM_QSTR(MP_QSTR_get), MP_ROM_PTR(&btree_get_obj) }, { MP_ROM_QSTR(MP_QSTR_put), MP_ROM_PTR(&btree_put_obj) }, { MP_ROM_QSTR(MP_QSTR_seq), MP_ROM_PTR(&btree_seq_obj) }, - { MP_ROM_QSTR(MP_QSTR_items), MP_ROM_PTR(&btree_range_obj) }, + { MP_ROM_QSTR(MP_QSTR_keys), MP_ROM_PTR(&btree_keys_obj) }, + { MP_ROM_QSTR(MP_QSTR_values), MP_ROM_PTR(&btree_values_obj) }, + { MP_ROM_QSTR(MP_QSTR_items), MP_ROM_PTR(&btree_items_obj) }, }; STATIC MP_DEFINE_CONST_DICT(btree_locals_dict, btree_locals_dict_table);