Skip to content

Commit 1d7fb2f

Browse files
chipacachipaca
committed
Implemented set.clear
1 parent 19b14d3 commit 1d7fb2f

4 files changed

Lines changed: 47 additions & 12 deletions

File tree

‎py/map.c‎

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -132,9 +132,29 @@ void mp_set_init(mp_set_t *set, int n) {
132132
set->table = m_new0(mp_obj_t, set->alloc);
133133
}
134134

135+
static void mp_set_rehash(mp_set_t *set) {
136+
int old_alloc = set->alloc;
137+
mp_obj_t *old_table = set->table;
138+
set->alloc = get_doubling_prime_greater_or_equal_to(set->alloc + 1);
139+
set->used = 0;
140+
set->table = m_new0(mp_obj_t, set->alloc);
141+
for (int i = 0; i < old_alloc; i++) {
142+
if (old_table[i] != NULL) {
143+
mp_set_lookup(set, old_table[i], true);
144+
}
145+
}
146+
m_del(mp_obj_t, old_table, old_alloc);
147+
}
148+
135149
mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, bool add_if_not_found) {
136150
int hash = mp_obj_hash(index);
137-
assert(set->alloc); /* FIXME: if alloc is ever 0 when doing a lookup, this'll fail: */
151+
if (set->alloc == 0) {
152+
if (add_if_not_found) {
153+
mp_set_rehash(set);
154+
} else {
155+
return NULL;
156+
}
157+
}
138158
int pos = hash % set->alloc;
139159
for (;;) {
140160
mp_obj_t elem = set->table[pos];
@@ -143,17 +163,7 @@ mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, bool add_if_not_found) {
143163
if (add_if_not_found) {
144164
if (set->used + 1 >= set->alloc) {
145165
// not enough room in table, rehash it
146-
int old_alloc = set->alloc;
147-
mp_obj_t *old_table = set->table;
148-
set->alloc = get_doubling_prime_greater_or_equal_to(set->alloc + 1);
149-
set->used = 0;
150-
set->table = m_new(mp_obj_t, set->alloc);
151-
for (int i = 0; i < old_alloc; i++) {
152-
if (old_table[i] != NULL) {
153-
mp_set_lookup(set, old_table[i], true);
154-
}
155-
}
156-
m_del(mp_obj_t, old_table, old_alloc);
166+
mp_set_rehash(set);
157167
// restart the search for the new element
158168
pos = hash % set->alloc;
159169
} else {
@@ -173,3 +183,13 @@ mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, bool add_if_not_found) {
173183
}
174184
}
175185
}
186+
187+
void mp_set_clear(mp_set_t *set) {
188+
set->used = 0;
189+
machine_uint_t a = set->alloc;
190+
set->alloc = 0;
191+
set->table = m_renew(mp_obj_t, set->table, a, set->alloc);
192+
for (uint i=0; i<set->alloc; i++) {
193+
set->table[i] = NULL;
194+
}
195+
}

‎py/map.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,3 +32,4 @@ void mp_map_clear(mp_map_t *map);
3232

3333
void mp_set_init(mp_set_t *set, int n);
3434
mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, bool add_if_not_found);
35+
void mp_set_clear(mp_set_t *set);

‎py/objset.c‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,13 +104,24 @@ static mp_obj_t set_add(mp_obj_t self_in, mp_obj_t item) {
104104
}
105105
static MP_DEFINE_CONST_FUN_OBJ_2(set_add_obj, set_add);
106106

107+
static mp_obj_t set_clear(mp_obj_t self_in) {
108+
assert(MP_OBJ_IS_TYPE(self_in, &set_type));
109+
mp_obj_set_t *self = self_in;
110+
111+
mp_set_clear(&self->set);
112+
113+
return mp_const_none;
114+
}
115+
static MP_DEFINE_CONST_FUN_OBJ_1(set_clear_obj, set_clear);
116+
107117

108118
/******************************************************************************/
109119
/* set constructors & public C API */
110120

111121

112122
static const mp_method_t set_type_methods[] = {
113123
{ "add", &set_add_obj },
124+
{ "clear", &set_clear_obj },
114125
{ NULL, NULL }, // end-of-list sentinel
115126
};
116127

‎tests/basics/tests/set_clear.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
s = {1, 2, 3, 4}
2+
print(s.clear())
3+
print(list(s))

0 commit comments

Comments
 (0)