extmod/modbtree: Cleverly implement "for key in btree:" syntax.

I.e. make it work like btree.keys(), while still not using a separate
iterator type.

1 files changed, 24 insertions, 6 deletions

diff --git a/extmod/modbtree.c b/extmod/modbtree.c
index b717172f4..581b0a01b 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,