blob: 5203f409ac59073d796455d1f43a4ced5be49278 [file] [log] [blame]
Andrew Hsieh9a7616f2013-05-21 20:32:42 +08001/* Range object implementation */
2
3#include "Python.h"
4
5typedef struct {
6 PyObject_HEAD
7 long start;
8 long step;
9 long len;
10} rangeobject;
11
12/* Return number of items in range (lo, hi, step). step != 0
13 * required. The result always fits in an unsigned long.
14 */
15static unsigned long
16get_len_of_range(long lo, long hi, long step)
17{
18 /* -------------------------------------------------------------
19 If step > 0 and lo >= hi, or step < 0 and lo <= hi, the range is empty.
20 Else for step > 0, if n values are in the range, the last one is
21 lo + (n-1)*step, which must be <= hi-1. Rearranging,
22 n <= (hi - lo - 1)/step + 1, so taking the floor of the RHS gives
23 the proper value. Since lo < hi in this case, hi-lo-1 >= 0, so
24 the RHS is non-negative and so truncation is the same as the
25 floor. Letting M be the largest positive long, the worst case
26 for the RHS numerator is hi=M, lo=-M-1, and then
27 hi-lo-1 = M-(-M-1)-1 = 2*M. Therefore unsigned long has enough
28 precision to compute the RHS exactly. The analysis for step < 0
29 is similar.
30 ---------------------------------------------------------------*/
31 assert(step != 0);
32 if (step > 0 && lo < hi)
33 return 1UL + (hi - 1UL - lo) / step;
34 else if (step < 0 && lo > hi)
35 return 1UL + (lo - 1UL - hi) / (0UL - step);
36 else
37 return 0UL;
38}
39
40/* Return a stop value suitable for reconstructing the xrange from
41 * a (start, stop, step) triple. Used in range_repr and range_reduce.
42 * Computes start + len * step, clipped to the range [LONG_MIN, LONG_MAX].
43 */
44static long
45get_stop_for_range(rangeobject *r)
46{
47 long last;
48
49 if (r->len == 0)
50 return r->start;
51
52 /* The tricky bit is avoiding overflow. We first compute the last entry in
53 the xrange, start + (len - 1) * step, which is guaranteed to lie within
54 the range of a long, and then add step to it. See the range_reverse
55 comments for an explanation of the casts below.
56 */
57 last = (long)(r->start + (unsigned long)(r->len - 1) * r->step);
58 if (r->step > 0)
59 return last > LONG_MAX - r->step ? LONG_MAX : last + r->step;
60 else
61 return last < LONG_MIN - r->step ? LONG_MIN : last + r->step;
62}
63
64static PyObject *
65range_new(PyTypeObject *type, PyObject *args, PyObject *kw)
66{
67 rangeobject *obj;
68 long ilow = 0, ihigh = 0, istep = 1;
69 unsigned long n;
70
71 if (!_PyArg_NoKeywords("xrange()", kw))
72 return NULL;
73
74 if (PyTuple_Size(args) <= 1) {
75 if (!PyArg_ParseTuple(args,
76 "l;xrange() requires 1-3 int arguments",
77 &ihigh))
78 return NULL;
79 }
80 else {
81 if (!PyArg_ParseTuple(args,
82 "ll|l;xrange() requires 1-3 int arguments",
83 &ilow, &ihigh, &istep))
84 return NULL;
85 }
86 if (istep == 0) {
87 PyErr_SetString(PyExc_ValueError, "xrange() arg 3 must not be zero");
88 return NULL;
89 }
90 n = get_len_of_range(ilow, ihigh, istep);
91 if (n > (unsigned long)LONG_MAX || (long)n > PY_SSIZE_T_MAX) {
92 PyErr_SetString(PyExc_OverflowError,
93 "xrange() result has too many items");
94 return NULL;
95 }
96
97 obj = PyObject_New(rangeobject, &PyRange_Type);
98 if (obj == NULL)
99 return NULL;
100 obj->start = ilow;
101 obj->len = (long)n;
102 obj->step = istep;
103 return (PyObject *) obj;
104}
105
106PyDoc_STRVAR(range_doc,
107"xrange(stop) -> xrange object\n\
108xrange(start, stop[, step]) -> xrange object\n\
109\n\
110Like range(), but instead of returning a list, returns an object that\n\
111generates the numbers in the range on demand. For looping, this is \n\
112slightly faster than range() and more memory efficient.");
113
114static PyObject *
115range_item(rangeobject *r, Py_ssize_t i)
116{
117 if (i < 0 || i >= r->len) {
118 PyErr_SetString(PyExc_IndexError,
119 "xrange object index out of range");
120 return NULL;
121 }
122 /* do calculation entirely using unsigned longs, to avoid
123 undefined behaviour due to signed overflow. */
124 return PyInt_FromLong((long)(r->start + (unsigned long)i * r->step));
125}
126
127static Py_ssize_t
128range_length(rangeobject *r)
129{
130 return (Py_ssize_t)(r->len);
131}
132
133static PyObject *
134range_repr(rangeobject *r)
135{
136 PyObject *rtn;
137
138 if (r->start == 0 && r->step == 1)
139 rtn = PyString_FromFormat("xrange(%ld)",
140 get_stop_for_range(r));
141
142 else if (r->step == 1)
143 rtn = PyString_FromFormat("xrange(%ld, %ld)",
144 r->start,
145 get_stop_for_range(r));
146
147 else
148 rtn = PyString_FromFormat("xrange(%ld, %ld, %ld)",
149 r->start,
150 get_stop_for_range(r),
151 r->step);
152 return rtn;
153}
154
155/* Pickling support */
156static PyObject *
157range_reduce(rangeobject *r, PyObject *args)
158{
159 return Py_BuildValue("(O(lll))", Py_TYPE(r),
160 r->start,
161 get_stop_for_range(r),
162 r->step);
163}
164
165static PySequenceMethods range_as_sequence = {
166 (lenfunc)range_length, /* sq_length */
167 0, /* sq_concat */
168 0, /* sq_repeat */
169 (ssizeargfunc)range_item, /* sq_item */
170 0, /* sq_slice */
171};
172
173static PyObject * range_iter(PyObject *seq);
174static PyObject * range_reverse(PyObject *seq);
175
176PyDoc_STRVAR(reverse_doc,
177"Returns a reverse iterator.");
178
179static PyMethodDef range_methods[] = {
180 {"__reversed__", (PyCFunction)range_reverse, METH_NOARGS, reverse_doc},
181 {"__reduce__", (PyCFunction)range_reduce, METH_VARARGS},
182 {NULL, NULL} /* sentinel */
183};
184
185PyTypeObject PyRange_Type = {
186 PyObject_HEAD_INIT(&PyType_Type)
187 0, /* Number of items for varobject */
188 "xrange", /* Name of this type */
189 sizeof(rangeobject), /* Basic object size */
190 0, /* Item size for varobject */
191 (destructor)PyObject_Del, /* tp_dealloc */
192 0, /* tp_print */
193 0, /* tp_getattr */
194 0, /* tp_setattr */
195 0, /* tp_compare */
196 (reprfunc)range_repr, /* tp_repr */
197 0, /* tp_as_number */
198 &range_as_sequence, /* tp_as_sequence */
199 0, /* tp_as_mapping */
200 0, /* tp_hash */
201 0, /* tp_call */
202 0, /* tp_str */
203 PyObject_GenericGetAttr, /* tp_getattro */
204 0, /* tp_setattro */
205 0, /* tp_as_buffer */
206 Py_TPFLAGS_DEFAULT, /* tp_flags */
207 range_doc, /* tp_doc */
208 0, /* tp_traverse */
209 0, /* tp_clear */
210 0, /* tp_richcompare */
211 0, /* tp_weaklistoffset */
212 range_iter, /* tp_iter */
213 0, /* tp_iternext */
214 range_methods, /* tp_methods */
215 0, /* tp_members */
216 0, /* tp_getset */
217 0, /* tp_base */
218 0, /* tp_dict */
219 0, /* tp_descr_get */
220 0, /* tp_descr_set */
221 0, /* tp_dictoffset */
222 0, /* tp_init */
223 0, /* tp_alloc */
224 range_new, /* tp_new */
225};
226
227/*********************** Xrange Iterator **************************/
228
229typedef struct {
230 PyObject_HEAD
231 long index;
232 long start;
233 long step;
234 long len;
235} rangeiterobject;
236
237static PyObject *
238rangeiter_next(rangeiterobject *r)
239{
240 if (r->index < r->len)
241 return PyInt_FromLong(r->start + (r->index++) * r->step);
242 return NULL;
243}
244
245static PyObject *
246rangeiter_len(rangeiterobject *r)
247{
248 return PyInt_FromLong(r->len - r->index);
249}
250
251PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
252
253static PyMethodDef rangeiter_methods[] = {
254 {"__length_hint__", (PyCFunction)rangeiter_len, METH_NOARGS, length_hint_doc},
255 {NULL, NULL} /* sentinel */
256};
257
258static PyTypeObject Pyrangeiter_Type = {
259 PyObject_HEAD_INIT(&PyType_Type)
260 0, /* ob_size */
261 "rangeiterator", /* tp_name */
262 sizeof(rangeiterobject), /* tp_basicsize */
263 0, /* tp_itemsize */
264 /* methods */
265 (destructor)PyObject_Del, /* tp_dealloc */
266 0, /* tp_print */
267 0, /* tp_getattr */
268 0, /* tp_setattr */
269 0, /* tp_compare */
270 0, /* tp_repr */
271 0, /* tp_as_number */
272 0, /* tp_as_sequence */
273 0, /* tp_as_mapping */
274 0, /* tp_hash */
275 0, /* tp_call */
276 0, /* tp_str */
277 PyObject_GenericGetAttr, /* tp_getattro */
278 0, /* tp_setattro */
279 0, /* tp_as_buffer */
280 Py_TPFLAGS_DEFAULT, /* tp_flags */
281 0, /* tp_doc */
282 0, /* tp_traverse */
283 0, /* tp_clear */
284 0, /* tp_richcompare */
285 0, /* tp_weaklistoffset */
286 PyObject_SelfIter, /* tp_iter */
287 (iternextfunc)rangeiter_next, /* tp_iternext */
288 rangeiter_methods, /* tp_methods */
289 0,
290};
291
292static PyObject *
293range_iter(PyObject *seq)
294{
295 rangeiterobject *it;
296
297 if (!PyRange_Check(seq)) {
298 PyErr_BadInternalCall();
299 return NULL;
300 }
301 it = PyObject_New(rangeiterobject, &Pyrangeiter_Type);
302 if (it == NULL)
303 return NULL;
304 it->index = 0;
305 it->start = ((rangeobject *)seq)->start;
306 it->step = ((rangeobject *)seq)->step;
307 it->len = ((rangeobject *)seq)->len;
308 return (PyObject *)it;
309}
310
311static PyObject *
312range_reverse(PyObject *seq)
313{
314 rangeiterobject *it;
315 long start, step, len;
316
317 if (!PyRange_Check(seq)) {
318 PyErr_BadInternalCall();
319 return NULL;
320 }
321 it = PyObject_New(rangeiterobject, &Pyrangeiter_Type);
322 if (it == NULL)
323 return NULL;
324
325 start = ((rangeobject *)seq)->start;
326 step = ((rangeobject *)seq)->step;
327 len = ((rangeobject *)seq)->len;
328
329 it->index = 0;
330 it->len = len;
331 /* the casts below guard against signed overflow by turning it
332 into unsigned overflow instead. The correctness of this
333 code still depends on conversion from unsigned long to long
334 wrapping modulo ULONG_MAX+1, which isn't guaranteed (see
335 C99 6.3.1.3p3) but seems to hold in practice for all
336 platforms we're likely to meet.
337
338 If step == LONG_MIN then we still end up with LONG_MIN
339 after negation; but this works out, since we've still got
340 the correct value modulo ULONG_MAX+1, and the range_item
341 calculation is also done modulo ULONG_MAX+1.
342 */
343 it->start = (long)(start + (unsigned long)(len-1) * step);
344 it->step = (long)(0UL-step);
345
346 return (PyObject *)it;
347}