#include "Python.h"

#include "structmember.h"

/* Itertools module written and maintained 

   by Raymond D. Hettinger <python@rcn.com>

   Copyright (c) 2003 Python Software Foundation.

   All rights reserved.

*/

/* groupby object ***********************************************************/

typedef struct {

	PyObject_HEAD

	PyObject *it;

	PyObject *keyfunc;

	PyObject *tgtkey;

	PyObject *currkey;

	PyObject *currvalue;

} groupbyobject;

static PyTypeObject groupby_type;

static PyObject *_grouper_create(groupbyobject *, PyObject *);

static PyObject *

groupby_new(PyTypeObject *type, PyObject *args, PyObject *kwds)

{

	static char *kwargs[] = {"iterable", "key", NULL};

	groupbyobject *gbo;

 	PyObject *it, *keyfunc = Py_None;

 	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O:groupby", kwargs,

					 &it, &keyfunc))

		return NULL;

	gbo = (groupbyobject *)type->tp_alloc(type, 0);

	if (gbo == NULL)

		return NULL;

	gbo->tgtkey = NULL;

	gbo->currkey = NULL;

	gbo->currvalue = NULL;

	gbo->keyfunc = keyfunc;

	Py_INCREF(keyfunc);

	gbo->it = PyObject_GetIter(it);

	if (gbo->it == NULL) {

		Py_DECREF(gbo);

		return NULL;

	}

	return (PyObject *)gbo;

}

static void

groupby_dealloc(groupbyobject *gbo)

{

	PyObject_GC_UnTrack(gbo);

	Py_XDECREF(gbo->it);

	Py_XDECREF(gbo->keyfunc);

	Py_XDECREF(gbo->tgtkey);

	Py_XDECREF(gbo->currkey);

	Py_XDECREF(gbo->currvalue);

	Py_TYPE(gbo)->tp_free(gbo);

}

static int

groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)

{

	Py_VISIT(gbo->it);

	Py_VISIT(gbo->keyfunc);

	Py_VISIT(gbo->tgtkey);

	Py_VISIT(gbo->currkey);

	Py_VISIT(gbo->currvalue);

	return 0;

}

static PyObject *

groupby_next(groupbyobject *gbo)

{

	PyObject *newvalue, *newkey, *r, *grouper, *tmp;

	/* skip to next iteration group */

	for (;;) {

		if (gbo->currkey == NULL)

			/* pass */;

		else if (gbo->tgtkey == NULL)

			break;

		else {

			int rcmp;

			rcmp = PyObject_RichCompareBool(gbo->tgtkey,

							gbo->currkey, Py_EQ);

			if (rcmp == -1)

				return NULL;

			else if (rcmp == 0)

				break;

		}

		newvalue = PyIter_Next(gbo->it);

		if (newvalue == NULL)

			return NULL;

		if (gbo->keyfunc == Py_None) {

			newkey = newvalue;

			Py_INCREF(newvalue);

		} else {

			newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,

							      newvalue, NULL);

			if (newkey == NULL) {

				Py_DECREF(newvalue);

				return NULL;

			}

		}

		tmp = gbo->currkey;

		gbo->currkey = newkey;

		Py_XDECREF(tmp);

		tmp = gbo->currvalue;

		gbo->currvalue = newvalue;

		Py_XDECREF(tmp);

	}

	Py_INCREF(gbo->currkey);

	tmp = gbo->tgtkey;

	gbo->tgtkey = gbo->currkey;

	Py_XDECREF(tmp);

	grouper = _grouper_create(gbo, gbo->tgtkey);

	if (grouper == NULL)

		return NULL;

	r = PyTuple_Pack(2, gbo->currkey, grouper);

	Py_DECREF(grouper);

	return r;

}

PyDoc_STRVAR(groupby_doc,

"groupby(iterable[, keyfunc]) -> create an iterator which returns\n\

(key, sub-iterator) grouped by each value of key(value).\n");

static PyTypeObject groupby_type = {

	PyVarObject_HEAD_INIT(NULL, 0)

	"itertools.groupby",		/* tp_name */

	sizeof(groupbyobject),		/* tp_basicsize */

	0,				/* tp_itemsize */

	/* methods */

	(destructor)groupby_dealloc,	/* tp_dealloc */

	0,				/* tp_print */

	0,				/* tp_getattr */

	0,				/* tp_setattr */

	0,				/* tp_compare */

	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 |

		Py_TPFLAGS_BASETYPE,	/* tp_flags */

	groupby_doc,			/* tp_doc */

	(traverseproc)groupby_traverse,	/* tp_traverse */

	0,				/* tp_clear */

	0,				/* tp_richcompare */

	0,				/* tp_weaklistoffset */

	PyObject_SelfIter,		/* tp_iter */

	(iternextfunc)groupby_next,	/* tp_iternext */

	0,				/* tp_methods */

	0,				/* tp_members */

	0,				/* tp_getset */

	0,				/* tp_base */

	0,				/* tp_dict */

	0,				/* tp_descr_get */

	0,				/* tp_descr_set */

	0,				/* tp_dictoffset */

	0,				/* tp_init */

	0,				/* tp_alloc */

	groupby_new,			/* tp_new */

	PyObject_GC_Del,		/* tp_free */

};

/* _grouper object (internal) ************************************************/

typedef struct {

	PyObject_HEAD

	PyObject *parent;

	PyObject *tgtkey;

} _grouperobject;

static PyTypeObject _grouper_type;

static PyObject *

_grouper_create(groupbyobject *parent, PyObject *tgtkey)

{

	_grouperobject *igo;

	igo = PyObject_GC_New(_grouperobject, &_grouper_type);

	if (igo == NULL)

		return NULL;

	igo->parent = (PyObject *)parent;

	Py_INCREF(parent);

	igo->tgtkey = tgtkey;

	Py_INCREF(tgtkey);

	PyObject_GC_Track(igo);

	return (PyObject *)igo;

}

static void

_grouper_dealloc(_grouperobject *igo)

{

	PyObject_GC_UnTrack(igo);

	Py_DECREF(igo->parent);

	Py_DECREF(igo->tgtkey);

	PyObject_GC_Del(igo);

}

static int

_grouper_traverse(_grouperobject *igo, visitproc visit, void *arg)

{

	Py_VISIT(igo->parent);

	Py_VISIT(igo->tgtkey);

	return 0;

}

static PyObject *

_grouper_next(_grouperobject *igo)

{

	groupbyobject *gbo = (groupbyobject *)igo->parent;

	PyObject *newvalue, *newkey, *r;

	int rcmp;

	if (gbo->currvalue == NULL) {

		newvalue = PyIter_Next(gbo->it);

		if (newvalue == NULL)

			return NULL;

		if (gbo->keyfunc == Py_None) {

			newkey = newvalue;

			Py_INCREF(newvalue);

		} else {

			newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc,

							      newvalue, NULL);

			if (newkey == NULL) {

				Py_DECREF(newvalue);

				return NULL;

			}

		}

		assert(gbo->currkey == NULL);

		gbo->currkey = newkey;

		gbo->currvalue = newvalue;

	}

	assert(gbo->currkey != NULL);

	rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);

	if (rcmp <= 0)

		/* got any error or current group is end */

		return NULL;

	r = gbo->currvalue;

	gbo->currvalue = NULL;

	Py_CLEAR(gbo->currkey);

	return r;

}

static PyTypeObject _grouper_type = {

	PyVarObject_HEAD_INIT(NULL, 0)

	"itertools._grouper",		/* tp_name */

	sizeof(_grouperobject),		/* tp_basicsize */

	0,				/* tp_itemsize */

	/* methods */

	(destructor)_grouper_dealloc,	/* tp_dealloc */

	0,				/* tp_print */

	0,				/* tp_getattr */

	0,				/* tp_setattr */

	0,				/* tp_compare */

	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)_grouper_traverse,/* tp_traverse */

	0,				/* tp_clear */

	0,				/* tp_richcompare */

	0,				/* tp_weaklistoffset */

	PyObject_SelfIter,		/* tp_iter */

	(iternextfunc)_grouper_next,	/* tp_iternext */

	0,				/* tp_methods */

	0,				/* tp_members */

	0,				/* tp_getset */

	0,				/* tp_base */

	0,				/* tp_dict */

	0,				/* tp_descr_get */

	0,				/* tp_descr_set */

	0,				/* tp_dictoffset */

	0,				/* tp_init */

	0,				/* tp_alloc */

	0,				/* tp_new */

	PyObject_GC_Del,		/* tp_free */

};

/* tee object and with supporting function and objects ***************/

/* The teedataobject pre-allocates space for LINKCELLS number of objects.

   To help the object fit neatly inside cache lines (space for 16 to 32

   pointers), the value should be a multiple of 16 minus  space for 

   the other structure members including PyHEAD overhead.  The larger the

   value, the less memory overhead per object and the less time spent

   allocating/deallocating new links.  The smaller the number, the less

   wasted space and the more rapid freeing of older data.

*/

#define LINKCELLS 57

typedef struct {

	PyObject_HEAD

	PyObject *it;

	int numread;

	PyObject *nextlink;

	PyObject *(values[LINKCELLS]);

} teedataobject;

typedef struct {

	PyObject_HEAD

	teedataobject *dataobj;

	int index;

	PyObject *weakreflist;

} teeobject;

static PyTypeObject teedataobject_type;

static PyObject *

teedataobject_new(PyObject *it)

{

	teedataobject *tdo;

	tdo = PyObject_GC_New(teedataobject, &teedataobject_type);

	if (tdo == NULL)

		return NULL;

	tdo->numread = 0;

	tdo->nextlink = NULL;

	Py_INCREF(it);

	tdo->it = it;

	PyObject_GC_Track(tdo);

	return (PyObject *)tdo;

}

static PyObject *

teedataobject_jumplink(teedataobject *tdo)

{

	if (tdo->nextlink == NULL)

		tdo->nextlink = teedataobject_new(tdo->it);

	Py_XINCREF(tdo->nextlink);

	return tdo->nextlink;

}

static PyObject *

teedataobject_getitem(teedataobject *tdo, int i)

{

	PyObject *value;

	assert(i < LINKCELLS);

	if (i < tdo->numread)

		value = tdo->values[i];

	else {

		/* this is the lead iterator, so fetch more data */

		assert(i == tdo->numread);

		value = PyIter_Next(tdo->it);

		if (value == NULL)

			return NULL;

		tdo->numread++;

		tdo->values[i] = value;

	}

	Py_INCREF(value);

	return value;

}

static int

teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)

{

	int i;

	Py_VISIT(tdo->it);

	for (i = 0; i < tdo->numread; i++)

		Py_VISIT(tdo->values[i]);

	Py_VISIT(tdo->nextlink);

	return 0;

}

static int

teedataobject_clear(teedataobject *tdo)

{

	int i;

	Py_CLEAR(tdo->it);

	for (i=0 ; i<tdo->numread ; i++)

		Py_CLEAR(tdo->values[i]);

	Py_CLEAR(tdo->nextlink);

	return 0;

}

static void

teedataobject_dealloc(teedataobject *tdo)

{

	PyObject_GC_UnTrack(tdo);

	teedataobject_clear(tdo);

	PyObject_GC_Del(tdo);

}

PyDoc_STRVAR(teedataobject_doc, "Data container common to multiple tee objects.");

static PyTypeObject teedataobject_type = {

	PyVarObject_HEAD_INIT(0, 0)	/* Must fill in type value later */

	"itertools.tee_dataobject",		/* tp_name */

	sizeof(teedataobject),			/* tp_basicsize */

	0,					/* tp_itemsize */

	/* methods */

	(destructor)teedataobject_dealloc,	/* tp_dealloc */

	0,					/* tp_print */

	0,					/* tp_getattr */

	0,					/* tp_setattr */

	0,					/* tp_compare */

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

	teedataobject_doc,			/* tp_doc */

	(traverseproc)teedataobject_traverse,	/* tp_traverse */

	(inquiry)teedataobject_clear,		/* tp_clear */

	0,					/* tp_richcompare */

	0,					/* tp_weaklistoffset */

	0,					/* tp_iter */

	0,					/* tp_iternext */

	0,					/* tp_methods */

	0,					/* tp_members */

	0,					/* tp_getset */

	0,					/* tp_base */

	0,					/* tp_dict */

	0,					/* tp_descr_get */

	0,					/* tp_descr_set */

	0,					/* tp_dictoffset */

	0,					/* tp_init */

	0,					/* tp_alloc */

	0,					/* tp_new */

	PyObject_GC_Del,			/* tp_free */

};

static PyTypeObject tee_type;

static PyObject *

tee_next(teeobject *to)

{

	PyObject *value, *link;

	if (to->index >= LINKCELLS) {

		link = teedataobject_jumplink(to->dataobj);

		Py_DECREF(to->dataobj);

		to->dataobj = (teedataobject *)link;

		to->index = 0;

	}

	value = teedataobject_getitem(to->dataobj, to->index);

	if (value == NULL)

		return NULL;

	to->index++;

	return value;

}

static int

tee_traverse(teeobject *to, visitproc visit, void *arg)

{

	Py_VISIT((PyObject *)to->dataobj);

	return 0;

}

static PyObject *

tee_copy(teeobject *to)

{

	teeobject *newto;

	newto = PyObject_GC_New(teeobject, &tee_type);

	if (newto == NULL)

		return NULL;

	Py_INCREF(to->dataobj);

	newto->dataobj = to->dataobj;

	newto->index = to->index;

	newto->weakreflist = NULL;

	PyObject_GC_Track(newto);

	return (PyObject *)newto;

}

PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");

static PyObject *

tee_fromiterable(PyObject *iterable)

{

	teeobject *to;

	PyObject *it = NULL;

	it = PyObject_GetIter(iterable);

	if (it == NULL)

		return NULL;

	if (PyObject_TypeCheck(it, &tee_type)) {

		to = (teeobject *)tee_copy((teeobject *)it);

		goto done;

	}

	to = PyObject_GC_New(teeobject, &tee_type);