diff options
| -rw-r--r-- | extmod/modbtree.c | 30 |
1 files changed, 24 insertions, 6 deletions
diff --git a/extmod/modbtree.c b/extmod/modbtree.c index b717172f46..581b0a01b8 100644 --- a/extmod/modbtree.c +++ b/extmod/modbtree.c @@ -49,6 +49,7 @@ typedef struct _mp_obj_btree_t { #define FLAG_ITER_VALUES 0x80 #define FLAG_ITER_ITEMS 0xc0 byte flags; + byte next_flags; } mp_obj_btree_t; STATIC const mp_obj_type_t btree_type; @@ -66,9 +67,9 @@ STATIC mp_obj_btree_t *btree_new(DB *db) { mp_obj_btree_t *o = m_new_obj(mp_obj_btree_t); o->base.type = &btree_type; o->db = db; - o->start_key = MP_OBJ_NULL; - o->end_key = MP_OBJ_NULL; - o->flags = 0; + o->start_key = mp_const_none; + o->end_key = mp_const_none; + o->next_flags = 0; return o; } @@ -132,7 +133,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_seq_obj, 2, 4, btree_seq); 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->flags = type; + self->next_flags = type; self->start_key = mp_const_none; self->end_key = mp_const_none; if (n_args > 1) { @@ -140,7 +141,7 @@ STATIC mp_obj_t btree_init_iter(size_t n_args, const mp_obj_t *args, byte type) if (n_args > 2) { self->end_key = args[2]; if (n_args > 3) { - self->flags |= MP_OBJ_SMALL_INT_VALUE(args[3]); + self->next_flags = type | MP_OBJ_SMALL_INT_VALUE(args[3]); } } } @@ -162,6 +163,23 @@ STATIC mp_obj_t btree_items(size_t n_args, const mp_obj_t *args) { } STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_items_obj, 1, 4, btree_items); +STATIC mp_obj_t btree_getiter(mp_obj_t self_in) { + mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in); + if (self->next_flags != 0) { + // If we're called immediately after keys(), values(), or items(), + // use their setup for iteration. + self->flags = self->next_flags; + self->next_flags = 0; + } else { + // Otherwise, iterate over all keys. + self->flags = FLAG_ITER_KEYS; + self->start_key = mp_const_none; + self->end_key = mp_const_none; + } + + return self_in; +} + STATIC mp_obj_t btree_iternext(mp_obj_t self_in) { mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in); DBT key, val; @@ -267,7 +285,7 @@ STATIC const mp_obj_type_t btree_type = { // Save on qstr's, reuse same as for module .name = MP_QSTR_btree, .print = btree_print, - .getiter = mp_identity, + .getiter = btree_getiter, .iternext = btree_iternext, .subscr = btree_subscr, .locals_dict = (void*)&btree_locals_dict, |