1 files changed, 234 insertions, 3 deletions

diff --git a/Objects/rangeobject.c b/Objects/rangeobject.c
index cebccff..214b455 100644
--- a/Objects/rangeobject.c
+++ b/Objects/rangeobject.c
@@ -1,6 +1,7 @@
 /* Range object implementation */
 
 #include "Python.h"
+#include "structmember.h"
 
 /* Support objects whose length is > PY_SSIZE_T_MAX.
 
@@ -610,6 +611,137 @@ range_contains(rangeobject *r, PyObject *ob)
                                        PY_ITERSEARCH_CONTAINS);
 }
 
+/* Compare two range objects.  Return 1 for equal, 0 for not equal
+   and -1 on error.  The algorithm is roughly the C equivalent of
+
+   if r0 is r1:
+       return True
+   if len(r0) != len(r1):
+       return False
+   if not len(r0):
+       return True
+   if r0.start != r1.start:
+       return False
+   if len(r0) == 1:
+       return True
+   return r0.step == r1.step
+*/
+static int
+range_equals(rangeobject *r0, rangeobject *r1)
+{
+    int cmp_result;
+    PyObject *one;
+
+    if (r0 == r1)
+        return 1;
+    cmp_result = PyObject_RichCompareBool(r0->length, r1->length, Py_EQ);
+    /* Return False or error to the caller. */
+    if (cmp_result != 1)
+        return cmp_result;
+    cmp_result = PyObject_Not(r0->length);
+    /* Return True or error to the caller. */
+    if (cmp_result != 0)
+        return cmp_result;
+    cmp_result = PyObject_RichCompareBool(r0->start, r1->start, Py_EQ);
+    /* Return False or error to the caller. */
+    if (cmp_result != 1)
+        return cmp_result;
+    one = PyLong_FromLong(1);
+    if (!one)
+        return -1;
+    cmp_result = PyObject_RichCompareBool(r0->length, one, Py_EQ);
+    Py_DECREF(one);
+    /* Return True or error to the caller. */
+    if (cmp_result != 0)
+        return cmp_result;
+    return PyObject_RichCompareBool(r0->step, r1->step, Py_EQ);
+}
+
+static PyObject *
+range_richcompare(PyObject *self, PyObject *other, int op)
+{
+    int result;
+
+    if (!PyRange_Check(other))
+        Py_RETURN_NOTIMPLEMENTED;
+    switch (op) {
+    case Py_NE:
+    case Py_EQ:
+        result = range_equals((rangeobject*)self, (rangeobject*)other);
+        if (result == -1)
+            return NULL;
+        if (op == Py_NE)
+            result = !result;
+        if (result)
+            Py_RETURN_TRUE;
+        else
+            Py_RETURN_FALSE;
+    case Py_LE:
+    case Py_GE:
+    case Py_LT:
+    case Py_GT:
+        Py_RETURN_NOTIMPLEMENTED;
+    default:
+        PyErr_BadArgument();
+        return NULL;
+    }
+}
+
+/* Hash function for range objects.  Rough C equivalent of
+
+   if not len(r):
+       return hash((len(r), None, None))
+   if len(r) == 1:
+       return hash((len(r), r.start, None))
+   return hash((len(r), r.start, r.step))
+*/
+static Py_hash_t
+range_hash(rangeobject *r)
+{
+    PyObject *t;
+    Py_hash_t result = -1;
+    int cmp_result;
+
+    t = PyTuple_New(3);
+    if (!t)
+        return -1;
+    Py_INCREF(r->length);
+    PyTuple_SET_ITEM(t, 0, r->length);
+    cmp_result = PyObject_Not(r->length);
+    if (cmp_result == -1)
+        goto end;
+    if (cmp_result == 1) {
+        Py_INCREF(Py_None);
+        Py_INCREF(Py_None);
+        PyTuple_SET_ITEM(t, 1, Py_None);
+        PyTuple_SET_ITEM(t, 2, Py_None);
+    }
+    else {
+        PyObject *one;
+        Py_INCREF(r->start);
+        PyTuple_SET_ITEM(t, 1, r->start);
+        one = PyLong_FromLong(1);
+        if (!one)
+            goto end;
+        cmp_result = PyObject_RichCompareBool(r->length, one, Py_EQ);
+        Py_DECREF(one);
+        if (cmp_result == -1)
+            goto end;
+        if (cmp_result == 1) {
+            Py_INCREF(Py_None);
+            PyTuple_SET_ITEM(t, 2, Py_None);
+        }
+        else {
+            Py_INCREF(r->step);
+            PyTuple_SET_ITEM(t, 2, r->step);
+        }
+    }
+    result = PyObject_Hash(t);
+  end:
+    Py_DECREF(t);
+    return result;
+}
+
 static PyObject *
 range_count(rangeobject *r, PyObject *ob)
 {
@@ -750,6 +882,13 @@ static PyMethodDef range_methods[] = {
     {NULL,              NULL}           /* sentinel */
 };
 
+static PyMemberDef range_members[] = {
+    {"start",   T_OBJECT_EX,    offsetof(rangeobject, start),   READONLY},
+    {"stop",    T_OBJECT_EX,    offsetof(rangeobject, stop),    READONLY},
+    {"step",    T_OBJECT_EX,    offsetof(rangeobject, step),    READONLY},
+    {0}
+};
+
 PyTypeObject PyRange_Type = {
         PyVarObject_HEAD_INIT(&PyType_Type, 0)
         "range",                /* Name of this type */
@@ -764,7 +903,7 @@ PyTypeObject PyRange_Type = {
         0,                      /* tp_as_number */
         &range_as_sequence,     /* tp_as_sequence */
         &range_as_mapping,      /* tp_as_mapping */
-        0,                      /* tp_hash */
+        (hashfunc)range_hash,   /* tp_hash */
         0,                      /* tp_call */
         0,                      /* tp_str */
         PyObject_GenericGetAttr,  /* tp_getattro */
@@ -774,12 +913,12 @@ PyTypeObject PyRange_Type = {
         range_doc,              /* tp_doc */
         0,                      /* tp_traverse */
         0,                      /* tp_clear */
-        0,                      /* tp_richcompare */
+        range_richcompare,      /* tp_richcompare */
         0,                      /* tp_weaklistoffset */
         range_iter,             /* tp_iter */
         0,                      /* tp_iternext */
         range_methods,          /* tp_methods */
-        0,                      /* tp_members */
+        range_members,          /* tp_members */
         0,                      /* tp_getset */
         0,                      /* tp_base */
         0,                      /* tp_dict */
@@ -826,9 +965,59 @@ rangeiter_len(rangeiterobject *r)
 PyDoc_STRVAR(length_hint_doc,
              "Private method returning an estimate of len(list(it)).");
 
+static PyObject *
+rangeiter_reduce(rangeiterobject *r)
+{
+    PyObject *start=NULL, *stop=NULL, *step=NULL;
+    PyObject *range;
+
+    /* create a range object for pickling */
+    start = PyLong_FromLong(r->start);
+    if (start == NULL)
+        goto err;
+    stop = PyLong_FromLong(r->start + r->len * r->step);
+    if (stop == NULL)
+        goto err;
+    step = PyLong_FromLong(r->step);
+    if (step == NULL)
+        goto err;
+    range = (PyObject*)make_range_object(&PyRange_Type,
+                               start, stop, step);
+    if (range == NULL)
+        goto err;
+    /* return the result */
+    return Py_BuildValue("N(N)i", _PyObject_GetBuiltin("iter"), range, r->index);
+err:
+    Py_XDECREF(start);
+    Py_XDECREF(stop);
+    Py_XDECREF(step);
+    return NULL;
+}
+
+static PyObject *
+rangeiter_setstate(rangeiterobject *r, PyObject *state)
+{
+    long index = PyLong_AsLong(state);
+    if (index == -1 && PyErr_Occurred())
+        return NULL;
+    if (index < 0 || index >= r->len) {
+        PyErr_SetString(PyExc_ValueError, "index out of range");
+        return NULL;
+    }
+    r->index = index;
+    Py_RETURN_NONE;
+}
+
+PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
+PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
+
 static PyMethodDef rangeiter_methods[] = {
     {"__length_hint__", (PyCFunction)rangeiter_len, METH_NOARGS,
         length_hint_doc},
+    {"__reduce__", (PyCFunction)rangeiter_reduce, METH_NOARGS,
+        reduce_doc},
+    {"__setstate__", (PyCFunction)rangeiter_setstate, METH_O,
+        setstate_doc},
     {NULL,              NULL}           /* sentinel */
 };
 
@@ -957,9 +1146,51 @@ longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
     return PyNumber_Subtract(r->len, r->index);
 }
 
+static PyObject *
+longrangeiter_reduce(longrangeiterobject *r)
+{
+    PyObject *product, *stop=NULL;
+    PyObject *range;
+
+    /* create a range object for pickling.  Must calculate the "stop" value */
+    product = PyNumber_Multiply(r->len, r->step);
+    if (product == NULL)
+        return NULL;
+    stop = PyNumber_Add(r->start, product);
+    Py_DECREF(product);
+    if (stop ==  NULL)
+        return NULL;
+    Py_INCREF(r->start);
+    Py_INCREF(r->step);
+    range =  (PyObject*)make_range_object(&PyRange_Type,
+                               r->start, stop, r->step);
+    if (range == NULL) {
+        Py_DECREF(r->start);
+        Py_DECREF(stop);
+        Py_DECREF(r->step);
+        return NULL;
+    }
+
+    /* return the result */
+    return Py_BuildValue("N(N)O", _PyObject_GetBuiltin("iter"), range, r->index);
+}
+
+static PyObject *
+longrangeiter_setstate(longrangeiterobject *r, PyObject *state)
+{
+    Py_CLEAR(r->index);
+    r->index = state;
+    Py_INCREF(r->index);
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef longrangeiter_methods[] = {
     {"__length_hint__", (PyCFunction)longrangeiter_len, METH_NOARGS,
         length_hint_doc},
+    {"__reduce__", (PyCFunction)longrangeiter_reduce, METH_NOARGS,
+        reduce_doc},
+    {"__setstate__", (PyCFunction)longrangeiter_setstate, METH_O,
+        setstate_doc},
     {NULL,              NULL}           /* sentinel */
 };