#include "Python.h"

#include "pycore_object.h" // _PyObject_GC_UNTRACK()

#include <stddef.h> // offsetof()

static PyObject _dummy_struct;

#define dummy (&_dummy_struct)

#ifndef LINEAR_PROBES

#define LINEAR_PROBES 9

#define PERTURB_SHIFT 5

static setentry *

set_lookkey(PySetObject *so, PyObject *key, Py_hash_t hash)

{

setentry *table;

setentry *entry;

size_t perturb;

size_t mask = so->mask;

size_t i = (size_t)hash & mask; /* Unsigned for defined overflow behavior */

size_t j;

int cmp;

entry = &so->table[i];

if (entry->key == NULL)

return entry;

perturb = hash;

while (1) {

if (entry->hash == hash) {

PyObject *startkey = entry->key;

/* startkey cannot be a dummy because the dummy hash field is -1 */

assert(startkey != dummy);

if (startkey == key)

return entry;

if (PyUnicode_CheckExact(startkey)

&& PyUnicode_CheckExact(key)

&& _PyUnicode_EQ(startkey, key))

return entry;

table = so->table;

Py_INCREF(startkey);

cmp = PyObject_RichCompareBool(startkey, key, Py_EQ);

Py_DECREF(startkey);

if (cmp < 0) /* unlikely */

return NULL;

if (table != so->table || entry->key != startkey) /* unlikely */

return set_lookkey(so, key, hash);

if (cmp > 0) /* likely */

return entry;

mask = so->mask; /* help avoid a register spill */

}

if (i + LINEAR_PROBES <= mask) {

for (j = 0 ; j < LINEAR_PROBES ; j++) {

entry++;

if (entry->hash == 0 && entry->key == NULL)

return entry;

if (entry->hash == hash) {

PyObject *startkey = entry->key;

assert(startkey != dummy);

if (startkey == key)

return entry;

if (PyUnicode_CheckExact(startkey)

&& PyUnicode_CheckExact(key)

&& _PyUnicode_EQ(startkey, key))

return entry;

table = so->table;

Py_INCREF(startkey);

cmp = PyObject_RichCompareBool(startkey, key, Py_EQ);

Py_DECREF(startkey);

if (cmp < 0)

return NULL;

if (table != so->table || entry->key != startkey)

return set_lookkey(so, key, hash);

if (cmp > 0)

return entry;

mask = so->mask;

}

}

}

perturb >>= PERTURB_SHIFT;

i = (i * 5 + 1 + perturb) & mask;

entry = &so->table[i];

if (entry->key == NULL)

return entry;

}

}

static int set_table_resize(PySetObject *, Py_ssize_t);

static int

set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash)

{

setentry *table;

setentry *freeslot;

setentry *entry;

size_t perturb;

size_t mask;

size_t i; /* Unsigned for defined overflow behavior */

size_t j;

int cmp;

/* Pre-increment is necessary to prevent arbitrary code in the rich

Py_INCREF(key);

restart:

mask = so->mask;

i = (size_t)hash & mask;

entry = &so->table[i];

if (entry->key == NULL)

goto found_unused;

freeslot = NULL;

perturb = hash;

while (1) {

if (entry->hash == hash) {

PyObject *startkey = entry->key;

/* startkey cannot be a dummy because the dummy hash field is -1 */

assert(startkey != dummy);

if (startkey == key)

goto found_active;

if (PyUnicode_CheckExact(startkey)

&& PyUnicode_CheckExact(key)

&& _PyUnicode_EQ(startkey, key))

goto found_active;

table = so->table;

Py_INCREF(startkey);

cmp = PyObject_RichCompareBool(startkey, key, Py_EQ);

Py_DECREF(startkey);

if (cmp > 0) /* likely */

goto found_active;

if (cmp < 0)

goto comparison_error;

/* Continuing the search from the current entry only makes

if (table != so->table || entry->key != startkey)

goto restart;

mask = so->mask; /* help avoid a register spill */

}

else if (entry->hash == -1)

freeslot = entry;

if (i + LINEAR_PROBES <= mask) {

for (j = 0 ; j < LINEAR_PROBES ; j++) {

entry++;

if (entry->hash == 0 && entry->key == NULL)

goto found_unused_or_dummy;

if (entry->hash == hash) {

PyObject *startkey = entry->key;

assert(startkey != dummy);

if (startkey == key)

goto found_active;

if (PyUnicode_CheckExact(startkey)

&& PyUnicode_CheckExact(key)

&& _PyUnicode_EQ(startkey, key))

goto found_active;

table = so->table;

Py_INCREF(startkey);

cmp = PyObject_RichCompareBool(startkey, key, Py_EQ);

Py_DECREF(startkey);

if (cmp > 0)

goto found_active;

if (cmp < 0)

goto comparison_error;

if (table != so->table || entry->key != startkey)

goto restart;

mask = so->mask;

}

else if (entry->hash == -1)

freeslot = entry;

}

}

perturb >>= PERTURB_SHIFT;

i = (i * 5 + 1 + perturb) & mask;

entry = &so->table[i];

if (entry->key == NULL)

goto found_unused_or_dummy;

}

found_unused_or_dummy:

if (freeslot == NULL)

goto found_unused;

so->used++;

freeslot->key = key;

freeslot->hash = hash;

return 0;

found_unused:

so->fill++;

so->used++;

entry->key = key;

entry->hash = hash;

if ((size_t)so->fill*5 < mask*3)

return 0;

return set_table_resize(so, so->used>50000 ? so->used*2 : so->used*4);

found_active:

Py_DECREF(key);

return 0;

comparison_error:

Py_DECREF(key);

return -1;

}

static void

set_insert_clean(setentry *table, size_t mask, PyObject *key, Py_hash_t hash)

{

setentry *entry;

size_t perturb = hash;

size_t i = (size_t)hash & mask;

size_t j;

while (1) {

entry = &table[i];

if (entry->key == NULL)

goto found_null;

if (i + LINEAR_PROBES <= mask) {

for (j = 0; j < LINEAR_PROBES; j++) {

entry++;

if (entry->key == NULL)

goto found_null;

}

}

perturb >>= PERTURB_SHIFT;

i = (i * 5 + 1 + perturb) & mask;

}

found_null:

entry->key = key;

entry->hash = hash;

}

static int

set_table_resize(PySetObject *so, Py_ssize_t minused)

{

setentry *oldtable, *newtable, *entry;

Py_ssize_t oldmask = so->mask;

size_t newmask;

int is_oldtable_malloced;

setentry small_copy[PySet_MINSIZE];

assert(minused >= 0);

/* Find the smallest table size > minused. */

/* XXX speed-up with intrinsics */

size_t newsize = PySet_MINSIZE;

while (newsize <= (size_t)minused) {

newsize <<= 1; // The largest possible value is PY_SSIZE_T_MAX + 1.

}

/* Get space for a new table. */

oldtable = so->table;

assert(oldtable != NULL);

is_oldtable_malloced = oldtable != so->smalltable;

if (newsize == PySet_MINSIZE) {

/* A large table is shrinking, or we can't get any smaller. */

newtable = so->smalltable;

if (newtable == oldtable) {

if (so->fill == so->used) {

/* No dummies, so no point doing anything. */

return 0;

}

/* We're not going to resize it, but rebuild the

assert(so->fill > so->used);

memcpy(small_copy, oldtable, sizeof(small_copy));

oldtable = small_copy;

}

}

else {

newtable = PyMem_NEW(setentry, newsize);

if (newtable == NULL) {

PyErr_NoMemory();

return -1;

}

}

/* Make the set empty, using the new table. */

assert(newtable != oldtable);

memset(newtable, 0, sizeof(setentry) * newsize);

so->mask = newsize - 1;

so->table = newtable;

/* Copy the data over; this is refcount-neutral for active entries;

newmask = (size_t)so->mask;

if (so->fill == so->used) {

for (entry = oldtable; entry <= oldtable + oldmask; entry++) {

if (entry->key != NULL) {

set_insert_clean(newtable, newmask, entry->key, entry->hash);

}

}

} else {

so->fill = so->used;

for (entry = oldtable; entry <= oldtable + oldmask; entry++) {

if (entry->key != NULL && entry->key != dummy) {

set_insert_clean(newtable, newmask, entry->key, entry->hash);

}

}

}

if (is_oldtable_malloced)

PyMem_DEL(oldtable);

return 0;

}

static int

set_contains_entry(PySetObject *so, PyObject *key, Py_hash_t hash)

{

setentry *entry;

entry = set_lookkey(so, key, hash);

if (entry != NULL)

return entry->key != NULL;

return -1;

}

#define DISCARD_NOTFOUND 0

#define DISCARD_FOUND 1

static int

set_discard_entry(PySetObject *so, PyObject *key, Py_hash_t hash)

{

setentry *entry;

PyObject *old_key;

entry = set_lookkey(so, key, hash);

if (entry == NULL)

return -1;

if (entry->key == NULL)

return DISCARD_NOTFOUND;

old_key = entry->key;

entry->key = dummy;

entry->hash = -1;

so->used--;

Py_DECREF(old_key);

return DISCARD_FOUND;

}

static int

set_add_key(PySetObject *so, PyObject *key)

{

Py_hash_t hash;

if (!PyUnicode_CheckExact(key) ||

(hash = ((PyASCIIObject *) key)->hash) == -1) {

hash = PyObject_Hash(key);

if (hash == -1)

return -1;

}

return set_add_entry(so, key, hash);

}

static int

set_contains_key(PySetObject *so, PyObject *key)

{

Py_hash_t hash;

if (!PyUnicode_CheckExact(key) ||

(hash = ((PyASCIIObject *) key)->hash) == -1) {

hash = PyObject_Hash(key);

if (hash == -1)

return -1;

}

return set_contains_entry(so, key, hash);

}

static int

set_discard_key(PySetObject *so, PyObject *key)

{

Py_hash_t hash;

if (!PyUnicode_CheckExact(key) ||

(hash = ((PyASCIIObject *) key)->hash) == -1) {

hash = PyObject_Hash(key);

if (hash == -1)

return -1;

}

return set_discard_entry(so, key, hash);

}

static void

set_empty_to_minsize(PySetObject *so)

{

memset(so->smalltable, 0, sizeof(so->smalltable));

so->fill = 0;

so->used = 0;

so->mask = PySet_MINSIZE - 1;

so->table = so->smalltable;

so->hash = -1;

}

static int

set_clear_internal(PySetObject *so)

{

setentry *entry;

setentry *table = so->table;

Py_ssize_t fill = so->fill;

Py_ssize_t used = so->used;

int table_is_malloced = table != so->smalltable;

setentry small_copy[PySet_MINSIZE];

assert (PyAnySet_Check(so));

assert(table != NULL);

/* This is delicate. During the process of clearing the set,

if (table_is_malloced)

set_empty_to_minsize(so);

else if (fill > 0) {

/* It's a small table with something that needs to be cleared.

memcpy(small_copy, table, sizeof(small_copy));

table = small_copy;

set_empty_to_minsize(so);

}

/* else it's a small table that's already empty */

/* Now we can finally clear things. If C had refcounts, we could

for (entry = table; used > 0; entry++) {

if (entry->key && entry->key != dummy) {

used--;

Py_DECREF(entry->key);

}

}

if (table_is_malloced)

PyMem_DEL(table);

return 0;

}

static int

set_next(PySetObject *so, Py_ssize_t *pos_ptr, setentry **entry_ptr)

{

Py_ssize_t i;

Py_ssize_t mask;

setentry *entry;

assert (PyAnySet_Check(so));

i = *pos_ptr;

assert(i >= 0);

mask = so->mask;

entry = &so->table[i];

while (i <= mask && (entry->key == NULL || entry->key == dummy)) {

i++;

entry++;

}

*pos_ptr = i+1;

if (i > mask)

return 0;

assert(entry != NULL);

*entry_ptr = entry;

return 1;

}

static void

set_dealloc(PySetObject *so)

{

setentry *entry;

Py_ssize_t used = so->used;

/* bpo-31095: UnTrack is needed before calling any callbacks */

PyObject_GC_UnTrack(so);

Py_TRASHCAN_BEGIN(so, set_dealloc)

if (so->weakreflist != NULL)

PyObject_ClearWeakRefs((PyObject *) so);

for (entry = so->table; used > 0; entry++) {

if (entry->key && entry->key != dummy) {

used--;

Py_DECREF(entry->key);

}

}

if (so->table != so->smalltable)

PyMem_DEL(so->table);

Py_TYPE(so)->tp_free(so);

Py_TRASHCAN_END

}

static PyObject *

set_repr(PySetObject *so)

{

PyObject *result=NULL, *keys, *listrepr, *tmp;

int status = Py_ReprEnter((PyObject*)so);

if (status != 0) {

if (status < 0)

return NULL;

return PyUnicode_FromFormat("%s(...)", Py_TYPE(so)->tp_name);

}

/* shortcut for the empty set */

if (!so->used) {

Py_ReprLeave((PyObject*)so);

return PyUnicode_FromFormat("%s()", Py_TYPE(so)->tp_name);

}

keys = PySequence_List((PyObject *)so);

if (keys == NULL)

goto done;

/* repr(keys)[1:-1] */

listrepr = PyObject_Repr(keys);

Py_DECREF(keys);

if (listrepr == NULL)

goto done;

tmp = PyUnicode_Substring(listrepr, 1, PyUnicode_GET_LENGTH(listrepr)-1);

Py_DECREF(listrepr);

if (tmp == NULL)

goto done;

listrepr = tmp;

if (!Py_IS_TYPE(so, &PySet_Type))

result = PyUnicode_FromFormat("%s({%U})",

Py_TYPE(so)->tp_name,

listrepr);

else

result = PyUnicode_FromFormat("{%U}", listrepr);

Py_DECREF(listrepr);

done:

Py_ReprLeave((PyObject*)so);

return result;

}

static Py_ssize_t

set_len(PyObject *so)

{

return ((PySetObject *)so)->used;

}

static int

set_merge(PySetObject *so, PyObject *otherset)

{

PySetObject *other;

PyObject *key;

Py_ssize_t i;

setentry *so_entry;

setentry *other_entry;

assert (PyAnySet_Check(so));

assert (PyAnySet_Check(otherset));

other = (PySetObject*)otherset;

if (other == so || other->used == 0)

/* a.update(a) or a.update(set()); nothing to do */

return 0;

/* Do one big resize at the start, rather than

if ((so->fill + other->used)*5 >= so->mask*3) {

if (set_table_resize(so, (so->used + other->used)*2) != 0)

return -1;

}

so_entry = so->table;

other_entry = other->table;

/* If our table is empty, and both tables have the same size, and

if (so->fill == 0 && so->mask == other->mask && other->fill == other->used) {

for (i = 0; i <= other->mask; i++, so_entry++, other_entry++) {

key = other_entry->key;

if (key != NULL) {

assert(so_entry->key == NULL);

Py_INCREF(key);

so_entry->key = key;

so_entry->hash = other_entry->hash;

}

}

so->fill = other->fill;

so->used = other->used;

return 0;

}

/* If our table is empty, we can use set_insert_clean() */

if (so->fill == 0) {

setentry *newtable = so->table;

size_t newmask = (size_t)so->mask;

so->fill = other->used;

so->used = other->used;

for (i = other->mask + 1; i > 0 ; i--, other_entry++) {

key = other_entry->key;

if (key != NULL && key != dummy) {

Py_INCREF(key);

set_insert_clean(newtable, newmask, key, other_entry->hash);

}

}

return 0;

}

/* We can't assure there are no duplicates, so do normal insertions */

for (i = 0; i <= other->mask; i++) {

other_entry = &other->table[i];

key = other_entry->key;

if (key != NULL && key != dummy) {

if (set_add_entry(so, key, other_entry->hash))

return -1;

}

}

return 0;

}

static PyObject *

set_pop(PySetObject *so, PyObject *Py_UNUSED(ignored))

{

/* Make sure the search finger is in bounds */

setentry *entry = so->table + (so->finger & so->mask);

setentry *limit = so->table + so->mask;

PyObject *key;

if (so->used == 0) {

PyErr_SetString(PyExc_KeyError, "pop from an empty set");

return NULL;

}

while (entry->key == NULL || entry->key==dummy) {

entry++;

if (entry > limit)

entry = so->table;

}

key = entry->key;

entry->key = dummy;

entry->hash = -1;

so->used--;

so->finger = entry - so->table + 1; /* next place to start */

return key;

}

PyDoc_STRVAR(pop_doc, "Remove and return an arbitrary set element.\n\

Raises KeyError if the set is empty.");

static int

set_traverse(PySetObject *so, visitproc visit, void *arg)

{

Py_ssize_t pos = 0;

setentry *entry;

while (set_next(so, &pos, &entry))

Py_VISIT(entry->key);

return 0;

}

static Py_uhash_t

_shuffle_bits(Py_uhash_t h)

{

return ((h ^ 89869747UL) ^ (h << 16)) * 3644798167UL;

}

static Py_hash_t

frozenset_hash(PyObject *self)

{

PySetObject *so = (PySetObject *)self;

Py_uhash_t hash = 0;

setentry *entry;

if (so->hash != -1)

return so->hash;

/* Xor-in shuffled bits from every entry's hash field because xor is

for (entry = so->table; entry <= &so->table[so->mask]; entry++)

hash ^= _shuffle_bits(entry->hash);

/* Remove the effect of an odd number of NULL entries */

if ((so->mask + 1 - so->fill) & 1)

hash ^= _shuffle_bits(0);

/* Remove the effect of an odd number of dummy entries */

if ((so->fill - so->used) & 1)

hash ^= _shuffle_bits(-1);

/* Factor in the number of active entries */

hash ^= ((Py_uhash_t)PySet_GET_SIZE(self) + 1) * 1927868237UL;

/* Disperse patterns arising in nested frozensets */

hash ^= (hash >> 11) ^ (hash >> 25);

hash = hash * 69069U + 907133923UL;

/* -1 is reserved as an error code */

if (hash == (Py_uhash_t)-1)

hash = 590923713UL;

so->hash = hash;

return hash;

}

typedef struct {

PyObject_HEAD

PySetObject *si_set; /* Set to NULL when iterator is exhausted */

Py_ssize_t si_used;

Py_ssize_t si_pos;

Py_ssize_t len;

} setiterobject;

static void

setiter_dealloc(setiterobject *si)

{

/* bpo-31095: UnTrack is needed before calling any callbacks */

_PyObject_GC_UNTRACK(si);

Py_XDECREF(si->si_set);

PyObject_GC_Del(si);

}

static int

setiter_traverse(setiterobject *si, visitproc visit, void *arg)

{

Py_VISIT(si->si_set);

return 0;

}

static PyObject *

setiter_len(setiterobject *si, PyObject *Py_UNUSED(ignored))

{

Py_ssize_t len = 0;

if (si->si_set != NULL && si->si_used == si->si_set->used)

len = si->len;

return PyLong_FromSsize_t(len);

}

PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");

static PyObject *setiter_iternext(setiterobject *si);

static PyObject *

setiter_reduce(setiterobject *si, PyObject *Py_UNUSED(ignored))

{

_Py_IDENTIFIER(iter);

/* copy the iterator state */

setiterobject tmp = *si;

Py_XINCREF(tmp.si_set);

/* iterate the temporary into a list */

PyObject *list = PySequence_List((PyObject*)&tmp);

Py_XDECREF(tmp.si_set);

if (list == NULL) {

return NULL;

}

return Py_BuildValue("N(N)", _PyEval_GetBuiltinId(&PyId_iter), list);

}

PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");

static PyMethodDef setiter_methods[] = {

{"__length_hint__", (PyCFunction)setiter_len, METH_NOARGS, length_hint_doc},

{"__reduce__", (PyCFunction)setiter_reduce, METH_NOARGS, reduce_doc},

{NULL, NULL} /* sentinel */

};

static PyObject *setiter_iternext(setiterobject *si)

{

PyObject *key;

Py_ssize_t i, mask;

setentry *entry;

PySetObject *so = si->si_set;

if (so == NULL)

return NULL;

assert (PyAnySet_Check(so));

if (si->si_used != so->used) {

PyErr_SetString(PyExc_RuntimeError,

"Set changed size during iteration");

si->si_used = -1; /* Make this state sticky */

return NULL;

}

i = si->si_pos;

assert(i>=0);

entry = so->table;

mask = so->mask;

while (i <= mask && (entry[i].key == NULL || entry[i].key == dummy))

i++;

si->si_pos = i+1;

if (i > mask)

goto fail;

si->len--;

key = entry[i].key;

Py_INCREF(key);

return key;

fail:

si->si_set = NULL;

Py_DECREF(so);

return NULL;

}

PyTypeObject PySetIter_Type = {

PyVarObject_HEAD_INIT(&PyType_Type, 0)

"set_iterator", /* tp_name */

sizeof(setiterobject), /* tp_basicsize */

0, /* tp_itemsize */

/* methods */

(destructor)setiter_dealloc, /* tp_dealloc */

0, /* tp_vectorcall_offset */

0, /* tp_getattr */

0, /* tp_setattr */

0, /* tp_as_async */

0, /* tp_repr */

0, /* tp_as_number */

0, /* tp_as_sequence */

0, /* tp_as_mapping */

0, /* tp_hash */

0, /* tp_call */

0, /* tp_str */

PyObject_GenericGetAttr, /* tp_getattro */

0, /* tp_setattro */

0, /* tp_as_buffer */

Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */

0, /* tp_doc */

(traverseproc)setiter_traverse, /* tp_traverse */

0, /* tp_clear */

0, /* tp_richcompare */

0, /* tp_weaklistoffset */

PyObject_SelfIter, /* tp_iter */

(iternextfunc)setiter_iternext, /* tp_iternext */

setiter_methods, /* tp_methods */

0,

};

static PyObject *

set_iter(PySetObject *so)

{

setiterobject *si = PyObject_GC_New(setiterobject, &PySetIter_Type);

if (si == NULL)

return NULL;

Py_INCREF(so);

si->si_set = so;

si->si_used = so->used;

si->si_pos = 0;

si->len = so->used;

_PyObject_GC_TRACK(si);

return (PyObject *)si;

}

static int

set_update_internal(PySetObject *so, PyObject *other)

{

PyObject *key, *it;

if (PyAnySet_Check(other))

return set_merge(so, other);

if (PyDict_CheckExact(other)) {

PyObject *value;

Py_ssize_t pos = 0;

Py_hash_t hash;

Py_ssize_t dictsize = PyDict_GET_SIZE(other);

/* Do one big resize at the start, rather than

if (dictsize < 0)

return -1;

if ((so->fill + dictsize)*5 >= so->mask*3) {

if (set_table_resize(so, (so->used + dictsize)*2) != 0)

return -1;

}

while (_PyDict_Next(other, &pos, &key, &value, &hash)) {

if (set_add_entry(so, key, hash))

return -1;

}

return 0;

}

it = PyObject_GetIter(other);

if (it == NULL)

return -1;

while ((key = PyIter_Next(it)) != NULL) {

if (set_add_key(so, key)) {

Py_DECREF(it);

Py_DECREF(key);

return -1;

}

Py_DECREF(key);

}

Py_DECREF(it);

if (PyErr_Occurred())

return -1;

return 0;

}