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