range: lean and mean

Objects/rangeobject.c

 /* Range object implementation */
 
 #include "Python.h"
+#include "structmember.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;
+    PyObject *length;
 } 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;
 }
+
+static PyObject* range_compute_length(PyObject *start,
+                                      PyObject *stop, PyObject *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;
+    PyObject *start = NULL, *stop = NULL, *step = NULL, *length = NULL;
 
     if (!_PyArg_NoKeywords("range()", kw))
         return NULL;
 
     if (PyTuple_Size(args) <= 1) {
         if (!PyArg_UnpackTuple(args, "range", 1, 1, &stop))
             goto Fail;
         stop = PyNumber_Index(stop);
         if (!stop)
             goto Fail;
         start = PyLong_FromLong(0);
         step = PyLong_FromLong(1);
         if (!start || !step)
             goto Fail;
     }
     else {
         if (!PyArg_UnpackTuple(args, "range", 2, 3,
                                &start, &stop, &step))
             goto Fail;
 
         /* Convert borrowed refs to owned refs */
         start = PyNumber_Index(start);
         stop = PyNumber_Index(stop);
         step = validate_step(step);
         if (!start || !stop || !step)
             goto Fail;
     }
 
+    length = range_compute_length(start, stop, step);
+    if (length == NULL)
+        goto Fail;
     obj = PyObject_New(rangeobject, &PyRange_Type);
     if (obj == NULL)
         goto Fail;
     obj->start = start;
     obj->stop = stop;
     obj->step = step;
+    obj->length = length;
     return (PyObject *) obj;
 
 Fail:
     Py_XDECREF(start);
     Py_XDECREF(stop);
     Py_XDECREF(step);
+    Py_XDECREF(length);
     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.");
+Return an arithmetic progression of numbers from start "
+             "to stop (exclusive) by step.");
 
 static void
 range_dealloc(rangeobject *r)
 {
     Py_DECREF(r->start);
     Py_DECREF(r->stop);
     Py_DECREF(r->step);
+    Py_DECREF(r->length);
     PyObject_Del(r);
 }
 
-/* Return number of items in range (lo, hi, step), when arguments are
- * PyInt or PyLong objects.  step > 0 required.  Return a value < 0 if
- * & only if the true value is too large to fit in a signed long.
- * Arguments MUST return 1 with either PyLong_Check() or
- * PyLong_Check().  Return -1 when there is an error.
- */
-static PyObject*
-range_length_obj(rangeobject *r)
+static PyObject *
+range_compute_length(PyObject *start, PyObject *stop, PyObject *step)
 {
     /* -------------------------------------------------------------
     Algorithm is equal to that of get_len_of_range(), but it operates
     on PyObjects (which are assumed to be PyLong or PyInt objects).
     ---------------------------------------------------------------*/
     int cmp_result, cmp_call;
     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_call = PyObject_Cmp(r->step, zero, &cmp_result);
+    cmp_call = PyObject_Cmp(step, zero, &cmp_result);
     Py_DECREF(zero);
     if (cmp_call == -1)
         return NULL;
 
     assert(cmp_result != 0);
     if (cmp_result > 0) {
-        lo = r->start;
-        hi = r->stop;
-        step = r->step;
+        lo = start;
+        hi = stop;
         Py_INCREF(step);
     } else {
-        lo = r->stop;
-        hi = r->start;
-        step = PyNumber_Negative(r->step);
+        lo = stop;
+        hi = start;
+        step = PyNumber_Negative(step);
         if (!step)
-            return NULL;
+            goto Fail;
     }
 
     /* if (lo >= hi), return length of 0. */
     if (PyObject_Compare(lo, hi) >= 0) {
         Py_XDECREF(step);
         return PyLong_FromLong(0);
     }
 
     if ((one = PyLong_FromLong(1L)) == NULL)
         goto Fail;
@@ skipping 22 matching lines @@
     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;
+    return PyLong_AsSsize_t(r->length);
 }
 
-/* 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);
 }
 
-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 PyMemberDef range_members[] = {
+    {"start", T_OBJECT, offsetof(rangeobject, start), READONLY},
+    {"stop", T_OBJECT, offsetof(rangeobject, stop), READONLY},
+    {"step", T_OBJECT, offsetof(rangeobject, step), READONLY},
+    {0}
 };
 
 static PyObject * range_iter(PyObject *seq);
 static PyObject * range_reverse(PyObject *seq);
 
 PyDoc_STRVAR(reverse_doc,
 "Returns a reverse iterator.");
+
+static PySequenceMethods range_as_sequence = {
+    (lenfunc)range_length,  /* sq_length */
+    0,                      /* sq_concat */
+    0,                      /* sq_repeat */
+    0,                      /* sq_item */
+    0,                      /* sq_slice */
+    0,                      /* sq_ass_item */
+    0,                      /* sq_ass_slice */
+    0,                      /* sq_contains */
+    0,                      /* sq_inplace_concat */
+    0,                      /* sq_inplace_repeat */
+};
 
 static PyMethodDef range_methods[] = {
     {"__reversed__",    (PyCFunction)range_reverse, METH_NOARGS,
         reverse_doc},
     {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_compare */
         (reprfunc)range_repr,   /* tp_repr */
         0,                      /* tp_as_number */
-        &range_as_sequence,     /* tp_as_sequence */
-        0,                      /* tp_as_mapping */
+        &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 */
+        range_members,          /* 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 */
 };
 
@@ skipping 24 matching lines @@
 rangeiter_len(rangeiterobject *r)
 {
     return PyLong_FromLong(r->len - r->index);
 }
 
 typedef struct {
     PyObject_HEAD
     PyObject *index;
     PyObject *start;
     PyObject *step;
-    PyObject *len;
+    PyObject *length;
 } longrangeiterobject;
 
 static PyObject *
 longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
 {
-    return PyNumber_Subtract(r->len, r->index);
+    return PyNumber_Subtract(r->length, 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},
@@ skipping 107 matching lines @@
         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);
+    Py_XDECREF(r->length);
     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)
+    if (PyObject_RichCompareBool(r->index, r->length, 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;
@@ skipping 47 matching lines @@
         (iternextfunc)longrangeiter_next,       /* tp_iternext */
         longrangeiter_methods,                  /* tp_methods */
         0,
 };
 
 static PyObject *
 range_iter(PyObject *seq)
 {
     rangeobject *r = (rangeobject *)seq;
     longrangeiterobject *it;
-    PyObject *tmp, *len;
     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;
-
-    /* Calculate length: (r->stop - r->start) / r->step */
-    tmp = PyNumber_Subtract(r->stop, r->start);
-    if (!tmp)
-        goto create_failure;
-    len = PyNumber_FloorDivide(tmp, r->step);
-    Py_DECREF(tmp);
-    if (!len)
-        goto create_failure;
-    it->len = len;
+    Py_INCREF(r->length);
+    it->length = r->length;
     it->index = PyLong_FromLong(0);
-    if (!it->index)
+    if (it->index == NULL)
         goto create_failure;
 
     return (PyObject *)it;
 
 create_failure:
     Py_DECREF(it);
     return NULL;
 }
 
 static PyObject *
@@ skipping 28 matching lines @@
             }
         }
     }
     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;
+    Py_INCREF(range->length);
+    len = range->length;
 
     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->length = 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;
 }