mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, mp_map_lookup_kind_t lookup_kind) {
int hash;
int pos;
if (set->alloc == 0) {
if (lookup_kind & MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
mp_set_rehash(set);
} else {
return NULL;
}
}
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 (set->used + 1 >= set->alloc) {
// not enough room in table, rehash it
mp_set_rehash(set);
// restart the search for the new element
pos = hash % set->alloc;
} else {
set->used += 1;
set->table[pos] = index;
return index;
}
} else if (lookup_kind & MP_MAP_LOOKUP_FIRST) {
pos++;
} else {
return MP_OBJ_NULL;
}
} else if ((lookup_kind & MP_MAP_LOOKUP_FIRST) || mp_obj_equal(elem, index)) {
// found it
if (lookup_kind & MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
set->used--;
set->table[pos] = NULL;
}
return elem;
} else {
// not yet found, keep searching in this table
pos = (pos + 1) % set->alloc;
}
}
}
mp_obj_t mp_set_lookup(mp_set_t *set, mp_obj_t index, mp_map_lookup_kind_t lookup_kind) {
// Note: lookup_kind can be MP_MAP_LOOKUP_ADD_IF_NOT_FOUND_OR_REMOVE_IF_FOUND which
// is handled by using bitwise operations.
if (set->alloc == 0) {
if (lookup_kind & MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
mp_set_rehash(set);
} else {
return MP_OBJ_NULL;
}
}
mp_uint_t hash = MP_OBJ_SMALL_INT_VALUE(mp_unary_op(MP_UNARY_OP_HASH, index));
size_t pos = hash % set->alloc;
size_t start_pos = pos;
mp_obj_t *avail_slot = NULL;
for (;;) {
mp_obj_t elem = set->table[pos];
if (elem == MP_OBJ_NULL) {
// found NULL slot, so index is not in table
if (lookup_kind & MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
if (avail_slot == NULL) {
avail_slot = &set->table[pos];
}
set->used++;
*avail_slot = index;
return index;
} else {
return MP_OBJ_NULL;
}
} else if (elem == MP_OBJ_SENTINEL) {
// found deleted slot, remember for later
if (avail_slot == NULL) {
avail_slot = &set->table[pos];
}
} else if (mp_obj_equal(elem, index)) {
// found index
if (lookup_kind & MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
// delete element
set->used--;
if (set->table[(pos + 1) % set->alloc] == MP_OBJ_NULL) {
// optimisation if next slot is empty
set->table[pos] = MP_OBJ_NULL;
} else {
set->table[pos] = MP_OBJ_SENTINEL;
}
}
return elem;
}
// not yet found, keep searching in this table
pos = (pos + 1) % set->alloc;
if (pos == start_pos) {
// search got back to starting position, so index is not in table
if (lookup_kind & MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
if (avail_slot != NULL) {
// there was an available slot, so use that
set->used++;
*avail_slot = index;
return index;
} else {
// not enough room in table, rehash it
mp_set_rehash(set);
// restart the search for the new element
start_pos = pos = hash % set->alloc;
}
} else {
return MP_OBJ_NULL;
}
}
}
}
