Implemented set.remove

3 files changed, 87 insertions, 10 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 {
diff --git a/py/map.h b/py/map.h
index 5b00e53f83..2db0ac3ebc 100644
--- a/py/map.h
+++ b/py/map.h
@@ -19,9 +19,10 @@ typedef struct _mp_set_t {
 } mp_set_t;
 
 typedef enum _mp_map_lookup_kind_t {
-    MP_MAP_LOOKUP,
-    MP_MAP_LOOKUP_ADD_IF_NOT_FOUND,
-    MP_MAP_LOOKUP_REMOVE_IF_FOUND,
+    MP_MAP_LOOKUP,                    // 0
+    MP_MAP_LOOKUP_ADD_IF_NOT_FOUND,   // 1
+    MP_MAP_LOOKUP_REMOVE_IF_FOUND,    // 2
+    MP_MAP_LOOKUP_FIRST = 4,
 } mp_map_lookup_kind_t;
 
 int get_doubling_prime_greater_or_equal_to(int x);
diff --git a/py/objset.c b/py/objset.c
index afc426b3b0..75bc7efb99 100644
--- a/py/objset.c
+++ b/py/objset.c
@@ -227,6 +227,70 @@ static mp_obj_t set_isdisjoint(mp_obj_t self_in, mp_obj_t other) {
 }
 static MP_DEFINE_CONST_FUN_OBJ_2(set_isdisjoint_obj, set_isdisjoint);
 
+static mp_obj_t set_issubset(mp_obj_t self_in, mp_obj_t other_in) {
+    mp_obj_set_t *self;
+    bool cleanup_self = false;
+    if (MP_OBJ_IS_TYPE(self_in, &set_type)) {
+        self = self_in;
+    } else {
+        self = set_make_new(NULL, 1, &self_in);
+        cleanup_self = true;
+    }
+
+    mp_obj_set_t *other;
+    bool cleanup_other = false;
+    if (MP_OBJ_IS_TYPE(other_in, &set_type)) {
+        other = other_in;
+    } else {
+        other = set_make_new(NULL, 1, &other_in);
+        cleanup_other = true;
+    }
+    mp_obj_t iter = set_getiter(self);
+    mp_obj_t next;
+    mp_obj_t out = mp_const_true;
+    while ((next = set_it_iternext(iter)) != mp_const_stop_iteration) {
+        if (!mp_set_lookup(&other->set, next, MP_MAP_LOOKUP)) {
+            out = mp_const_false;
+            break;
+        }
+    }
+    if (cleanup_self) {
+        set_clear(self);
+    }
+    if (cleanup_other) {
+        set_clear(other);
+    }
+    return out;
+}
+static MP_DEFINE_CONST_FUN_OBJ_2(set_issubset_obj, set_issubset);
+
+static mp_obj_t set_issuperset(mp_obj_t self_in, mp_obj_t other_in) {
+    return set_issubset(other_in, self_in);
+}
+static MP_DEFINE_CONST_FUN_OBJ_2(set_issuperset_obj, set_issuperset);
+
+static mp_obj_t set_pop(mp_obj_t self_in) {
+    assert(MP_OBJ_IS_TYPE(self_in, &set_type));
+    mp_obj_set_t *self = self_in;
+
+    if (self->set.used == 0) {
+        nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "pop from an empty set"));
+    }
+    mp_obj_t obj = mp_set_lookup(&self->set, NULL,
+                         MP_MAP_LOOKUP_REMOVE_IF_FOUND | MP_MAP_LOOKUP_FIRST);
+    return obj;
+}
+static MP_DEFINE_CONST_FUN_OBJ_1(set_pop_obj, set_pop);
+
+static mp_obj_t set_remove(mp_obj_t self_in, mp_obj_t item) {
+    assert(MP_OBJ_IS_TYPE(self_in, &set_type));
+    mp_obj_set_t *self = self_in;
+    if (mp_set_lookup(&self->set, item, MP_MAP_LOOKUP_REMOVE_IF_FOUND) == MP_OBJ_NULL) {
+        nlr_jump(mp_obj_new_exception(MP_QSTR_KeyError));
+    }
+    return mp_const_none;
+}
+static MP_DEFINE_CONST_FUN_OBJ_2(set_remove_obj, set_remove);
 
 /******************************************************************************/
 /* set constructors & public C API                                            */
@@ -242,6 +306,10 @@ static const mp_method_t set_type_methods[] = {
     { "intersection", &set_intersect_obj },
     { "intersection_update", &set_intersect_update_obj },
     { "isdisjoint", &set_isdisjoint_obj },
+    { "issubset", &set_issubset_obj },
+    { "issuperset", &set_issuperset_obj },
+    { "pop", &set_pop_obj },
+    { "remove", &set_remove_obj },
     { NULL, NULL }, // end-of-list sentinel
 };