1 files changed, 15 insertions, 7 deletions

diff --git a/py/map.c b/py/map.c
index 3dee058cd3..1ce763ab0e 100644
--- a/py/map.c
+++ b/py/map.c
@@ -147,20 +147,27 @@ static void mp_set_rehash(mp_set_t *set) {
 }
 
 mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, mp_map_lookup_kind_t lookup_kind) {
-    int hash = mp_obj_hash(index);
+    int hash;
+    int pos;
     if (set->alloc == 0) {
-        if (lookup_kind == MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
+        if (lookup_kind & MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
             mp_set_rehash(set);
         } else {
             return NULL;
         }
     }
-    int pos = hash % set->alloc;
+    if (lookup_kind & MP_MAP_LOOKUP_FIRST) {
+        hash = 0;
+        pos = 0;
+    } else {
+        hash = mp_obj_hash(index);;
+        pos = hash % set->alloc;
+    }
     for (;;) {
         mp_obj_t elem = set->table[pos];
         if (elem == MP_OBJ_NULL) {
             // not in table
-            if (lookup_kind == MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
+            if (lookup_kind & MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
                 if (set->used + 1 >= set->alloc) {
                     // not enough room in table, rehash it
                     mp_set_rehash(set);
@@ -171,15 +178,16 @@ mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, mp_map_lookup_kind_t looku
                     set->table[pos] = index;
                     return index;
                 }
+            } else if (lookup_kind & MP_MAP_LOOKUP_FIRST) {
+                pos++;
             } else {
                 return MP_OBJ_NULL;
             }
-        } else if (mp_obj_equal(elem, index)) {
+        } else if (lookup_kind & MP_MAP_LOOKUP_FIRST || mp_obj_equal(elem, index)) {
             // found it
-            if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
+            if (lookup_kind & MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
                 set->used--;
                 set->table[pos] = NULL;
-                return elem;
             }
             return elem;
         } else {