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