Commit ae00d334 authored by John R. Lenton's avatar John R. Lenton
Browse files

Implemented set.remove

parent 4a08067c
...@@ -147,20 +147,27 @@ static void mp_set_rehash(mp_set_t *set) { ...@@ -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) { 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 (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); mp_set_rehash(set);
} else { } else {
return NULL; 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 (;;) { for (;;) {
mp_obj_t elem = set->table[pos]; mp_obj_t elem = set->table[pos];
if (elem == MP_OBJ_NULL) { if (elem == MP_OBJ_NULL) {
// not in table // 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) { if (set->used + 1 >= set->alloc) {
// not enough room in table, rehash it // not enough room in table, rehash it
mp_set_rehash(set); 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 ...@@ -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; set->table[pos] = index;
return index; return index;
} }
} else if (lookup_kind & MP_MAP_LOOKUP_FIRST) {
pos++;
} else { } else {
return MP_OBJ_NULL; 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 // found it
if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) { if (lookup_kind & MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
set->used--; set->used--;
set->table[pos] = NULL; set->table[pos] = NULL;
return elem;
} }
return elem; return elem;
} else { } else {
......
...@@ -19,9 +19,10 @@ typedef struct _mp_set_t { ...@@ -19,9 +19,10 @@ typedef struct _mp_set_t {
} mp_set_t; } mp_set_t;
typedef enum _mp_map_lookup_kind_t { typedef enum _mp_map_lookup_kind_t {
MP_MAP_LOOKUP, MP_MAP_LOOKUP, // 0
MP_MAP_LOOKUP_ADD_IF_NOT_FOUND, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND, // 1
MP_MAP_LOOKUP_REMOVE_IF_FOUND, MP_MAP_LOOKUP_REMOVE_IF_FOUND, // 2
MP_MAP_LOOKUP_FIRST = 4,
} mp_map_lookup_kind_t; } mp_map_lookup_kind_t;
int get_doubling_prime_greater_or_equal_to(int x); int get_doubling_prime_greater_or_equal_to(int x);
......
...@@ -227,6 +227,70 @@ static mp_obj_t set_isdisjoint(mp_obj_t self_in, mp_obj_t other) { ...@@ -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_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 */ /* set constructors & public C API */
...@@ -242,6 +306,10 @@ static const mp_method_t set_type_methods[] = { ...@@ -242,6 +306,10 @@ static const mp_method_t set_type_methods[] = {
{ "intersection", &set_intersect_obj }, { "intersection", &set_intersect_obj },
{ "intersection_update", &set_intersect_update_obj }, { "intersection_update", &set_intersect_update_obj },
{ "isdisjoint", &set_isdisjoint_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 { NULL, NULL }, // end-of-list sentinel
}; };
......
sets = [set(), {1}, {1, 2, 3}, {3, 4, 5}, {5, 6, 7}]
for i in sets:
for j in sets:
print(i.issubset(j))
print(i.issuperset(j))
s = {1}
print(s.pop())
try:
print(s.pop(), "!!!")
except KeyError:
pass
else:
print("Failed to raise KeyError")
s = {1}
print(s.remove(1))
print(list(s))
try:
print(s.remove(1), "!!!")
except KeyError:
pass
else:
print("failed to raise KeyError")
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment