xref: /third_party/python/Objects/rangeobject.c (revision 7db96d56)
1/* Range object implementation */
2
3#include "Python.h"
4#include "pycore_abstract.h"      // _PyIndex_Check()
5#include "pycore_long.h"          // _PyLong_GetZero()
6#include "pycore_tuple.h"         // _PyTuple_ITEMS()
7#include "structmember.h"         // PyMemberDef
8
9/* Support objects whose length is > PY_SSIZE_T_MAX.
10
11   This could be sped up for small PyLongs if they fit in a Py_ssize_t.
12   This only matters on Win64.  Though we could use long long which
13   would presumably help perf.
14*/
15
16typedef struct {
17    PyObject_HEAD
18    PyObject *start;
19    PyObject *stop;
20    PyObject *step;
21    PyObject *length;
22} rangeobject;
23
24/* Helper function for validating step.  Always returns a new reference or
25   NULL on error.
26*/
27static PyObject *
28validate_step(PyObject *step)
29{
30    /* No step specified, use a step of 1. */
31    if (!step)
32        return PyLong_FromLong(1);
33
34    step = PyNumber_Index(step);
35    if (step && _PyLong_Sign(step) == 0) {
36        PyErr_SetString(PyExc_ValueError,
37                        "range() arg 3 must not be zero");
38        Py_CLEAR(step);
39    }
40
41    return step;
42}
43
44static PyObject *
45compute_range_length(PyObject *start, PyObject *stop, PyObject *step);
46
47static rangeobject *
48make_range_object(PyTypeObject *type, PyObject *start,
49                  PyObject *stop, PyObject *step)
50{
51    rangeobject *obj = NULL;
52    PyObject *length;
53    length = compute_range_length(start, stop, step);
54    if (length == NULL) {
55        return NULL;
56    }
57    obj = PyObject_New(rangeobject, type);
58    if (obj == NULL) {
59        Py_DECREF(length);
60        return NULL;
61    }
62    obj->start = start;
63    obj->stop = stop;
64    obj->step = step;
65    obj->length = length;
66    return obj;
67}
68
69/* XXX(nnorwitz): should we error check if the user passes any empty ranges?
70   range(-10)
71   range(0, -5)
72   range(0, 5, -1)
73*/
74static PyObject *
75range_from_array(PyTypeObject *type, PyObject *const *args, Py_ssize_t num_args)
76{
77    rangeobject *obj;
78    PyObject *start = NULL, *stop = NULL, *step = NULL;
79
80    switch (num_args) {
81        case 3:
82            step = args[2];
83            /* fallthrough */
84        case 2:
85            /* Convert borrowed refs to owned refs */
86            start = PyNumber_Index(args[0]);
87            if (!start) {
88                return NULL;
89            }
90            stop = PyNumber_Index(args[1]);
91            if (!stop) {
92                Py_DECREF(start);
93                return NULL;
94            }
95            step = validate_step(step);  /* Caution, this can clear exceptions */
96            if (!step) {
97                Py_DECREF(start);
98                Py_DECREF(stop);
99                return NULL;
100            }
101            break;
102        case 1:
103            stop = PyNumber_Index(args[0]);
104            if (!stop) {
105                return NULL;
106            }
107            start = _PyLong_GetZero();
108            Py_INCREF(start);
109            step = _PyLong_GetOne();
110            Py_INCREF(step);
111            break;
112        case 0:
113            PyErr_SetString(PyExc_TypeError,
114                            "range expected at least 1 argument, got 0");
115            return NULL;
116        default:
117            PyErr_Format(PyExc_TypeError,
118                         "range expected at most 3 arguments, got %zd",
119                         num_args);
120            return NULL;
121    }
122    obj = make_range_object(type, start, stop, step);
123    if (obj != NULL) {
124        return (PyObject *) obj;
125    }
126
127    /* Failed to create object, release attributes */
128    Py_DECREF(start);
129    Py_DECREF(stop);
130    Py_DECREF(step);
131    return NULL;
132}
133
134static PyObject *
135range_new(PyTypeObject *type, PyObject *args, PyObject *kw)
136{
137    if (!_PyArg_NoKeywords("range", kw))
138        return NULL;
139
140    return range_from_array(type, _PyTuple_ITEMS(args), PyTuple_GET_SIZE(args));
141}
142
143
144static PyObject *
145range_vectorcall(PyTypeObject *type, PyObject *const *args,
146                 size_t nargsf, PyObject *kwnames)
147{
148    Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
149    if (!_PyArg_NoKwnames("range", kwnames)) {
150        return NULL;
151    }
152    return range_from_array(type, args, nargs);
153}
154
155PyDoc_STRVAR(range_doc,
156"range(stop) -> range object\n\
157range(start, stop[, step]) -> range object\n\
158\n\
159Return an object that produces a sequence of integers from start (inclusive)\n\
160to stop (exclusive) by step.  range(i, j) produces i, i+1, i+2, ..., j-1.\n\
161start defaults to 0, and stop is omitted!  range(4) produces 0, 1, 2, 3.\n\
162These are exactly the valid indices for a list of 4 elements.\n\
163When step is given, it specifies the increment (or decrement).");
164
165static void
166range_dealloc(rangeobject *r)
167{
168    Py_DECREF(r->start);
169    Py_DECREF(r->stop);
170    Py_DECREF(r->step);
171    Py_DECREF(r->length);
172    PyObject_Free(r);
173}
174
175/* Return number of items in range (lo, hi, step) as a PyLong object,
176 * when arguments are PyLong objects.  Arguments MUST return 1 with
177 * PyLong_Check().  Return NULL when there is an error.
178 */
179static PyObject*
180compute_range_length(PyObject *start, PyObject *stop, PyObject *step)
181{
182    /* -------------------------------------------------------------
183    Algorithm is equal to that of get_len_of_range(), but it operates
184    on PyObjects (which are assumed to be PyLong objects).
185    ---------------------------------------------------------------*/
186    int cmp_result;
187    PyObject *lo, *hi;
188    PyObject *diff = NULL;
189    PyObject *tmp1 = NULL, *tmp2 = NULL, *result;
190                /* holds sub-expression evaluations */
191
192    PyObject *zero = _PyLong_GetZero();  // borrowed reference
193    PyObject *one = _PyLong_GetOne();  // borrowed reference
194
195    cmp_result = PyObject_RichCompareBool(step, zero, Py_GT);
196    if (cmp_result == -1)
197        return NULL;
198
199    if (cmp_result == 1) {
200        lo = start;
201        hi = stop;
202        Py_INCREF(step);
203    } else {
204        lo = stop;
205        hi = start;
206        step = PyNumber_Negative(step);
207        if (!step)
208            return NULL;
209    }
210
211    /* if (lo >= hi), return length of 0. */
212    cmp_result = PyObject_RichCompareBool(lo, hi, Py_GE);
213    if (cmp_result != 0) {
214        Py_DECREF(step);
215        if (cmp_result < 0)
216            return NULL;
217        result = zero;
218        Py_INCREF(result);
219        return result;
220    }
221
222    if ((tmp1 = PyNumber_Subtract(hi, lo)) == NULL)
223        goto Fail;
224
225    if ((diff = PyNumber_Subtract(tmp1, one)) == NULL)
226        goto Fail;
227
228    if ((tmp2 = PyNumber_FloorDivide(diff, step)) == NULL)
229        goto Fail;
230
231    if ((result = PyNumber_Add(tmp2, one)) == NULL)
232        goto Fail;
233
234    Py_DECREF(tmp2);
235    Py_DECREF(diff);
236    Py_DECREF(step);
237    Py_DECREF(tmp1);
238    return result;
239
240  Fail:
241    Py_DECREF(step);
242    Py_XDECREF(tmp2);
243    Py_XDECREF(diff);
244    Py_XDECREF(tmp1);
245    return NULL;
246}
247
248static Py_ssize_t
249range_length(rangeobject *r)
250{
251    return PyLong_AsSsize_t(r->length);
252}
253
254static PyObject *
255compute_item(rangeobject *r, PyObject *i)
256{
257    PyObject *incr, *result;
258    /* PyLong equivalent to:
259     *    return r->start + (i * r->step)
260     */
261    if (r->step == _PyLong_GetOne()) {
262        result = PyNumber_Add(r->start, i);
263    }
264    else {
265        incr = PyNumber_Multiply(i, r->step);
266        if (!incr) {
267            return NULL;
268        }
269        result = PyNumber_Add(r->start, incr);
270        Py_DECREF(incr);
271    }
272    return result;
273}
274
275static PyObject *
276compute_range_item(rangeobject *r, PyObject *arg)
277{
278    PyObject *zero = _PyLong_GetZero();  // borrowed reference
279    int cmp_result;
280    PyObject *i, *result;
281
282    /* PyLong equivalent to:
283     *   if (arg < 0) {
284     *     i = r->length + arg
285     *   } else {
286     *     i = arg
287     *   }
288     */
289    cmp_result = PyObject_RichCompareBool(arg, zero, Py_LT);
290    if (cmp_result == -1) {
291        return NULL;
292    }
293    if (cmp_result == 1) {
294        i = PyNumber_Add(r->length, arg);
295        if (!i) {
296          return NULL;
297        }
298    } else {
299        i = arg;
300        Py_INCREF(i);
301    }
302
303    /* PyLong equivalent to:
304     *   if (i < 0 || i >= r->length) {
305     *     <report index out of bounds>
306     *   }
307     */
308    cmp_result = PyObject_RichCompareBool(i, zero, Py_LT);
309    if (cmp_result == 0) {
310        cmp_result = PyObject_RichCompareBool(i, r->length, Py_GE);
311    }
312    if (cmp_result == -1) {
313       Py_DECREF(i);
314       return NULL;
315    }
316    if (cmp_result == 1) {
317        Py_DECREF(i);
318        PyErr_SetString(PyExc_IndexError,
319                        "range object index out of range");
320        return NULL;
321    }
322
323    result = compute_item(r, i);
324    Py_DECREF(i);
325    return result;
326}
327
328static PyObject *
329range_item(rangeobject *r, Py_ssize_t i)
330{
331    PyObject *res, *arg = PyLong_FromSsize_t(i);
332    if (!arg) {
333        return NULL;
334    }
335    res = compute_range_item(r, arg);
336    Py_DECREF(arg);
337    return res;
338}
339
340static PyObject *
341compute_slice(rangeobject *r, PyObject *_slice)
342{
343    PySliceObject *slice = (PySliceObject *) _slice;
344    rangeobject *result;
345    PyObject *start = NULL, *stop = NULL, *step = NULL;
346    PyObject *substart = NULL, *substop = NULL, *substep = NULL;
347    int error;
348
349    error = _PySlice_GetLongIndices(slice, r->length, &start, &stop, &step);
350    if (error == -1)
351        return NULL;
352
353    substep = PyNumber_Multiply(r->step, step);
354    if (substep == NULL) goto fail;
355    Py_CLEAR(step);
356
357    substart = compute_item(r, start);
358    if (substart == NULL) goto fail;
359    Py_CLEAR(start);
360
361    substop = compute_item(r, stop);
362    if (substop == NULL) goto fail;
363    Py_CLEAR(stop);
364
365    result = make_range_object(Py_TYPE(r), substart, substop, substep);
366    if (result != NULL) {
367        return (PyObject *) result;
368    }
369fail:
370    Py_XDECREF(start);
371    Py_XDECREF(stop);
372    Py_XDECREF(step);
373    Py_XDECREF(substart);
374    Py_XDECREF(substop);
375    Py_XDECREF(substep);
376    return NULL;
377}
378
379/* Assumes (PyLong_CheckExact(ob) || PyBool_Check(ob)) */
380static int
381range_contains_long(rangeobject *r, PyObject *ob)
382{
383    PyObject *zero = _PyLong_GetZero();  // borrowed reference
384    int cmp1, cmp2, cmp3;
385    PyObject *tmp1 = NULL;
386    PyObject *tmp2 = NULL;
387    int result = -1;
388
389    /* Check if the value can possibly be in the range. */
390
391    cmp1 = PyObject_RichCompareBool(r->step, zero, Py_GT);
392    if (cmp1 == -1)
393        goto end;
394    if (cmp1 == 1) { /* positive steps: start <= ob < stop */
395        cmp2 = PyObject_RichCompareBool(r->start, ob, Py_LE);
396        cmp3 = PyObject_RichCompareBool(ob, r->stop, Py_LT);
397    }
398    else { /* negative steps: stop < ob <= start */
399        cmp2 = PyObject_RichCompareBool(ob, r->start, Py_LE);
400        cmp3 = PyObject_RichCompareBool(r->stop, ob, Py_LT);
401    }
402
403    if (cmp2 == -1 || cmp3 == -1) /* TypeError */
404        goto end;
405    if (cmp2 == 0 || cmp3 == 0) { /* ob outside of range */
406        result = 0;
407        goto end;
408    }
409
410    /* Check that the stride does not invalidate ob's membership. */
411    tmp1 = PyNumber_Subtract(ob, r->start);
412    if (tmp1 == NULL)
413        goto end;
414    tmp2 = PyNumber_Remainder(tmp1, r->step);
415    if (tmp2 == NULL)
416        goto end;
417    /* result = ((int(ob) - start) % step) == 0 */
418    result = PyObject_RichCompareBool(tmp2, zero, Py_EQ);
419  end:
420    Py_XDECREF(tmp1);
421    Py_XDECREF(tmp2);
422    return result;
423}
424
425static int
426range_contains(rangeobject *r, PyObject *ob)
427{
428    if (PyLong_CheckExact(ob) || PyBool_Check(ob))
429        return range_contains_long(r, ob);
430
431    return (int)_PySequence_IterSearch((PyObject*)r, ob,
432                                       PY_ITERSEARCH_CONTAINS);
433}
434
435/* Compare two range objects.  Return 1 for equal, 0 for not equal
436   and -1 on error.  The algorithm is roughly the C equivalent of
437
438   if r0 is r1:
439       return True
440   if len(r0) != len(r1):
441       return False
442   if not len(r0):
443       return True
444   if r0.start != r1.start:
445       return False
446   if len(r0) == 1:
447       return True
448   return r0.step == r1.step
449*/
450static int
451range_equals(rangeobject *r0, rangeobject *r1)
452{
453    int cmp_result;
454
455    if (r0 == r1)
456        return 1;
457    cmp_result = PyObject_RichCompareBool(r0->length, r1->length, Py_EQ);
458    /* Return False or error to the caller. */
459    if (cmp_result != 1)
460        return cmp_result;
461    cmp_result = PyObject_Not(r0->length);
462    /* Return True or error to the caller. */
463    if (cmp_result != 0)
464        return cmp_result;
465    cmp_result = PyObject_RichCompareBool(r0->start, r1->start, Py_EQ);
466    /* Return False or error to the caller. */
467    if (cmp_result != 1)
468        return cmp_result;
469    cmp_result = PyObject_RichCompareBool(r0->length, _PyLong_GetOne(), Py_EQ);
470    /* Return True or error to the caller. */
471    if (cmp_result != 0)
472        return cmp_result;
473    return PyObject_RichCompareBool(r0->step, r1->step, Py_EQ);
474}
475
476static PyObject *
477range_richcompare(PyObject *self, PyObject *other, int op)
478{
479    int result;
480
481    if (!PyRange_Check(other))
482        Py_RETURN_NOTIMPLEMENTED;
483    switch (op) {
484    case Py_NE:
485    case Py_EQ:
486        result = range_equals((rangeobject*)self, (rangeobject*)other);
487        if (result == -1)
488            return NULL;
489        if (op == Py_NE)
490            result = !result;
491        if (result)
492            Py_RETURN_TRUE;
493        else
494            Py_RETURN_FALSE;
495    case Py_LE:
496    case Py_GE:
497    case Py_LT:
498    case Py_GT:
499        Py_RETURN_NOTIMPLEMENTED;
500    default:
501        PyErr_BadArgument();
502        return NULL;
503    }
504}
505
506/* Hash function for range objects.  Rough C equivalent of
507
508   if not len(r):
509       return hash((len(r), None, None))
510   if len(r) == 1:
511       return hash((len(r), r.start, None))
512   return hash((len(r), r.start, r.step))
513*/
514static Py_hash_t
515range_hash(rangeobject *r)
516{
517    PyObject *t;
518    Py_hash_t result = -1;
519    int cmp_result;
520
521    t = PyTuple_New(3);
522    if (!t)
523        return -1;
524    Py_INCREF(r->length);
525    PyTuple_SET_ITEM(t, 0, r->length);
526    cmp_result = PyObject_Not(r->length);
527    if (cmp_result == -1)
528        goto end;
529    if (cmp_result == 1) {
530        Py_INCREF(Py_None);
531        Py_INCREF(Py_None);
532        PyTuple_SET_ITEM(t, 1, Py_None);
533        PyTuple_SET_ITEM(t, 2, Py_None);
534    }
535    else {
536        Py_INCREF(r->start);
537        PyTuple_SET_ITEM(t, 1, r->start);
538        cmp_result = PyObject_RichCompareBool(r->length, _PyLong_GetOne(), Py_EQ);
539        if (cmp_result == -1)
540            goto end;
541        if (cmp_result == 1) {
542            Py_INCREF(Py_None);
543            PyTuple_SET_ITEM(t, 2, Py_None);
544        }
545        else {
546            Py_INCREF(r->step);
547            PyTuple_SET_ITEM(t, 2, r->step);
548        }
549    }
550    result = PyObject_Hash(t);
551  end:
552    Py_DECREF(t);
553    return result;
554}
555
556static PyObject *
557range_count(rangeobject *r, PyObject *ob)
558{
559    if (PyLong_CheckExact(ob) || PyBool_Check(ob)) {
560        int result = range_contains_long(r, ob);
561        if (result == -1)
562            return NULL;
563        return PyLong_FromLong(result);
564    } else {
565        Py_ssize_t count;
566        count = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_COUNT);
567        if (count == -1)
568            return NULL;
569        return PyLong_FromSsize_t(count);
570    }
571}
572
573static PyObject *
574range_index(rangeobject *r, PyObject *ob)
575{
576    int contains;
577
578    if (!PyLong_CheckExact(ob) && !PyBool_Check(ob)) {
579        Py_ssize_t index;
580        index = _PySequence_IterSearch((PyObject*)r, ob, PY_ITERSEARCH_INDEX);
581        if (index == -1)
582            return NULL;
583        return PyLong_FromSsize_t(index);
584    }
585
586    contains = range_contains_long(r, ob);
587    if (contains == -1)
588        return NULL;
589
590    if (contains) {
591        PyObject *idx = PyNumber_Subtract(ob, r->start);
592        if (idx == NULL) {
593            return NULL;
594        }
595
596        if (r->step == _PyLong_GetOne()) {
597            return idx;
598        }
599
600        /* idx = (ob - r.start) // r.step */
601        PyObject *sidx = PyNumber_FloorDivide(idx, r->step);
602        Py_DECREF(idx);
603        return sidx;
604    }
605
606    /* object is not in the range */
607    PyErr_Format(PyExc_ValueError, "%R is not in range", ob);
608    return NULL;
609}
610
611static PySequenceMethods range_as_sequence = {
612    (lenfunc)range_length,      /* sq_length */
613    0,                          /* sq_concat */
614    0,                          /* sq_repeat */
615    (ssizeargfunc)range_item,   /* sq_item */
616    0,                          /* sq_slice */
617    0,                          /* sq_ass_item */
618    0,                          /* sq_ass_slice */
619    (objobjproc)range_contains, /* sq_contains */
620};
621
622static PyObject *
623range_repr(rangeobject *r)
624{
625    Py_ssize_t istep;
626
627    /* Check for special case values for printing.  We don't always
628       need the step value.  We don't care about overflow. */
629    istep = PyNumber_AsSsize_t(r->step, NULL);
630    if (istep == -1 && PyErr_Occurred()) {
631        assert(!PyErr_ExceptionMatches(PyExc_OverflowError));
632        return NULL;
633    }
634
635    if (istep == 1)
636        return PyUnicode_FromFormat("range(%R, %R)", r->start, r->stop);
637    else
638        return PyUnicode_FromFormat("range(%R, %R, %R)",
639                                    r->start, r->stop, r->step);
640}
641
642/* Pickling support */
643static PyObject *
644range_reduce(rangeobject *r, PyObject *args)
645{
646    return Py_BuildValue("(O(OOO))", Py_TYPE(r),
647                         r->start, r->stop, r->step);
648}
649
650static PyObject *
651range_subscript(rangeobject* self, PyObject* item)
652{
653    if (_PyIndex_Check(item)) {
654        PyObject *i, *result;
655        i = PyNumber_Index(item);
656        if (!i)
657            return NULL;
658        result = compute_range_item(self, i);
659        Py_DECREF(i);
660        return result;
661    }
662    if (PySlice_Check(item)) {
663        return compute_slice(self, item);
664    }
665    PyErr_Format(PyExc_TypeError,
666                 "range indices must be integers or slices, not %.200s",
667                 Py_TYPE(item)->tp_name);
668    return NULL;
669}
670
671
672static PyMappingMethods range_as_mapping = {
673        (lenfunc)range_length,       /* mp_length */
674        (binaryfunc)range_subscript, /* mp_subscript */
675        (objobjargproc)0,            /* mp_ass_subscript */
676};
677
678static int
679range_bool(rangeobject* self)
680{
681    return PyObject_IsTrue(self->length);
682}
683
684static PyNumberMethods range_as_number = {
685    .nb_bool = (inquiry)range_bool,
686};
687
688static PyObject * range_iter(PyObject *seq);
689static PyObject * range_reverse(PyObject *seq, PyObject *Py_UNUSED(ignored));
690
691PyDoc_STRVAR(reverse_doc,
692"Return a reverse iterator.");
693
694PyDoc_STRVAR(count_doc,
695"rangeobject.count(value) -> integer -- return number of occurrences of value");
696
697PyDoc_STRVAR(index_doc,
698"rangeobject.index(value) -> integer -- return index of value.\n"
699"Raise ValueError if the value is not present.");
700
701static PyMethodDef range_methods[] = {
702    {"__reversed__",    range_reverse,              METH_NOARGS, reverse_doc},
703    {"__reduce__",      (PyCFunction)range_reduce,  METH_VARARGS},
704    {"count",           (PyCFunction)range_count,   METH_O,      count_doc},
705    {"index",           (PyCFunction)range_index,   METH_O,      index_doc},
706    {NULL,              NULL}           /* sentinel */
707};
708
709static PyMemberDef range_members[] = {
710    {"start",   T_OBJECT_EX,    offsetof(rangeobject, start),   READONLY},
711    {"stop",    T_OBJECT_EX,    offsetof(rangeobject, stop),    READONLY},
712    {"step",    T_OBJECT_EX,    offsetof(rangeobject, step),    READONLY},
713    {0}
714};
715
716PyTypeObject PyRange_Type = {
717        PyVarObject_HEAD_INIT(&PyType_Type, 0)
718        "range",                /* Name of this type */
719        sizeof(rangeobject),    /* Basic object size */
720        0,                      /* Item size for varobject */
721        (destructor)range_dealloc, /* tp_dealloc */
722        0,                      /* tp_vectorcall_offset */
723        0,                      /* tp_getattr */
724        0,                      /* tp_setattr */
725        0,                      /* tp_as_async */
726        (reprfunc)range_repr,   /* tp_repr */
727        &range_as_number,       /* tp_as_number */
728        &range_as_sequence,     /* tp_as_sequence */
729        &range_as_mapping,      /* tp_as_mapping */
730        (hashfunc)range_hash,   /* tp_hash */
731        0,                      /* tp_call */
732        0,                      /* tp_str */
733        PyObject_GenericGetAttr,  /* tp_getattro */
734        0,                      /* tp_setattro */
735        0,                      /* tp_as_buffer */
736        Py_TPFLAGS_DEFAULT | Py_TPFLAGS_SEQUENCE,  /* tp_flags */
737        range_doc,              /* tp_doc */
738        0,                      /* tp_traverse */
739        0,                      /* tp_clear */
740        range_richcompare,      /* tp_richcompare */
741        0,                      /* tp_weaklistoffset */
742        range_iter,             /* tp_iter */
743        0,                      /* tp_iternext */
744        range_methods,          /* tp_methods */
745        range_members,          /* tp_members */
746        0,                      /* tp_getset */
747        0,                      /* tp_base */
748        0,                      /* tp_dict */
749        0,                      /* tp_descr_get */
750        0,                      /* tp_descr_set */
751        0,                      /* tp_dictoffset */
752        0,                      /* tp_init */
753        0,                      /* tp_alloc */
754        range_new,              /* tp_new */
755        .tp_vectorcall = (vectorcallfunc)range_vectorcall
756};
757
758/*********************** range Iterator **************************/
759
760/* There are 2 types of iterators, one for C longs, the other for
761   Python ints (ie, PyObjects).  This should make iteration fast
762   in the normal case, but possible for any numeric value.
763*/
764
765typedef struct {
766        PyObject_HEAD
767        long    index;
768        long    start;
769        long    step;
770        long    len;
771} rangeiterobject;
772
773static PyObject *
774rangeiter_next(rangeiterobject *r)
775{
776    if (r->index < r->len)
777        /* cast to unsigned to avoid possible signed overflow
778           in intermediate calculations. */
779        return PyLong_FromLong((long)(r->start +
780                                      (unsigned long)(r->index++) * r->step));
781    return NULL;
782}
783
784static PyObject *
785rangeiter_len(rangeiterobject *r, PyObject *Py_UNUSED(ignored))
786{
787    return PyLong_FromLong(r->len - r->index);
788}
789
790PyDoc_STRVAR(length_hint_doc,
791             "Private method returning an estimate of len(list(it)).");
792
793static PyObject *
794rangeiter_reduce(rangeiterobject *r, PyObject *Py_UNUSED(ignored))
795{
796    PyObject *start=NULL, *stop=NULL, *step=NULL;
797    PyObject *range;
798
799    /* create a range object for pickling */
800    start = PyLong_FromLong(r->start);
801    if (start == NULL)
802        goto err;
803    stop = PyLong_FromLong(r->start + r->len * r->step);
804    if (stop == NULL)
805        goto err;
806    step = PyLong_FromLong(r->step);
807    if (step == NULL)
808        goto err;
809    range = (PyObject*)make_range_object(&PyRange_Type,
810                               start, stop, step);
811    if (range == NULL)
812        goto err;
813    /* return the result */
814    return Py_BuildValue(
815            "N(N)l", _PyEval_GetBuiltin(&_Py_ID(iter)), range, r->index);
816err:
817    Py_XDECREF(start);
818    Py_XDECREF(stop);
819    Py_XDECREF(step);
820    return NULL;
821}
822
823static PyObject *
824rangeiter_setstate(rangeiterobject *r, PyObject *state)
825{
826    long index = PyLong_AsLong(state);
827    if (index == -1 && PyErr_Occurred())
828        return NULL;
829    /* silently clip the index value */
830    if (index < 0)
831        index = 0;
832    else if (index > r->len)
833        index = r->len; /* exhausted iterator */
834    r->index = index;
835    Py_RETURN_NONE;
836}
837
838PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
839PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
840
841static PyMethodDef rangeiter_methods[] = {
842    {"__length_hint__", (PyCFunction)rangeiter_len, METH_NOARGS,
843        length_hint_doc},
844    {"__reduce__", (PyCFunction)rangeiter_reduce, METH_NOARGS,
845        reduce_doc},
846    {"__setstate__", (PyCFunction)rangeiter_setstate, METH_O,
847        setstate_doc},
848    {NULL,              NULL}           /* sentinel */
849};
850
851PyTypeObject PyRangeIter_Type = {
852        PyVarObject_HEAD_INIT(&PyType_Type, 0)
853        "range_iterator",                        /* tp_name */
854        sizeof(rangeiterobject),                /* tp_basicsize */
855        0,                                      /* tp_itemsize */
856        /* methods */
857        (destructor)PyObject_Del,               /* tp_dealloc */
858        0,                                      /* tp_vectorcall_offset */
859        0,                                      /* tp_getattr */
860        0,                                      /* tp_setattr */
861        0,                                      /* tp_as_async */
862        0,                                      /* tp_repr */
863        0,                                      /* tp_as_number */
864        0,                                      /* tp_as_sequence */
865        0,                                      /* tp_as_mapping */
866        0,                                      /* tp_hash */
867        0,                                      /* tp_call */
868        0,                                      /* tp_str */
869        PyObject_GenericGetAttr,                /* tp_getattro */
870        0,                                      /* tp_setattro */
871        0,                                      /* tp_as_buffer */
872        Py_TPFLAGS_DEFAULT,                     /* tp_flags */
873        0,                                      /* tp_doc */
874        0,                                      /* tp_traverse */
875        0,                                      /* tp_clear */
876        0,                                      /* tp_richcompare */
877        0,                                      /* tp_weaklistoffset */
878        PyObject_SelfIter,                      /* tp_iter */
879        (iternextfunc)rangeiter_next,           /* tp_iternext */
880        rangeiter_methods,                      /* tp_methods */
881        0,                                      /* tp_members */
882};
883
884/* Return number of items in range (lo, hi, step).  step != 0
885 * required.  The result always fits in an unsigned long.
886 */
887static unsigned long
888get_len_of_range(long lo, long hi, long step)
889{
890    /* -------------------------------------------------------------
891    If step > 0 and lo >= hi, or step < 0 and lo <= hi, the range is empty.
892    Else for step > 0, if n values are in the range, the last one is
893    lo + (n-1)*step, which must be <= hi-1.  Rearranging,
894    n <= (hi - lo - 1)/step + 1, so taking the floor of the RHS gives
895    the proper value.  Since lo < hi in this case, hi-lo-1 >= 0, so
896    the RHS is non-negative and so truncation is the same as the
897    floor.  Letting M be the largest positive long, the worst case
898    for the RHS numerator is hi=M, lo=-M-1, and then
899    hi-lo-1 = M-(-M-1)-1 = 2*M.  Therefore unsigned long has enough
900    precision to compute the RHS exactly.  The analysis for step < 0
901    is similar.
902    ---------------------------------------------------------------*/
903    assert(step != 0);
904    if (step > 0 && lo < hi)
905        return 1UL + (hi - 1UL - lo) / step;
906    else if (step < 0 && lo > hi)
907        return 1UL + (lo - 1UL - hi) / (0UL - step);
908    else
909        return 0UL;
910}
911
912/* Initialize a rangeiter object.  If the length of the rangeiter object
913   is not representable as a C long, OverflowError is raised. */
914
915static PyObject *
916fast_range_iter(long start, long stop, long step, long len)
917{
918    rangeiterobject *it = PyObject_New(rangeiterobject, &PyRangeIter_Type);
919    if (it == NULL)
920        return NULL;
921    it->start = start;
922    it->step = step;
923    it->len = len;
924    it->index = 0;
925    return (PyObject *)it;
926}
927
928typedef struct {
929    PyObject_HEAD
930    PyObject *index;
931    PyObject *start;
932    PyObject *step;
933    PyObject *len;
934} longrangeiterobject;
935
936static PyObject *
937longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
938{
939    return PyNumber_Subtract(r->len, r->index);
940}
941
942static PyObject *
943longrangeiter_reduce(longrangeiterobject *r, PyObject *Py_UNUSED(ignored))
944{
945    PyObject *product, *stop=NULL;
946    PyObject *range;
947
948    /* create a range object for pickling.  Must calculate the "stop" value */
949    product = PyNumber_Multiply(r->len, r->step);
950    if (product == NULL)
951        return NULL;
952    stop = PyNumber_Add(r->start, product);
953    Py_DECREF(product);
954    if (stop ==  NULL)
955        return NULL;
956    Py_INCREF(r->start);
957    Py_INCREF(r->step);
958    range =  (PyObject*)make_range_object(&PyRange_Type,
959                               r->start, stop, r->step);
960    if (range == NULL) {
961        Py_DECREF(r->start);
962        Py_DECREF(stop);
963        Py_DECREF(r->step);
964        return NULL;
965    }
966
967    /* return the result */
968    return Py_BuildValue(
969            "N(N)O", _PyEval_GetBuiltin(&_Py_ID(iter)), range, r->index);
970}
971
972static PyObject *
973longrangeiter_setstate(longrangeiterobject *r, PyObject *state)
974{
975    PyObject *zero = _PyLong_GetZero();  // borrowed reference
976    int cmp;
977
978    /* clip the value */
979    cmp = PyObject_RichCompareBool(state, zero, Py_LT);
980    if (cmp < 0)
981        return NULL;
982    if (cmp > 0) {
983        state = zero;
984    }
985    else {
986        cmp = PyObject_RichCompareBool(r->len, state, Py_LT);
987        if (cmp < 0)
988            return NULL;
989        if (cmp > 0)
990            state = r->len;
991    }
992    Py_INCREF(state);
993    Py_XSETREF(r->index, state);
994    Py_RETURN_NONE;
995}
996
997static PyMethodDef longrangeiter_methods[] = {
998    {"__length_hint__", (PyCFunction)longrangeiter_len, METH_NOARGS,
999        length_hint_doc},
1000    {"__reduce__", (PyCFunction)longrangeiter_reduce, METH_NOARGS,
1001        reduce_doc},
1002    {"__setstate__", (PyCFunction)longrangeiter_setstate, METH_O,
1003        setstate_doc},
1004    {NULL,              NULL}           /* sentinel */
1005};
1006
1007static void
1008longrangeiter_dealloc(longrangeiterobject *r)
1009{
1010    Py_XDECREF(r->index);
1011    Py_XDECREF(r->start);
1012    Py_XDECREF(r->step);
1013    Py_XDECREF(r->len);
1014    PyObject_Free(r);
1015}
1016
1017static PyObject *
1018longrangeiter_next(longrangeiterobject *r)
1019{
1020    PyObject *product, *new_index, *result;
1021    if (PyObject_RichCompareBool(r->index, r->len, Py_LT) != 1)
1022        return NULL;
1023
1024    new_index = PyNumber_Add(r->index, _PyLong_GetOne());
1025    if (!new_index)
1026        return NULL;
1027
1028    product = PyNumber_Multiply(r->index, r->step);
1029    if (!product) {
1030        Py_DECREF(new_index);
1031        return NULL;
1032    }
1033
1034    result = PyNumber_Add(r->start, product);
1035    Py_DECREF(product);
1036    if (result) {
1037        Py_SETREF(r->index, new_index);
1038    }
1039    else {
1040        Py_DECREF(new_index);
1041    }
1042
1043    return result;
1044}
1045
1046PyTypeObject PyLongRangeIter_Type = {
1047        PyVarObject_HEAD_INIT(&PyType_Type, 0)
1048        "longrange_iterator",                   /* tp_name */
1049        sizeof(longrangeiterobject),            /* tp_basicsize */
1050        0,                                      /* tp_itemsize */
1051        /* methods */
1052        (destructor)longrangeiter_dealloc,      /* tp_dealloc */
1053        0,                                      /* tp_vectorcall_offset */
1054        0,                                      /* tp_getattr */
1055        0,                                      /* tp_setattr */
1056        0,                                      /* tp_as_async */
1057        0,                                      /* tp_repr */
1058        0,                                      /* tp_as_number */
1059        0,                                      /* tp_as_sequence */
1060        0,                                      /* tp_as_mapping */
1061        0,                                      /* tp_hash */
1062        0,                                      /* tp_call */
1063        0,                                      /* tp_str */
1064        PyObject_GenericGetAttr,                /* tp_getattro */
1065        0,                                      /* tp_setattro */
1066        0,                                      /* tp_as_buffer */
1067        Py_TPFLAGS_DEFAULT,                     /* tp_flags */
1068        0,                                      /* tp_doc */
1069        0,                                      /* tp_traverse */
1070        0,                                      /* tp_clear */
1071        0,                                      /* tp_richcompare */
1072        0,                                      /* tp_weaklistoffset */
1073        PyObject_SelfIter,                      /* tp_iter */
1074        (iternextfunc)longrangeiter_next,       /* tp_iternext */
1075        longrangeiter_methods,                  /* tp_methods */
1076        0,
1077};
1078
1079static PyObject *
1080range_iter(PyObject *seq)
1081{
1082    rangeobject *r = (rangeobject *)seq;
1083    longrangeiterobject *it;
1084    long lstart, lstop, lstep;
1085    unsigned long ulen;
1086
1087    assert(PyRange_Check(seq));
1088
1089    /* If all three fields and the length convert to long, use the int
1090     * version */
1091    lstart = PyLong_AsLong(r->start);
1092    if (lstart == -1 && PyErr_Occurred()) {
1093        PyErr_Clear();
1094        goto long_range;
1095    }
1096    lstop = PyLong_AsLong(r->stop);
1097    if (lstop == -1 && PyErr_Occurred()) {
1098        PyErr_Clear();
1099        goto long_range;
1100    }
1101    lstep = PyLong_AsLong(r->step);
1102    if (lstep == -1 && PyErr_Occurred()) {
1103        PyErr_Clear();
1104        goto long_range;
1105    }
1106    ulen = get_len_of_range(lstart, lstop, lstep);
1107    if (ulen > (unsigned long)LONG_MAX) {
1108        goto long_range;
1109    }
1110    /* check for potential overflow of lstart + ulen * lstep */
1111    if (ulen) {
1112        if (lstep > 0) {
1113            if (lstop > LONG_MAX - (lstep - 1))
1114                goto long_range;
1115        }
1116        else {
1117            if (lstop < LONG_MIN + (-1 - lstep))
1118                goto long_range;
1119        }
1120    }
1121    return fast_range_iter(lstart, lstop, lstep, (long)ulen);
1122
1123  long_range:
1124    it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
1125    if (it == NULL)
1126        return NULL;
1127
1128    it->start = r->start;
1129    it->step = r->step;
1130    it->len = r->length;
1131    it->index = _PyLong_GetZero();
1132    Py_INCREF(it->start);
1133    Py_INCREF(it->step);
1134    Py_INCREF(it->len);
1135    Py_INCREF(it->index);
1136    return (PyObject *)it;
1137}
1138
1139static PyObject *
1140range_reverse(PyObject *seq, PyObject *Py_UNUSED(ignored))
1141{
1142    rangeobject *range = (rangeobject*) seq;
1143    longrangeiterobject *it;
1144    PyObject *sum, *diff, *product;
1145    long lstart, lstop, lstep, new_start, new_stop;
1146    unsigned long ulen;
1147
1148    assert(PyRange_Check(seq));
1149
1150    /* reversed(range(start, stop, step)) can be expressed as
1151       range(start+(n-1)*step, start-step, -step), where n is the number of
1152       integers in the range.
1153
1154       If each of start, stop, step, -step, start-step, and the length
1155       of the iterator is representable as a C long, use the int
1156       version.  This excludes some cases where the reversed range is
1157       representable as a range_iterator, but it's good enough for
1158       common cases and it makes the checks simple. */
1159
1160    lstart = PyLong_AsLong(range->start);
1161    if (lstart == -1 && PyErr_Occurred()) {
1162        PyErr_Clear();
1163        goto long_range;
1164    }
1165    lstop = PyLong_AsLong(range->stop);
1166    if (lstop == -1 && PyErr_Occurred()) {
1167        PyErr_Clear();
1168        goto long_range;
1169    }
1170    lstep = PyLong_AsLong(range->step);
1171    if (lstep == -1 && PyErr_Occurred()) {
1172        PyErr_Clear();
1173        goto long_range;
1174    }
1175    /* check for possible overflow of -lstep */
1176    if (lstep == LONG_MIN)
1177        goto long_range;
1178
1179    /* check for overflow of lstart - lstep:
1180
1181       for lstep > 0, need only check whether lstart - lstep < LONG_MIN.
1182       for lstep < 0, need only check whether lstart - lstep > LONG_MAX
1183
1184       Rearrange these inequalities as:
1185
1186           lstart - LONG_MIN < lstep  (lstep > 0)
1187           LONG_MAX - lstart < -lstep  (lstep < 0)
1188
1189       and compute both sides as unsigned longs, to avoid the
1190       possibility of undefined behaviour due to signed overflow. */
1191
1192    if (lstep > 0) {
1193         if ((unsigned long)lstart - LONG_MIN < (unsigned long)lstep)
1194            goto long_range;
1195    }
1196    else {
1197        if (LONG_MAX - (unsigned long)lstart < 0UL - lstep)
1198            goto long_range;
1199    }
1200
1201    ulen = get_len_of_range(lstart, lstop, lstep);
1202    if (ulen > (unsigned long)LONG_MAX)
1203        goto long_range;
1204
1205    new_stop = lstart - lstep;
1206    new_start = (long)(new_stop + ulen * lstep);
1207    return fast_range_iter(new_start, new_stop, -lstep, (long)ulen);
1208
1209long_range:
1210    it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
1211    if (it == NULL)
1212        return NULL;
1213    it->index = it->start = it->step = NULL;
1214
1215    /* start + (len - 1) * step */
1216    it->len = range->length;
1217    Py_INCREF(it->len);
1218
1219    diff = PyNumber_Subtract(it->len, _PyLong_GetOne());
1220    if (!diff)
1221        goto create_failure;
1222
1223    product = PyNumber_Multiply(diff, range->step);
1224    Py_DECREF(diff);
1225    if (!product)
1226        goto create_failure;
1227
1228    sum = PyNumber_Add(range->start, product);
1229    Py_DECREF(product);
1230    it->start = sum;
1231    if (!it->start)
1232        goto create_failure;
1233
1234    it->step = PyNumber_Negative(range->step);
1235    if (!it->step)
1236        goto create_failure;
1237
1238    it->index = _PyLong_GetZero();
1239    Py_INCREF(it->index);
1240    return (PyObject *)it;
1241
1242create_failure:
1243    Py_DECREF(it);
1244    return NULL;
1245}
1246