/* set object implementation 

   Written and maintained by Raymond D. Hettinger <python@rcn.com>

   Derived from Lib/sets.py and Objects/dictobject.c.

   Copyright (c) 2003-2007 Python Software Foundation.

   All rights reserved.

*/

#include "Python.h"

#include "structmember.h"

/* Set a key error with the specified argument, wrapping it in a

 * tuple automatically so that tuple keys are not unpacked as the

 * exception arguments. */

static void

set_key_error(PyObject *arg)

{

	PyObject *tup;

	tup = PyTuple_Pack(1, arg);

	if (!tup)

		return; /* caller will expect error to be set anyway */

	PyErr_SetObject(PyExc_KeyError, tup);

	Py_DECREF(tup);

}

/* This must be >= 1. */

#define PERTURB_SHIFT 5

/* Object used as dummy key to fill deleted entries */

static PyObject *dummy = NULL; /* Initialized by first call to make_new_set() */

#ifdef Py_REF_DEBUG

PyObject *

_PySet_Dummy(void)

{

	return dummy;

}

#endif

#define INIT_NONZERO_SET_SLOTS(so) do {				\

	(so)->table = (so)->smalltable;				\

	(so)->mask = PySet_MINSIZE - 1;				\

	(so)->hash = -1;					\

    } while(0)

#define EMPTY_TO_MINSIZE(so) do {				\

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

	(so)->used = (so)->fill = 0;				\

	INIT_NONZERO_SET_SLOTS(so);				\

    } while(0)

/* Reuse scheme to save calls to malloc, free, and memset */

#ifndef PySet_MAXFREELIST

#define PySet_MAXFREELIST 80

#endif

static PySetObject *free_list[PySet_MAXFREELIST];

static int numfree = 0;

/*

The basic lookup function used by all operations.

This is based on Algorithm D from Knuth Vol. 3, Sec. 6.4.

Open addressing is preferred over chaining since the link overhead for

chaining would be substantial (100% with typical malloc overhead).

The initial probe index is computed as hash mod the table size. Subsequent

probe indices are computed as explained in Objects/dictobject.c.

All arithmetic on hash should ignore overflow.

Unlike the dictionary implementation, the lookkey functions can return

NULL if the rich comparison returns an error.

*/

static setentry *

set_lookkey(PySetObject *so, PyObject *key, register long hash)

{

	register Py_ssize_t i;

	register size_t perturb;

	register setentry *freeslot;

	register size_t mask = so->mask;

	setentry *table = so->table;

	register setentry *entry;

	register int cmp;

	PyObject *startkey;

	i = hash & mask;

	entry = &table[i];

	if (entry->key == NULL || entry->key == key)

		return entry;

	if (entry->key == dummy)

		freeslot = entry;

	else {

		if (entry->hash == hash) {

			startkey = entry->key;

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

				if (cmp > 0)

					return entry;

			}

			else {

				/* The compare did major nasty stuff to the

				 * set:  start over.

 				 */

 				return set_lookkey(so, key, hash);

 			}

		}

		freeslot = NULL;

	}

	/* In the loop, key == dummy is by far (factor of 100s) the

	   least likely outcome, so test for that last. */

	for (perturb = hash; ; perturb >>= PERTURB_SHIFT) {

		i = (i << 2) + i + perturb + 1;

		entry = &table[i & mask];

		if (entry->key == NULL) {

			if (freeslot != NULL)

				entry = freeslot;

			break;

		}

		if (entry->key == key)

			break;

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

			startkey = entry->key;

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

				if (cmp > 0)

					break;

			}

			else {

				/* The compare did major nasty stuff to the

				 * set:  start over.

 				 */

 				return set_lookkey(so, key, hash);

 			}

		}

		else if (entry->key == dummy && freeslot == NULL)

			freeslot = entry;

	}

	return entry;

}

/*

 * Hacked up version of set_lookkey which can assume keys are always strings;

 * This means we can always use _PyString_Eq directly and not have to check to

 * see if the comparison altered the table.

 */

static setentry *

set_lookkey_string(PySetObject *so, PyObject *key, register long hash)

{

	register Py_ssize_t i;

	register size_t perturb;

	register setentry *freeslot;

	register size_t mask = so->mask;

	setentry *table = so->table;

	register setentry *entry;

	/* Make sure this function doesn't have to handle non-string keys,

	   including subclasses of str; e.g., one reason to subclass

	   strings is to override __eq__, and for speed we don't cater to

	   that here. */

	if (!PyString_CheckExact(key)) {

		so->lookup = set_lookkey;

		return set_lookkey(so, key, hash);

	}

	i = hash & mask;

	entry = &table[i];

	if (entry->key == NULL || entry->key == key)

		return entry;

	if (entry->key == dummy)

		freeslot = entry;

	else {

		if (entry->hash == hash && _PyString_Eq(entry->key, key))

			return entry;

		freeslot = NULL;

	}

	/* In the loop, key == dummy is by far (factor of 100s) the

	   least likely outcome, so test for that last. */

	for (perturb = hash; ; perturb >>= PERTURB_SHIFT) {

		i = (i << 2) + i + perturb + 1;

		entry = &table[i & mask];

		if (entry->key == NULL)

			return freeslot == NULL ? entry : freeslot;

		if (entry->key == key

		    || (entry->hash == hash

			&& entry->key != dummy

			&& _PyString_Eq(entry->key, key)))

			return entry;

		if (entry->key == dummy && freeslot == NULL)

			freeslot = entry;

	}

	assert(0);	/* NOT REACHED */

	return 0;

}

/*

Internal routine to insert a new key into the table.

Used by the public insert routine.

Eats a reference to key.

*/

static int

set_insert_key(register PySetObject *so, PyObject *key, long hash)

{

	register setentry *entry;

	typedef setentry *(*lookupfunc)(PySetObject *, PyObject *, long);

	assert(so->lookup != NULL);

	entry = so->lookup(so, key, hash);

	if (entry == NULL)

		return -1;

	if (entry->key == NULL) {

		/* UNUSED */

		so->fill++; 

		entry->key = key;

		entry->hash = hash;

		so->used++;

	} else if (entry->key == dummy) {

		/* DUMMY */

		entry->key = key;

		entry->hash = hash;

		so->used++;

		Py_DECREF(dummy);

	} else {

		/* ACTIVE */

		Py_DECREF(key);

	}

	return 0;

}

/*

Internal routine used by set_table_resize() to insert an item which is

known to be absent from the set.  This routine also assumes that

the set contains no deleted entries.  Besides the performance benefit,

using set_insert_clean() in set_table_resize() is dangerous (SF bug #1456209).

Note that no refcounts are changed by this routine; if needed, the caller

is responsible for incref'ing `key`.

*/

static void

set_insert_clean(register PySetObject *so, PyObject *key, long hash)

{

	register size_t i;

	register size_t perturb;

	register size_t mask = (size_t)so->mask;

	setentry *table = so->table;

	register setentry *entry;

	i = hash & mask;

	entry = &table[i];

	for (perturb = hash; entry->key != NULL; perturb >>= PERTURB_SHIFT) {

		i = (i << 2) + i + perturb + 1;

		entry = &table[i & mask];

	}

	so->fill++;

	entry->key = key;

	entry->hash = hash;

	so->used++;

}

/*

Restructure the table by allocating a new table and reinserting all

keys again.  When entries have been deleted, the new table may

actually be smaller than the old one.

*/

static int

set_table_resize(PySetObject *so, Py_ssize_t minused)

{

	Py_ssize_t newsize;

	setentry *oldtable, *newtable, *entry;

	Py_ssize_t i;

	int is_oldtable_malloced;

	setentry small_copy[PySet_MINSIZE];

	assert(minused >= 0);

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

	for (newsize = PySet_MINSIZE;

	     newsize <= minused && newsize > 0;

	     newsize <<= 1)

		;

	if (newsize <= 0) {

		PyErr_NoMemory();

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

			   table anyway to purge old dummy entries.

			   Subtle:  This is *necessary* if fill==size,

			   as set_lookkey needs at least one virgin slot to

			   terminate failing searches.  If fill < size, it's

			   merely desirable, as dummies slow searches. */

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

	so->table = newtable;

	so->mask = newsize - 1;

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

	so->used = 0;

	i = so->fill;

	so->fill = 0;

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

	   dummy entries aren't copied over, of course */

	for (entry = oldtable; i > 0; entry++) {

		if (entry->key == NULL) {

			/* UNUSED */

			;

		} else if (entry->key == dummy) {

			/* DUMMY */

			--i;

			assert(entry->key == dummy);

			Py_DECREF(entry->key);

		} else {

			/* ACTIVE */

			--i;

			set_insert_clean(so, entry->key, entry->hash);

		}

	}

	if (is_oldtable_malloced)

		PyMem_DEL(oldtable);

	return 0;

}

/* CAUTION: set_add_key/entry() must guarantee it won't resize the table */

static int

set_add_entry(register PySetObject *so, setentry *entry)

{

	register Py_ssize_t n_used;

	assert(so->fill <= so->mask);  /* at least one empty slot */

	n_used = so->used;

	Py_INCREF(entry->key);

	if (set_insert_key(so, entry->key, entry->hash) == -1) {

		Py_DECREF(entry->key);

		return -1;

	}

	if (!(so->used > n_used && so->fill*3 >= (so->mask+1)*2))

		return 0;

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

}

static int

set_add_key(register PySetObject *so, PyObject *key)

{

	register long hash;

	register Py_ssize_t n_used;

	if (!PyString_CheckExact(key) ||

	    (hash = ((PyStringObject *) key)->ob_shash) == -1) {

		hash = PyObject_Hash(key);

		if (hash == -1)

			return -1;

	}

	assert(so->fill <= so->mask);  /* at least one empty slot */

	n_used = so->used;

	Py_INCREF(key);

	if (set_insert_key(so, key, hash) == -1) {

		Py_DECREF(key);

		return -1;

	}

	if (!(so->used > n_used && so->fill*3 >= (so->mask+1)*2))

		return 0;

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

}

#define DISCARD_NOTFOUND 0

#define DISCARD_FOUND 1

static int

set_discard_entry(PySetObject *so, setentry *oldentry)

{	register setentry *entry;

	PyObject *old_key;

	entry = (so->lookup)(so, oldentry->key, oldentry->hash);

	if (entry == NULL)

		return -1;

	if (entry->key == NULL  ||  entry->key == dummy)

		return DISCARD_NOTFOUND;

	old_key = entry->key;

	Py_INCREF(dummy);

	entry->key = dummy;

	so->used--;

	Py_DECREF(old_key);

	return DISCARD_FOUND;

}

static int

set_discard_key(PySetObject *so, PyObject *key)

{

	register long hash;

	register setentry *entry;

	PyObject *old_key;

	assert (PyAnySet_Check(so));

	if (!PyString_CheckExact(key) ||

	    (hash = ((PyStringObject *) key)->ob_shash) == -1) {

		hash = PyObject_Hash(key);

		if (hash == -1)

			return -1;

	}

	entry = (so->lookup)(so, key, hash);

	if (entry == NULL)

		return -1;

	if (entry->key == NULL  ||  entry->key == dummy)

		return DISCARD_NOTFOUND;

	old_key = entry->key;

	Py_INCREF(dummy);

	entry->key = dummy;

	so->used--;

	Py_DECREF(old_key);

	return DISCARD_FOUND;

}

static int

set_clear_internal(PySetObject *so)

{

	setentry *entry, *table;

	int table_is_malloced;

	Py_ssize_t fill;

	setentry small_copy[PySet_MINSIZE];

#ifdef Py_DEBUG

	Py_ssize_t i, n;

	assert (PyAnySet_Check(so));

	n = so->mask + 1;

	i = 0;

#endif

	table = so->table;

	assert(table != NULL);

	table_is_malloced = table != so->smalltable;

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

	 * decrefs can cause the set to mutate.  To avoid fatal confusion

	 * (voice of experience), we have to make the set empty before

	 * clearing the slots, and never refer to anything via so->ref while

	 * clearing.

	 */

	fill = so->fill;

	if (table_is_malloced)

		EMPTY_TO_MINSIZE(so);

	else if (fill > 0) {

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

		 * Afraid the only safe way is to copy the set entries into

		 * another small table first.

		 */

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

		table = small_copy;

		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

	 * assert that the refcount on table is 1 now, i.e. that this function

	 * has unique access to it, so decref side-effects can't alter it.

	 */

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

#ifdef Py_DEBUG

		assert(i < n);

		++i;

#endif

		if (entry->key) {

			--fill;

			Py_DECREF(entry->key);

		}

#ifdef Py_DEBUG

		else

			assert(entry->key == NULL);

#endif

	}

	if (table_is_malloced)

		PyMem_DEL(table);

	return 0;

}

/*

 * Iterate over a set table.  Use like so:

 *

 *     Py_ssize_t pos;

 *     setentry *entry;

 *     pos = 0;   # important!  pos should not otherwise be changed by you

 *     while (set_next(yourset, &pos, &entry)) {

 *              Refer to borrowed reference in entry->key.

 *     }

 *

 * CAUTION:  In general, it isn't safe to use set_next in a loop that

 * mutates the table.  

 */

static int

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