#include "Python.h"

#include "pycore_ceval.h" // _PyEval_GetBuiltin()

#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION, Py_END_CRITICAL_SECTION

#include "pycore_dict.h" // _PyDict_Contains_KnownHash()

#include "pycore_modsupport.h" // _PyArg_NoKwnames()

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

#include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_LOAD_SSIZE_RELAXED()

#include "pycore_pyerrors.h" // _PyErr_SetKeyError()

#include "pycore_setobject.h" // _PySet_NextEntry() definition

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

#include "clinic/setobject.c.h"

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 = hash;

size_t mask = so->mask;

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

int probes;

int cmp;

while (1) {

entry = &so->table[i];

probes = (i + LINEAR_PROBES <= mask) ? LINEAR_PROBES: 0;

do {

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;

}

entry++;

} while (probes--);

perturb >>= PERTURB_SHIFT;

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

}

}

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 */

int probes;

int cmp;

_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);

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

Py_INCREF(key);

restart:

mask = so->mask;

i = (size_t)hash & mask;

freeslot = NULL;

perturb = hash;

while (1) {

entry = &so->table[i];

probes = (i + LINEAR_PROBES <= mask) ? LINEAR_PROBES: 0;

do {

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) {

assert (entry->key == dummy);

freeslot = entry;

}

entry++;

} while (probes--);

perturb >>= PERTURB_SHIFT;

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

}

found_unused_or_dummy:

if (freeslot == NULL)

goto found_unused;

FT_ATOMIC_STORE_SSIZE_RELAXED(so->used, so->used + 1);

freeslot->key = key;

freeslot->hash = hash;

return 0;

found_unused:

so->fill++;

FT_ATOMIC_STORE_SSIZE_RELAXED(so->used, so->used + 1);

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_Free(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;

FT_ATOMIC_STORE_SSIZE_RELAXED(so->used, so->used - 1);

Py_DECREF(old_key);

return DISCARD_FOUND;

}

static int

set_add_key(PySetObject *so, PyObject *key)

{

Py_hash_t hash = _PyObject_HashFast(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 = _PyObject_HashFast(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 = _PyObject_HashFast(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;

FT_ATOMIC_STORE_SSIZE_RELAXED(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_Free(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_Free(so->table);

Py_TYPE(so)->tp_free(so);

Py_TRASHCAN_END

}

static PyObject *

set_repr_lock_held(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 (!PySet_CheckExact(so))

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 PyObject *

set_repr(PySetObject *so)

{

PyObject *result;

Py_BEGIN_CRITICAL_SECTION(so);

result = set_repr_lock_held(so);

Py_END_CRITICAL_SECTION();

return result;

}

static Py_ssize_t

set_len(PySetObject *so)

{

return FT_ATOMIC_LOAD_SSIZE_RELAXED(so->used);

}

static int

set_merge_lock_held(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));

_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);

_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(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);

so_entry->key = Py_NewRef(key);

so_entry->hash = other_entry->hash;

}

}

so->fill = other->fill;

FT_ATOMIC_STORE_SSIZE_RELAXED(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;

FT_ATOMIC_STORE_SSIZE_RELAXED(so->used, other->used);

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

key = other_entry->key;

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

set_insert_clean(newtable, newmask, Py_NewRef(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_impl(PySetObject *so)

{

/* 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;

FT_ATOMIC_STORE_SSIZE_RELAXED(so->used, so->used - 1);

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

return key;

}

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))

{

/* 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_GetBuiltin(&_Py_ID(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 = NULL;

Py_ssize_t i, mask;

setentry *entry;

PySetObject *so = si->si_set;

if (so == NULL)

return NULL;

assert (PyAnySet_Check(so));

Py_ssize_t so_used = FT_ATOMIC_LOAD_SSIZE(so->used);

Py_ssize_t si_used = FT_ATOMIC_LOAD_SSIZE(si->si_used);

if (si_used != so_used) {

PyErr_SetString(PyExc_RuntimeError,

"Set changed size during iteration");

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

return NULL;

}

Py_BEGIN_CRITICAL_SECTION(so);

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++;

}

if (i <= mask) {

key = Py_NewRef(entry[i].key);

}

Py_END_CRITICAL_SECTION();

si->si_pos = i+1;

if (key == NULL) {

si->si_set = NULL;

Py_DECREF(so);

return NULL;

}

si->len--;

return key;

}

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)

{

Py_ssize_t size = set_len(so);

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

if (si == NULL)

return NULL;

si->si_set = (PySetObject*)Py_NewRef(so);

si->si_used = size;

si->si_pos = 0;

si->len = size;

_PyObject_GC_TRACK(si);

return (PyObject *)si;

}

static int

set_update_dict_lock_held(PySetObject *so, PyObject *other)

{

assert(PyDict_CheckExact(other));

_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);

_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(other);

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

Py_ssize_t dictsize = PyDict_GET_SIZE(other);

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

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

return -1;

}

}

Py_ssize_t pos = 0;

PyObject *key;

PyObject *value;

Py_hash_t hash;

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

if (set_add_entry(so, key, hash)) {

return -1;

}

}

return 0;

}

static int

set_update_iterable_lock_held(PySetObject *so, PyObject *other)

{

_Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);

PyObject *it = PyObject_GetIter(other);

if (it == NULL) {

return -1;

}

PyObject *key;

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;

}

static int

set_update_lock_held(PySetObject *so, PyObject *other)

{

if (PyAnySet_Check(other)) {

return set_merge_lock_held(so, other);

}

else if (PyDict_CheckExact(other)) {

return set_update_dict_lock_held(so, other);

}

return set_update_iterable_lock_held(so, other);

}

static int

set_update_local(PySetObject *so, PyObject *other)

{

assert(Py_REFCNT(so) == 1);

if (PyAnySet_Check(other)) {

int rv;

Py_BEGIN_CRITICAL_SECTION(other);

rv = set_merge_lock_held(so, other);

Py_END_CRITICAL_SECTION();

return rv;

}

else if (PyDict_CheckExact(other)) {

int rv;

Py_BEGIN_CRITICAL_SECTION(other);

rv = set_update_dict_lock_held(so, other);

Py_END_CRITICAL_SECTION();

return rv;

}

return set_update_iterable_lock_held(so, other);