blob: 88ca69834857da77d3c22ddc7a4fcb9fd5b87f80 [file] [log] [blame]
/* Range object implementation */
#include "Python.h"
/* Support objects whose length is > PY_SSIZE_T_MAX.
This could be sped up for small PyLongs if they fit in an Py_ssize_t.
This only matters on Win64. Though we could use PY_LONG_LONG which
would presumably help perf.
*/
typedef struct {
PyObject_HEAD
PyObject *start;
PyObject *stop;
PyObject *step;
} rangeobject;
/* Helper function for validating step. Always returns a new reference or
NULL on error.
*/
static PyObject *
validate_step(PyObject *step)
{
/* No step specified, use a step of 1. */
if (!step)
return PyLong_FromLong(1);
step = PyNumber_Index(step);
if (step) {
Py_ssize_t istep = PyNumber_AsSsize_t(step, NULL);
if (istep == -1 && PyErr_Occurred()) {
/* Ignore OverflowError, we know the value isn't 0. */
PyErr_Clear();
}
else if (istep == 0) {
PyErr_SetString(PyExc_ValueError,
"range() arg 3 must not be zero");
Py_CLEAR(step);
}
}
return step;
}
/* XXX(nnorwitz): should we error check if the user passes any empty ranges?
range(-10)
range(0, -5)
range(0, 5, -1)
*/
static PyObject *
range_new(PyTypeObject *type, PyObject *args, PyObject *kw)
{
rangeobject *obj = NULL;
PyObject *start = NULL, *stop = NULL, *step = NULL;
if (!_PyArg_NoKeywords("range()", kw))
return NULL;
if (PyTuple_Size(args) <= 1) {
if (!PyArg_UnpackTuple(args, "range", 1, 1, &stop))
return NULL;
stop = PyNumber_Index(stop);
if (!stop)
return NULL;
start = PyLong_FromLong(0);
if (!start) {
Py_DECREF(stop);
return NULL;
}
step = PyLong_FromLong(1);
if (!step) {
Py_DECREF(stop);
Py_DECREF(start);
return NULL;
}
}
else {
if (!PyArg_UnpackTuple(args, "range", 2, 3,
&start, &stop, &step))
return NULL;
/* Convert borrowed refs to owned refs */
start = PyNumber_Index(start);
if (!start)
return NULL;
stop = PyNumber_Index(stop);
if (!stop) {
Py_DECREF(start);
return NULL;
}
step = validate_step(step); /* Caution, this can clear exceptions */
if (!step) {
Py_DECREF(start);
Py_DECREF(stop);
return NULL;
}
}
obj = PyObject_New(rangeobject, &PyRange_Type);
if (obj == NULL)
goto Fail;
obj->start = start;
obj->stop = stop;
obj->step = step;
return (PyObject *) obj;
Fail:
Py_XDECREF(start);
Py_XDECREF(stop);
Py_XDECREF(step);
return NULL;
}
PyDoc_STRVAR(range_doc,
"range([start,] stop[, step]) -> range object\n\
\n\
Returns an iterator that generates the numbers in the range on demand.");
static void
range_dealloc(rangeobject *r)
{
Py_DECREF(r->start);
Py_DECREF(r->stop);
Py_DECREF(r->step);
PyObject_Del(r);
}
/* Return number of items in range (lo, hi, step) as a PyLong object,
* when arguments are PyLong objects. Arguments MUST return 1 with
* PyLong_Check(). Return NULL when there is an error.
*/
static PyObject*
range_length_obj(rangeobject *r)
{
/* -------------------------------------------------------------
Algorithm is equal to that of get_len_of_range(), but it operates
on PyObjects (which are assumed to be PyLong objects).
---------------------------------------------------------------*/
int cmp_result;
PyObject *lo, *hi;
PyObject *step = NULL;
PyObject *diff = NULL;
PyObject *one = NULL;
PyObject *tmp1 = NULL, *tmp2 = NULL, *result;
/* holds sub-expression evaluations */
PyObject *zero = PyLong_FromLong(0);
if (zero == NULL)
return NULL;
cmp_result = PyObject_RichCompareBool(r->step, zero, Py_GT);
Py_DECREF(zero);
if (cmp_result == -1)
return NULL;
if (cmp_result == 1) {
lo = r->start;
hi = r->stop;
step = r->step;
Py_INCREF(step);
} else {
lo = r->stop;
hi = r->start;
step = PyNumber_Negative(r->step);
if (!step)
return NULL;
}
/* if (lo >= hi), return length of 0. */
if (PyObject_RichCompareBool(lo, hi, Py_GE) == 1) {
Py_XDECREF(step);
return PyLong_FromLong(0);
}
if ((one = PyLong_FromLong(1L)) == NULL)
goto Fail;
if ((tmp1 = PyNumber_Subtract(hi, lo)) == NULL)
goto Fail;
if ((diff = PyNumber_Subtract(tmp1, one)) == NULL)
goto Fail;
if ((tmp2 = PyNumber_FloorDivide(diff, step)) == NULL)
goto Fail;
if ((result = PyNumber_Add(tmp2, one)) == NULL)
goto Fail;
Py_DECREF(tmp2);
Py_DECREF(diff);
Py_DECREF(step);
Py_DECREF(tmp1);
Py_DECREF(one);
return result;
Fail:
Py_XDECREF(tmp2);
Py_XDECREF(diff);
Py_XDECREF(step);
Py_XDECREF(tmp1);
Py_XDECREF(one);
return NULL;
}
static Py_ssize_t
range_length(rangeobject *r)
{
PyObject *len = range_length_obj(r);
Py_ssize_t result = -1;
if (len) {
result = PyLong_AsSsize_t(len);
Py_DECREF(len);
}
return result;
}
/* range(...)[x] is necessary for: seq[:] = range(...) */
static PyObject *
range_item(rangeobject *r, Py_ssize_t i)
{
Py_ssize_t len = range_length(r);
PyObject *rem, *incr, *result;
/* XXX(nnorwitz): should negative indices be supported? */
/* XXX(nnorwitz): should support range[x] where x > PY_SSIZE_T_MAX? */
if (i < 0 || i >= len) {
if (!PyErr_Occurred())
PyErr_SetString(PyExc_IndexError,
"range object index out of range");
return NULL;
}
/* XXX(nnorwitz): optimize for short ints. */
rem = PyLong_FromSsize_t(i);
if (!rem)
return NULL;
incr = PyNumber_Multiply(rem, r->step);
Py_DECREF(rem);
if (!incr)
return NULL;
result = PyNumber_Add(r->start, incr);
Py_DECREF(incr);
return result;
}
static PyObject *
range_repr(rangeobject *r)
{
Py_ssize_t istep;
/* Check for special case values for printing. We don't always
need the step value. We don't care about errors
(it means overflow), so clear the errors. */
istep = PyNumber_AsSsize_t(r->step, NULL);
if (istep != 1 || (istep == -1 && PyErr_Occurred())) {
PyErr_Clear();
}
if (istep == 1)
return PyUnicode_FromFormat("range(%R, %R)", r->start, r->stop);
else
return PyUnicode_FromFormat("range(%R, %R, %R)",
r->start, r->stop, r->step);
}
/* Pickling support */
static PyObject *
range_reduce(rangeobject *r, PyObject *args)
{
return Py_BuildValue("(O(OOO))", Py_TYPE(r),
r->start, r->stop, r->step);
}
static PySequenceMethods range_as_sequence = {
(lenfunc)range_length, /* sq_length */
0, /* sq_concat */
0, /* sq_repeat */
(ssizeargfunc)range_item, /* sq_item */
0, /* sq_slice */
};
static PyObject * range_iter(PyObject *seq);
static PyObject * range_reverse(PyObject *seq);
PyDoc_STRVAR(reverse_doc,
"Returns a reverse iterator.");
static PyMethodDef range_methods[] = {
{"__reversed__", (PyCFunction)range_reverse, METH_NOARGS,
reverse_doc},
{"__reduce__", (PyCFunction)range_reduce, METH_VARARGS},
{NULL, NULL} /* sentinel */
};
PyTypeObject PyRange_Type = {
PyVarObject_HEAD_INIT(&PyType_Type, 0)
"range", /* Name of this type */
sizeof(rangeobject), /* Basic object size */
0, /* Item size for varobject */
(destructor)range_dealloc, /* tp_dealloc */
0, /* tp_print */
0, /* tp_getattr */
0, /* tp_setattr */
0, /* tp_reserved */
(reprfunc)range_repr, /* tp_repr */
0, /* tp_as_number */
&range_as_sequence, /* 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, /* tp_flags */
range_doc, /* tp_doc */
0, /* tp_traverse */
0, /* tp_clear */
0, /* tp_richcompare */
0, /* tp_weaklistoffset */
range_iter, /* tp_iter */
0, /* tp_iternext */
range_methods, /* 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 */
range_new, /* tp_new */
};
/*********************** range Iterator **************************/
/* There are 2 types of iterators, one for C longs, the other for
Python longs (ie, PyObjects). This should make iteration fast
in the normal case, but possible for any numeric value.
*/
typedef struct {
PyObject_HEAD
long index;
long start;
long step;
long len;
} rangeiterobject;
static PyObject *
rangeiter_next(rangeiterobject *r)
{
if (r->index < r->len)
return PyLong_FromLong(r->start + (r->index++) * r->step);
return NULL;
}
static PyObject *
rangeiter_len(rangeiterobject *r)
{
return PyLong_FromLong(r->len - r->index);
}
typedef struct {
PyObject_HEAD
PyObject *index;
PyObject *start;
PyObject *step;
PyObject *len;
} longrangeiterobject;
static PyObject *
longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
{
return PyNumber_Subtract(r->len, r->index);
}
static PyObject *rangeiter_new(PyTypeObject *, PyObject *args, PyObject *kw);
PyDoc_STRVAR(length_hint_doc,
"Private method returning an estimate of len(list(it)).");
static PyMethodDef rangeiter_methods[] = {
{"__length_hint__", (PyCFunction)rangeiter_len, METH_NOARGS,
length_hint_doc},
{NULL, NULL} /* sentinel */
};
PyTypeObject PyRangeIter_Type = {
PyVarObject_HEAD_INIT(&PyType_Type, 0)
"range_iterator", /* tp_name */
sizeof(rangeiterobject), /* tp_basicsize */
0, /* tp_itemsize */
/* methods */
(destructor)PyObject_Del, /* tp_dealloc */
0, /* tp_print */
0, /* tp_getattr */
0, /* tp_setattr */
0, /* tp_reserved */
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, /* tp_flags */
0, /* tp_doc */
0, /* tp_traverse */
0, /* tp_clear */
0, /* tp_richcompare */
0, /* tp_weaklistoffset */
PyObject_SelfIter, /* tp_iter */
(iternextfunc)rangeiter_next, /* tp_iternext */
rangeiter_methods, /* 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 */
rangeiter_new, /* tp_new */
};
/* Return number of items in range (lo, hi, step). step > 0
* required. Return a value < 0 if & only if the true value is too
* large to fit in a signed long.
*/
static long
get_len_of_range(long lo, long hi, long step)
{
/* -------------------------------------------------------------
If lo >= hi, the range is empty.
Else if n values are in the range, the last one is
lo + (n-1)*step, which must be <= hi-1. Rearranging,
n <= (hi - lo - 1)/step + 1, so taking the floor of the RHS gives
the proper value. Since lo < hi in this case, hi-lo-1 >= 0, so
the RHS is non-negative and so truncation is the same as the
floor. Letting M be the largest positive long, the worst case
for the RHS numerator is hi=M, lo=-M-1, and then
hi-lo-1 = M-(-M-1)-1 = 2*M. Therefore unsigned long has enough
precision to compute the RHS exactly.
---------------------------------------------------------------*/
long n = 0;
if (lo < hi) {
unsigned long uhi = (unsigned long)hi;
unsigned long ulo = (unsigned long)lo;
unsigned long diff = uhi - ulo - 1;
n = (long)(diff / (unsigned long)step + 1);
}
return n;
}
static PyObject *
int_range_iter(long start, long stop, long step)
{
rangeiterobject *it = PyObject_New(rangeiterobject, &PyRangeIter_Type);
if (it == NULL)
return NULL;
it->start = start;
it->step = step;
if (step > 0)
it->len = get_len_of_range(start, stop, step);
else
it->len = get_len_of_range(stop, start, -step);
it->index = 0;
return (PyObject *)it;
}
static PyObject *
rangeiter_new(PyTypeObject *type, PyObject *args, PyObject *kw)
{
long start, stop, step;
if (!_PyArg_NoKeywords("rangeiter()", kw))
return NULL;
if (!PyArg_ParseTuple(args, "lll;rangeiter() requires 3 int arguments",
&start, &stop, &step))
return NULL;
return int_range_iter(start, stop, step);
}
static PyMethodDef longrangeiter_methods[] = {
{"__length_hint__", (PyCFunction)longrangeiter_len, METH_NOARGS,
length_hint_doc},
{NULL, NULL} /* sentinel */
};
static void
longrangeiter_dealloc(longrangeiterobject *r)
{
Py_XDECREF(r->index);
Py_XDECREF(r->start);
Py_XDECREF(r->step);
Py_XDECREF(r->len);
PyObject_Del(r);
}
static PyObject *
longrangeiter_next(longrangeiterobject *r)
{
PyObject *one, *product, *new_index, *result;
if (PyObject_RichCompareBool(r->index, r->len, Py_LT) != 1)
return NULL;
one = PyLong_FromLong(1);
if (!one)
return NULL;
product = PyNumber_Multiply(r->index, r->step);
if (!product) {
Py_DECREF(one);
return NULL;
}
new_index = PyNumber_Add(r->index, one);
Py_DECREF(one);
if (!new_index) {
Py_DECREF(product);
return NULL;
}
result = PyNumber_Add(r->start, product);
Py_DECREF(product);
if (result) {
Py_DECREF(r->index);
r->index = new_index;
}
return result;
}
PyTypeObject PyLongRangeIter_Type = {
PyVarObject_HEAD_INIT(&PyType_Type, 0)
"longrange_iterator", /* tp_name */
sizeof(longrangeiterobject), /* tp_basicsize */
0, /* tp_itemsize */
/* methods */
(destructor)longrangeiter_dealloc, /* tp_dealloc */
0, /* tp_print */
0, /* tp_getattr */
0, /* tp_setattr */
0, /* tp_reserved */
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, /* tp_flags */
0, /* tp_doc */
0, /* tp_traverse */
0, /* tp_clear */
0, /* tp_richcompare */
0, /* tp_weaklistoffset */
PyObject_SelfIter, /* tp_iter */
(iternextfunc)longrangeiter_next, /* tp_iternext */
longrangeiter_methods, /* tp_methods */
0,
};
static PyObject *
range_iter(PyObject *seq)
{
rangeobject *r = (rangeobject *)seq;
longrangeiterobject *it;
long lstart, lstop, lstep;
assert(PyRange_Check(seq));
/* If all three fields convert to long, use the int version */
lstart = PyLong_AsLong(r->start);
if (lstart != -1 || !PyErr_Occurred()) {
lstop = PyLong_AsLong(r->stop);
if (lstop != -1 || !PyErr_Occurred()) {
lstep = PyLong_AsLong(r->step);
if (lstep != -1 || !PyErr_Occurred())
return int_range_iter(lstart, lstop, lstep);
}
}
/* Some conversion failed, so there is an error set. Clear it,
and try again with a long range. */
PyErr_Clear();
it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
if (it == NULL)
return NULL;
/* Do all initialization here, so we can DECREF on failure. */
it->start = r->start;
it->step = r->step;
Py_INCREF(it->start);
Py_INCREF(it->step);
it->len = it->index = NULL;
it->len = range_length_obj(r);
if (!it->len)
goto create_failure;
it->index = PyLong_FromLong(0);
if (!it->index)
goto create_failure;
return (PyObject *)it;
create_failure:
Py_DECREF(it);
return NULL;
}
static PyObject *
range_reverse(PyObject *seq)
{
rangeobject *range = (rangeobject*) seq;
longrangeiterobject *it;
PyObject *one, *sum, *diff, *len = NULL, *product;
long lstart, lstop, lstep;
/* XXX(nnorwitz): do the calc for the new start/stop first,
then if they fit, call the proper iter()?
*/
assert(PyRange_Check(seq));
/* If all three fields convert to long, use the int version */
lstart = PyLong_AsLong(range->start);
if (lstart != -1 || !PyErr_Occurred()) {
lstop = PyLong_AsLong(range->stop);
if (lstop != -1 || !PyErr_Occurred()) {
lstep = PyLong_AsLong(range->step);
if (lstep != -1 || !PyErr_Occurred()) {
/* XXX(nnorwitz): need to check for overflow and simplify. */
long len = get_len_of_range(lstart, lstop, lstep);
long new_start = lstart + (len - 1) * lstep;
long new_stop = lstart;
if (lstep > 0)
new_stop -= 1;
else
new_stop += 1;
return int_range_iter(new_start, new_stop, -lstep);
}
}
}
PyErr_Clear();
it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
if (it == NULL)
return NULL;
/* start + (len - 1) * step */
len = range_length_obj(range);
if (!len)
goto create_failure;
one = PyLong_FromLong(1);
if (!one)
goto create_failure;
diff = PyNumber_Subtract(len, one);
Py_DECREF(one);
if (!diff)
goto create_failure;
product = PyNumber_Multiply(len, range->step);
if (!product)
goto create_failure;
sum = PyNumber_Add(range->start, product);
Py_DECREF(product);
it->start = sum;
if (!it->start)
goto create_failure;
it->step = PyNumber_Negative(range->step);
if (!it->step) {
Py_DECREF(it->start);
PyObject_Del(it);
return NULL;
}
/* Steal reference to len. */
it->len = len;
it->index = PyLong_FromLong(0);
if (!it->index) {
Py_DECREF(it);
return NULL;
}
return (PyObject *)it;
create_failure:
Py_XDECREF(len);
PyObject_Del(it);
return NULL;
}