range: lean and mean

Objects/rangeobject.c

Index: Objects/rangeobject.c
===================================================================
--- Objects/rangeobject.c	(revision 62662)
+++ Objects/rangeobject.c	(working copy)
@@ -1,6 +1,7 @@
 /* Range object implementation */
 
 #include "Python.h"
+#include "structmember.h"
 
 /* Support objects whose length is > PY_SSIZE_T_MAX.
 
@@ -14,6 +15,7 @@
     PyObject *start;
     PyObject *stop;
     PyObject *step;
+    PyObject *length;
 } rangeobject;
 
 /* Helper function for validating step.  Always returns a new reference or
@@ -43,6 +45,9 @@
     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)
@@ -52,7 +57,7 @@
 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;
@@ -81,25 +86,31 @@
             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)
@@ -107,17 +118,15 @@
     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)
+/*Return number of items in range (lo, hi, step), when arguments are
+* PyLong objects. step > 0 required. Return a PyLong with the length.
+* Return NULL when there is an error.*/
+static PyObject *
+range_compute_length(PyObject *start, PyObject *stop, PyObject *step)
 {
     /* -------------------------------------------------------------
     Algorithm is equal to that of get_len_of_range(), but it operates
@@ -125,7 +134,6 @@
     ---------------------------------------------------------------*/
     int cmp_result, cmp_call;
     PyObject *lo, *hi;
-    PyObject *step = NULL;
     PyObject *diff = NULL;
     PyObject *one = NULL;
     PyObject *tmp1 = NULL, *tmp2 = NULL, *result;
@@ -134,23 +142,22 @@
     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. */
@@ -193,46 +200,11 @@
 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;
@@ -252,12 +224,11 @@
                                     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);
@@ -266,6 +237,19 @@
 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},
@@ -284,8 +268,8 @@
 	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 */
@@ -301,8 +285,8 @@
 	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 */
@@ -347,13 +331,13 @@
     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);
@@ -481,7 +465,7 @@
     Py_XDECREF(r->index);
     Py_XDECREF(r->start);
     Py_XDECREF(r->step);
-    Py_XDECREF(r->len);
+    Py_XDECREF(r->length);
     PyObject_Del(r);
 }
 
@@ -489,7 +473,7 @@
 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);
@@ -557,7 +541,6 @@
 {
     rangeobject *r = (rangeobject *)seq;
     longrangeiterobject *it;
-    PyObject *tmp, *len;
     long lstart, lstop, lstep;
 
     assert(PyRange_Check(seq));
@@ -586,19 +569,10 @@
     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;
@@ -647,9 +621,8 @@
         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)
@@ -677,7 +650,7 @@
     }
 
     /* Steal reference to len. */
-    it->len = len;
+    it->length = len;
 
     it->index = PyLong_FromLong(0);
     if (!it->index) {