Tagging 3.0a4.
[python/dscho.git] / Objects / rangeobject.c
blob5ea5d6fb87dc1b24ce8e03ec5688ccb7e5ee751d
1 /* Range object implementation */
3 #include "Python.h"
5 /* Support objects whose length is > PY_SSIZE_T_MAX.
7 This could be sped up for small PyLongs if they fit in an Py_ssize_t.
8 This only matters on Win64. Though we could use PY_LONG_LONG which
9 would presumably help perf.
12 typedef struct {
13 PyObject_HEAD
14 PyObject *start;
15 PyObject *stop;
16 PyObject *step;
17 } rangeobject;
19 /* Helper function for validating step. Always returns a new reference or
20 NULL on error.
22 static PyObject *
23 validate_step(PyObject *step)
25 /* No step specified, use a step of 1. */
26 if (!step)
27 return PyLong_FromLong(1);
29 step = PyNumber_Index(step);
30 if (step) {
31 Py_ssize_t istep = PyNumber_AsSsize_t(step, NULL);
32 if (istep == -1 && PyErr_Occurred()) {
33 /* Ignore OverflowError, we know the value isn't 0. */
34 PyErr_Clear();
36 else if (istep == 0) {
37 PyErr_SetString(PyExc_ValueError,
38 "range() arg 3 must not be zero");
39 Py_CLEAR(step);
43 return step;
46 /* XXX(nnorwitz): should we error check if the user passes any empty ranges?
47 range(-10)
48 range(0, -5)
49 range(0, 5, -1)
51 static PyObject *
52 range_new(PyTypeObject *type, PyObject *args, PyObject *kw)
54 rangeobject *obj = NULL;
55 PyObject *start = NULL, *stop = NULL, *step = NULL;
57 if (!_PyArg_NoKeywords("range()", kw))
58 return NULL;
60 if (PyTuple_Size(args) <= 1) {
61 if (!PyArg_UnpackTuple(args, "range", 1, 1, &stop))
62 goto Fail;
63 stop = PyNumber_Index(stop);
64 if (!stop)
65 goto Fail;
66 start = PyLong_FromLong(0);
67 step = PyLong_FromLong(1);
68 if (!start || !step)
69 goto Fail;
71 else {
72 if (!PyArg_UnpackTuple(args, "range", 2, 3,
73 &start, &stop, &step))
74 goto Fail;
76 /* Convert borrowed refs to owned refs */
77 start = PyNumber_Index(start);
78 stop = PyNumber_Index(stop);
79 step = validate_step(step);
80 if (!start || !stop || !step)
81 goto Fail;
84 obj = PyObject_New(rangeobject, &PyRange_Type);
85 if (obj == NULL)
86 goto Fail;
87 obj->start = start;
88 obj->stop = stop;
89 obj->step = step;
90 return (PyObject *) obj;
92 Fail:
93 Py_XDECREF(start);
94 Py_XDECREF(stop);
95 Py_XDECREF(step);
96 return NULL;
99 PyDoc_STRVAR(range_doc,
100 "range([start,] stop[, step]) -> range object\n\
102 Returns an iterator that generates the numbers in the range on demand.");
104 static void
105 range_dealloc(rangeobject *r)
107 Py_DECREF(r->start);
108 Py_DECREF(r->stop);
109 Py_DECREF(r->step);
110 PyObject_Del(r);
113 /* Return number of items in range (lo, hi, step), when arguments are
114 * PyInt or PyLong objects. step > 0 required. Return a value < 0 if
115 * & only if the true value is too large to fit in a signed long.
116 * Arguments MUST return 1 with either PyLong_Check() or
117 * PyLong_Check(). Return -1 when there is an error.
119 static PyObject*
120 range_length_obj(rangeobject *r)
122 /* -------------------------------------------------------------
123 Algorithm is equal to that of get_len_of_range(), but it operates
124 on PyObjects (which are assumed to be PyLong or PyInt objects).
125 ---------------------------------------------------------------*/
126 int cmp_result, cmp_call;
127 PyObject *lo, *hi;
128 PyObject *step = NULL;
129 PyObject *diff = NULL;
130 PyObject *one = NULL;
131 PyObject *tmp1 = NULL, *tmp2 = NULL, *result;
132 /* holds sub-expression evaluations */
134 PyObject *zero = PyLong_FromLong(0);
135 if (zero == NULL)
136 return NULL;
137 cmp_call = PyObject_Cmp(r->step, zero, &cmp_result);
138 Py_DECREF(zero);
139 if (cmp_call == -1)
140 return NULL;
142 assert(cmp_result != 0);
143 if (cmp_result > 0) {
144 lo = r->start;
145 hi = r->stop;
146 step = r->step;
147 Py_INCREF(step);
148 } else {
149 lo = r->stop;
150 hi = r->start;
151 step = PyNumber_Negative(r->step);
152 if (!step)
153 return NULL;
156 /* if (lo >= hi), return length of 0. */
157 if (PyObject_Compare(lo, hi) >= 0) {
158 Py_XDECREF(step);
159 return PyLong_FromLong(0);
162 if ((one = PyLong_FromLong(1L)) == NULL)
163 goto Fail;
165 if ((tmp1 = PyNumber_Subtract(hi, lo)) == NULL)
166 goto Fail;
168 if ((diff = PyNumber_Subtract(tmp1, one)) == NULL)
169 goto Fail;
171 if ((tmp2 = PyNumber_FloorDivide(diff, step)) == NULL)
172 goto Fail;
174 if ((result = PyNumber_Add(tmp2, one)) == NULL)
175 goto Fail;
177 Py_DECREF(tmp2);
178 Py_DECREF(diff);
179 Py_DECREF(step);
180 Py_DECREF(tmp1);
181 Py_DECREF(one);
182 return result;
184 Fail:
185 Py_XDECREF(tmp2);
186 Py_XDECREF(diff);
187 Py_XDECREF(step);
188 Py_XDECREF(tmp1);
189 Py_XDECREF(one);
190 return NULL;
193 static Py_ssize_t
194 range_length(rangeobject *r)
196 PyObject *len = range_length_obj(r);
197 Py_ssize_t result = -1;
198 if (len) {
199 result = PyLong_AsSsize_t(len);
200 Py_DECREF(len);
202 return result;
205 /* range(...)[x] is necessary for: seq[:] = range(...) */
207 static PyObject *
208 range_item(rangeobject *r, Py_ssize_t i)
210 Py_ssize_t len = range_length(r);
211 PyObject *rem, *incr, *result;
213 /* XXX(nnorwitz): should negative indices be supported? */
214 /* XXX(nnorwitz): should support range[x] where x > PY_SSIZE_T_MAX? */
215 if (i < 0 || i >= len) {
216 if (!PyErr_Occurred())
217 PyErr_SetString(PyExc_IndexError,
218 "range object index out of range");
219 return NULL;
222 /* XXX(nnorwitz): optimize for short ints. */
223 rem = PyLong_FromSsize_t(i);
224 if (!rem)
225 return NULL;
226 incr = PyNumber_Multiply(rem, r->step);
227 Py_DECREF(rem);
228 if (!incr)
229 return NULL;
230 result = PyNumber_Add(r->start, incr);
231 Py_DECREF(incr);
232 return result;
235 static PyObject *
236 range_repr(rangeobject *r)
238 Py_ssize_t istep;
240 /* Check for special case values for printing. We don't always
241 need the step value. We don't care about errors
242 (it means overflow), so clear the errors. */
243 istep = PyNumber_AsSsize_t(r->step, NULL);
244 if (istep != 1 || (istep == -1 && PyErr_Occurred())) {
245 PyErr_Clear();
248 if (istep == 1)
249 return PyUnicode_FromFormat("range(%R, %R)", r->start, r->stop);
250 else
251 return PyUnicode_FromFormat("range(%R, %R, %R)",
252 r->start, r->stop, r->step);
255 static PySequenceMethods range_as_sequence = {
256 (lenfunc)range_length, /* sq_length */
257 0, /* sq_concat */
258 0, /* sq_repeat */
259 (ssizeargfunc)range_item, /* sq_item */
260 0, /* sq_slice */
263 static PyObject * range_iter(PyObject *seq);
264 static PyObject * range_reverse(PyObject *seq);
266 PyDoc_STRVAR(reverse_doc,
267 "Returns a reverse iterator.");
269 static PyMethodDef range_methods[] = {
270 {"__reversed__", (PyCFunction)range_reverse, METH_NOARGS,
271 reverse_doc},
272 {NULL, NULL} /* sentinel */
275 PyTypeObject PyRange_Type = {
276 PyVarObject_HEAD_INIT(&PyType_Type, 0)
277 "range", /* Name of this type */
278 sizeof(rangeobject), /* Basic object size */
279 0, /* Item size for varobject */
280 (destructor)range_dealloc, /* tp_dealloc */
281 0, /* tp_print */
282 0, /* tp_getattr */
283 0, /* tp_setattr */
284 0, /* tp_compare */
285 (reprfunc)range_repr, /* tp_repr */
286 0, /* tp_as_number */
287 &range_as_sequence, /* tp_as_sequence */
288 0, /* tp_as_mapping */
289 0, /* tp_hash */
290 0, /* tp_call */
291 0, /* tp_str */
292 PyObject_GenericGetAttr, /* tp_getattro */
293 0, /* tp_setattro */
294 0, /* tp_as_buffer */
295 Py_TPFLAGS_DEFAULT, /* tp_flags */
296 range_doc, /* tp_doc */
297 0, /* tp_traverse */
298 0, /* tp_clear */
299 0, /* tp_richcompare */
300 0, /* tp_weaklistoffset */
301 range_iter, /* tp_iter */
302 0, /* tp_iternext */
303 range_methods, /* tp_methods */
304 0, /* tp_members */
305 0, /* tp_getset */
306 0, /* tp_base */
307 0, /* tp_dict */
308 0, /* tp_descr_get */
309 0, /* tp_descr_set */
310 0, /* tp_dictoffset */
311 0, /* tp_init */
312 0, /* tp_alloc */
313 range_new, /* tp_new */
316 /*********************** range Iterator **************************/
318 /* There are 2 types of iterators, one for C longs, the other for
319 Python longs (ie, PyObjects). This should make iteration fast
320 in the normal case, but possible for any numeric value.
323 typedef struct {
324 PyObject_HEAD
325 long index;
326 long start;
327 long step;
328 long len;
329 } rangeiterobject;
331 static PyObject *
332 rangeiter_next(rangeiterobject *r)
334 if (r->index < r->len)
335 return PyLong_FromLong(r->start + (r->index++) * r->step);
336 return NULL;
339 static PyObject *
340 rangeiter_len(rangeiterobject *r)
342 return PyLong_FromLong(r->len - r->index);
345 typedef struct {
346 PyObject_HEAD
347 PyObject *index;
348 PyObject *start;
349 PyObject *step;
350 PyObject *len;
351 } longrangeiterobject;
353 static PyObject *
354 longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
356 return PyNumber_Subtract(r->len, r->index);
359 static PyObject *rangeiter_new(PyTypeObject *, PyObject *args, PyObject *kw);
361 PyDoc_STRVAR(length_hint_doc,
362 "Private method returning an estimate of len(list(it)).");
364 static PyMethodDef rangeiter_methods[] = {
365 {"__length_hint__", (PyCFunction)rangeiter_len, METH_NOARGS,
366 length_hint_doc},
367 {NULL, NULL} /* sentinel */
370 PyTypeObject PyRangeIter_Type = {
371 PyVarObject_HEAD_INIT(&PyType_Type, 0)
372 "range_iterator", /* tp_name */
373 sizeof(rangeiterobject), /* tp_basicsize */
374 0, /* tp_itemsize */
375 /* methods */
376 (destructor)PyObject_Del, /* tp_dealloc */
377 0, /* tp_print */
378 0, /* tp_getattr */
379 0, /* tp_setattr */
380 0, /* tp_compare */
381 0, /* tp_repr */
382 0, /* tp_as_number */
383 0, /* tp_as_sequence */
384 0, /* tp_as_mapping */
385 0, /* tp_hash */
386 0, /* tp_call */
387 0, /* tp_str */
388 PyObject_GenericGetAttr, /* tp_getattro */
389 0, /* tp_setattro */
390 0, /* tp_as_buffer */
391 Py_TPFLAGS_DEFAULT, /* tp_flags */
392 0, /* tp_doc */
393 0, /* tp_traverse */
394 0, /* tp_clear */
395 0, /* tp_richcompare */
396 0, /* tp_weaklistoffset */
397 PyObject_SelfIter, /* tp_iter */
398 (iternextfunc)rangeiter_next, /* tp_iternext */
399 rangeiter_methods, /* tp_methods */
400 0, /* tp_members */
401 0, /* tp_getset */
402 0, /* tp_base */
403 0, /* tp_dict */
404 0, /* tp_descr_get */
405 0, /* tp_descr_set */
406 0, /* tp_dictoffset */
407 0, /* tp_init */
408 0, /* tp_alloc */
409 rangeiter_new, /* tp_new */
412 /* Return number of items in range/xrange (lo, hi, step). step > 0
413 * required. Return a value < 0 if & only if the true value is too
414 * large to fit in a signed long.
416 static long
417 get_len_of_range(long lo, long hi, long step)
419 /* -------------------------------------------------------------
420 If lo >= hi, the range is empty.
421 Else if n values are in the range, the last one is
422 lo + (n-1)*step, which must be <= hi-1. Rearranging,
423 n <= (hi - lo - 1)/step + 1, so taking the floor of the RHS gives
424 the proper value. Since lo < hi in this case, hi-lo-1 >= 0, so
425 the RHS is non-negative and so truncation is the same as the
426 floor. Letting M be the largest positive long, the worst case
427 for the RHS numerator is hi=M, lo=-M-1, and then
428 hi-lo-1 = M-(-M-1)-1 = 2*M. Therefore unsigned long has enough
429 precision to compute the RHS exactly.
430 ---------------------------------------------------------------*/
431 long n = 0;
432 if (lo < hi) {
433 unsigned long uhi = (unsigned long)hi;
434 unsigned long ulo = (unsigned long)lo;
435 unsigned long diff = uhi - ulo - 1;
436 n = (long)(diff / (unsigned long)step + 1);
438 return n;
441 static PyObject *
442 int_range_iter(long start, long stop, long step)
444 rangeiterobject *it = PyObject_New(rangeiterobject, &PyRangeIter_Type);
445 if (it == NULL)
446 return NULL;
447 it->start = start;
448 it->step = step;
449 if (step > 0)
450 it->len = get_len_of_range(start, stop, step);
451 else
452 it->len = get_len_of_range(stop, start, -step);
453 it->index = 0;
454 return (PyObject *)it;
457 static PyObject *
458 rangeiter_new(PyTypeObject *type, PyObject *args, PyObject *kw)
460 long start, stop, step;
462 if (!_PyArg_NoKeywords("rangeiter()", kw))
463 return NULL;
465 if (!PyArg_ParseTuple(args, "lll;rangeiter() requires 3 int arguments",
466 &start, &stop, &step))
467 return NULL;
469 return int_range_iter(start, stop, step);
472 static PyMethodDef longrangeiter_methods[] = {
473 {"__length_hint__", (PyCFunction)longrangeiter_len, METH_NOARGS,
474 length_hint_doc},
475 {NULL, NULL} /* sentinel */
478 static void
479 longrangeiter_dealloc(longrangeiterobject *r)
481 Py_XDECREF(r->index);
482 Py_XDECREF(r->start);
483 Py_XDECREF(r->step);
484 Py_XDECREF(r->len);
485 PyObject_Del(r);
488 static PyObject *
489 longrangeiter_next(longrangeiterobject *r)
491 PyObject *one, *product, *new_index, *result;
492 if (PyObject_RichCompareBool(r->index, r->len, Py_LT) != 1)
493 return NULL;
495 one = PyLong_FromLong(1);
496 if (!one)
497 return NULL;
499 product = PyNumber_Multiply(r->index, r->step);
500 if (!product) {
501 Py_DECREF(one);
502 return NULL;
505 new_index = PyNumber_Add(r->index, one);
506 Py_DECREF(one);
507 if (!new_index) {
508 Py_DECREF(product);
509 return NULL;
512 result = PyNumber_Add(r->start, product);
513 Py_DECREF(product);
514 if (result) {
515 Py_DECREF(r->index);
516 r->index = new_index;
519 return result;
522 PyTypeObject PyLongRangeIter_Type = {
523 PyVarObject_HEAD_INIT(&PyType_Type, 0)
524 "longrange_iterator", /* tp_name */
525 sizeof(longrangeiterobject), /* tp_basicsize */
526 0, /* tp_itemsize */
527 /* methods */
528 (destructor)longrangeiter_dealloc, /* tp_dealloc */
529 0, /* tp_print */
530 0, /* tp_getattr */
531 0, /* tp_setattr */
532 0, /* tp_compare */
533 0, /* tp_repr */
534 0, /* tp_as_number */
535 0, /* tp_as_sequence */
536 0, /* tp_as_mapping */
537 0, /* tp_hash */
538 0, /* tp_call */
539 0, /* tp_str */
540 PyObject_GenericGetAttr, /* tp_getattro */
541 0, /* tp_setattro */
542 0, /* tp_as_buffer */
543 Py_TPFLAGS_DEFAULT, /* tp_flags */
544 0, /* tp_doc */
545 0, /* tp_traverse */
546 0, /* tp_clear */
547 0, /* tp_richcompare */
548 0, /* tp_weaklistoffset */
549 PyObject_SelfIter, /* tp_iter */
550 (iternextfunc)longrangeiter_next, /* tp_iternext */
551 longrangeiter_methods, /* tp_methods */
555 static PyObject *
556 range_iter(PyObject *seq)
558 rangeobject *r = (rangeobject *)seq;
559 longrangeiterobject *it;
560 PyObject *tmp, *len;
561 long lstart, lstop, lstep;
563 assert(PyRange_Check(seq));
565 /* If all three fields convert to long, use the int version */
566 lstart = PyLong_AsLong(r->start);
567 if (lstart != -1 || !PyErr_Occurred()) {
568 lstop = PyLong_AsLong(r->stop);
569 if (lstop != -1 || !PyErr_Occurred()) {
570 lstep = PyLong_AsLong(r->step);
571 if (lstep != -1 || !PyErr_Occurred())
572 return int_range_iter(lstart, lstop, lstep);
575 /* Some conversion failed, so there is an error set. Clear it,
576 and try again with a long range. */
577 PyErr_Clear();
579 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
580 if (it == NULL)
581 return NULL;
583 /* Do all initialization here, so we can DECREF on failure. */
584 it->start = r->start;
585 it->step = r->step;
586 Py_INCREF(it->start);
587 Py_INCREF(it->step);
589 it->len = it->index = NULL;
591 /* Calculate length: (r->stop - r->start) / r->step */
592 tmp = PyNumber_Subtract(r->stop, r->start);
593 if (!tmp)
594 goto create_failure;
595 len = PyNumber_FloorDivide(tmp, r->step);
596 Py_DECREF(tmp);
597 if (!len)
598 goto create_failure;
599 it->len = len;
600 it->index = PyLong_FromLong(0);
601 if (!it->index)
602 goto create_failure;
604 return (PyObject *)it;
606 create_failure:
607 Py_DECREF(it);
608 return NULL;
611 static PyObject *
612 range_reverse(PyObject *seq)
614 rangeobject *range = (rangeobject*) seq;
615 longrangeiterobject *it;
616 PyObject *one, *sum, *diff, *len = NULL, *product;
617 long lstart, lstop, lstep;
619 /* XXX(nnorwitz): do the calc for the new start/stop first,
620 then if they fit, call the proper iter()?
622 assert(PyRange_Check(seq));
624 /* If all three fields convert to long, use the int version */
625 lstart = PyLong_AsLong(range->start);
626 if (lstart != -1 || !PyErr_Occurred()) {
627 lstop = PyLong_AsLong(range->stop);
628 if (lstop != -1 || !PyErr_Occurred()) {
629 lstep = PyLong_AsLong(range->step);
630 if (lstep != -1 || !PyErr_Occurred()) {
631 /* XXX(nnorwitz): need to check for overflow and simplify. */
632 long len = get_len_of_range(lstart, lstop, lstep);
633 long new_start = lstart + (len - 1) * lstep;
634 long new_stop = lstart;
635 if (lstep > 0)
636 new_stop -= 1;
637 else
638 new_stop += 1;
639 return int_range_iter(new_start, new_stop, -lstep);
643 PyErr_Clear();
645 it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
646 if (it == NULL)
647 return NULL;
649 /* start + (len - 1) * step */
650 len = range_length_obj(range);
651 if (!len)
652 goto create_failure;
654 one = PyLong_FromLong(1);
655 if (!one)
656 goto create_failure;
658 diff = PyNumber_Subtract(len, one);
659 Py_DECREF(one);
660 if (!diff)
661 goto create_failure;
663 product = PyNumber_Multiply(len, range->step);
664 if (!product)
665 goto create_failure;
667 sum = PyNumber_Add(range->start, product);
668 Py_DECREF(product);
669 it->start = sum;
670 if (!it->start)
671 goto create_failure;
672 it->step = PyNumber_Negative(range->step);
673 if (!it->step) {
674 Py_DECREF(it->start);
675 PyObject_Del(it);
676 return NULL;
679 /* Steal reference to len. */
680 it->len = len;
682 it->index = PyLong_FromLong(0);
683 if (!it->index) {
684 Py_DECREF(it);
685 return NULL;
688 return (PyObject *)it;
690 create_failure:
691 Py_XDECREF(len);
692 PyObject_Del(it);
693 return NULL;