Skip to content

Commit 2fa0b0d

Browse files
serhiy-storchakaserhiy-storchaka
committed
bpo-45026: More compact range iterator
1 parent 989f6a3 commit 2fa0b0d

File tree

2 files changed

+55
-40
lines changed

2 files changed

+55
-40
lines changed

‎Lib/test/test_range.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -421,6 +421,24 @@ def test_large_exhausted_iterator_pickling(self):
421421
self.assertEqual(list(i), [])
422422
self.assertEqual(list(i2), [])
423423

424+
def test_iterator_unpickle_compat(self):
425+
testcases = [
426+
b'c__builtin__\niter\n(c__builtin__\nxrange\n(I10\nI20\nI2\ntRtRI2\nb.',
427+
b'c__builtin__\niter\n(c__builtin__\nxrange\n(K\nK\x14K\x02tRtRK\x02b.',
428+
b'\x80\x02c__builtin__\niter\nc__builtin__\nxrange\nK\nK\x14K\x02\x87R\x85RK\x02b.',
429+
b'\x80\x03cbuiltins\niter\ncbuiltins\nrange\nK\nK\x14K\x02\x87R\x85RK\x02b.',
430+
b'\x80\x04\x951\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x8c\x04iter\x93\x8c\x08builtins\x8c\x05range\x93K\nK\x14K\x02\x87R\x85RK\x02b.',
431+
432+
b'c__builtin__\niter\n(c__builtin__\nxrange\n(L-36893488147419103232L\nI20\nI2\ntRtRL18446744073709551623L\nb.',
433+
b'c__builtin__\niter\n(c__builtin__\nxrange\n(L-36893488147419103232L\nK\x14K\x02tRtRL18446744073709551623L\nb.',
434+
b'\x80\x02c__builtin__\niter\nc__builtin__\nxrange\n\x8a\t\x00\x00\x00\x00\x00\x00\x00\x00\xfeK\x14K\x02\x87R\x85R\x8a\t\x07\x00\x00\x00\x00\x00\x00\x00\x01b.',
435+
b'\x80\x03cbuiltins\niter\ncbuiltins\nrange\n\x8a\t\x00\x00\x00\x00\x00\x00\x00\x00\xfeK\x14K\x02\x87R\x85R\x8a\t\x07\x00\x00\x00\x00\x00\x00\x00\x01b.',
436+
b'\x80\x04\x95C\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x8c\x04iter\x93\x8c\x08builtins\x8c\x05range\x93\x8a\t\x00\x00\x00\x00\x00\x00\x00\x00\xfeK\x14K\x02\x87R\x85R\x8a\t\x07\x00\x00\x00\x00\x00\x00\x00\x01b.',
437+
]
438+
for t in testcases:
439+
it = pickle.loads(t)
440+
self.assertEqual(list(it), [14, 16, 18])
441+
424442
def test_odd_bug(self):
425443
# This used to raise a "SystemError: NULL result without error"
426444
# because the range validation step was eating the exception

‎Objects/rangeobject.c

Lines changed: 37 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -766,7 +766,6 @@ PyTypeObject PyRange_Type = {
766766

767767
typedef struct {
768768
PyObject_HEAD
769-
long index;
770769
long start;
771770
long step;
772771
long len;
@@ -775,18 +774,19 @@ typedef struct {
775774
static PyObject *
776775
rangeiter_next(rangeiterobject *r)
777776
{
778-
if (r->index < r->len)
779-
/* cast to unsigned to avoid possible signed overflow
780-
in intermediate calculations. */
781-
return PyLong_FromLong((long)(r->start +
782-
(unsigned long)(r->index++) * r->step));
777+
if (r->len > 0) {
778+
long result = r->start;
779+
r->start += r->step;
780+
r->len--;
781+
return PyLong_FromLong(result);
782+
}
783783
return NULL;
784784
}
785785

786786
static PyObject *
787787
rangeiter_len(rangeiterobject *r, PyObject *Py_UNUSED(ignored))
788788
{
789-
return PyLong_FromLong(r->len - r->index);
789+
return PyLong_FromLong(r->len);
790790
}
791791

792792
PyDoc_STRVAR(length_hint_doc,
@@ -813,8 +813,8 @@ rangeiter_reduce(rangeiterobject *r, PyObject *Py_UNUSED(ignored))
813813
if (range == NULL)
814814
goto err;
815815
/* return the result */
816-
return Py_BuildValue("N(N)i", _PyEval_GetBuiltinId(&PyId_iter),
817-
range, r->index);
816+
return Py_BuildValue("N(N)O", _PyEval_GetBuiltinId(&PyId_iter),
817+
range, Py_None);
818818
err:
819819
Py_XDECREF(start);
820820
Py_XDECREF(stop);
@@ -833,7 +833,8 @@ rangeiter_setstate(rangeiterobject *r, PyObject *state)
833833
index = 0;
834834
else if (index > r->len)
835835
index = r->len; /* exhausted iterator */
836-
r->index = index;
836+
r->start += index * r->step;
837+
r->len -= index;
837838
Py_RETURN_NONE;
838839
}
839840

@@ -931,13 +932,11 @@ fast_range_iter(long start, long stop, long step)
931932
return NULL;
932933
}
933934
it->len = (long)ulen;
934-
it->index = 0;
935935
return (PyObject *)it;
936936
}
937937

938938
typedef struct {
939939
PyObject_HEAD
940-
PyObject *index;
941940
PyObject *start;
942941
PyObject *step;
943942
PyObject *len;
@@ -946,7 +945,8 @@ typedef struct {
946945
static PyObject *
947946
longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
948947
{
949-
return PyNumber_Subtract(r->len, r->index);
948+
Py_INCREF(r->len);
949+
return r->len;
950950
}
951951

952952
static PyObject *
@@ -976,7 +976,7 @@ longrangeiter_reduce(longrangeiterobject *r, PyObject *Py_UNUSED(ignored))
976976

977977
/* return the result */
978978
return Py_BuildValue("N(N)O", _PyEval_GetBuiltinId(&PyId_iter),
979-
range, r->index);
979+
range, Py_None);
980980
}
981981

982982
static PyObject *
@@ -999,8 +999,18 @@ longrangeiter_setstate(longrangeiterobject *r, PyObject *state)
999999
if (cmp > 0)
10001000
state = r->len;
10011001
}
1002-
Py_INCREF(state);
1003-
Py_XSETREF(r->index, state);
1002+
PyObject *new_len = PyNumber_Subtract(r->len, state);
1003+
if (new_len == NULL)
1004+
return NULL;
1005+
Py_SETREF(r->len, new_len);
1006+
PyObject *product = PyNumber_Multiply(state, r->step);
1007+
if (product == NULL)
1008+
return NULL;
1009+
PyObject *new_start = PyNumber_Add(r->start, product);
1010+
Py_DECREF(product);
1011+
if (new_start == NULL)
1012+
return NULL;
1013+
Py_SETREF(r->start, new_start);
10041014
Py_RETURN_NONE;
10051015
}
10061016

@@ -1017,7 +1027,6 @@ static PyMethodDef longrangeiter_methods[] = {
10171027
static void
10181028
longrangeiter_dealloc(longrangeiterobject *r)
10191029
{
1020-
Py_XDECREF(r->index);
10211030
Py_XDECREF(r->start);
10221031
Py_XDECREF(r->step);
10231032
Py_XDECREF(r->len);
@@ -1027,29 +1036,21 @@ longrangeiter_dealloc(longrangeiterobject *r)
10271036
static PyObject *
10281037
longrangeiter_next(longrangeiterobject *r)
10291038
{
1030-
PyObject *product, *new_index, *result;
1031-
if (PyObject_RichCompareBool(r->index, r->len, Py_LT) != 1)
1039+
if (PyObject_RichCompareBool(r->len, _PyLong_GetZero(), Py_GT) != 1)
10321040
return NULL;
10331041

1034-
new_index = PyNumber_Add(r->index, _PyLong_GetOne());
1035-
if (!new_index)
1042+
PyObject *new_start = PyNumber_Add(r->start, r->step);
1043+
if (new_start == NULL) {
10361044
return NULL;
1037-
1038-
product = PyNumber_Multiply(r->index, r->step);
1039-
if (!product) {
1040-
Py_DECREF(new_index);
1041-
return NULL;
1042-
}
1043-
1044-
result = PyNumber_Add(r->start, product);
1045-
Py_DECREF(product);
1046-
if (result) {
1047-
Py_SETREF(r->index, new_index);
10481045
}
1049-
else {
1050-
Py_DECREF(new_index);
1046+
PyObject *new_len = PyNumber_Subtract(r->len, _PyLong_GetOne());
1047+
if (new_len == NULL) {
1048+
Py_DECREF(new_start);
1049+
return NULL;
10511050
}
1052-
1051+
PyObject *result = r->start;
1052+
r->start = new_start;
1053+
Py_SETREF(r->len, new_len);
10531054
return result;
10541055
}
10551056

@@ -1128,11 +1129,9 @@ range_iter(PyObject *seq)
11281129
it->start = r->start;
11291130
it->step = r->step;
11301131
it->len = r->length;
1131-
it->index = _PyLong_GetZero();
11321132
Py_INCREF(it->start);
11331133
Py_INCREF(it->step);
11341134
Py_INCREF(it->len);
1135-
Py_INCREF(it->index);
11361135
return (PyObject *)it;
11371136
}
11381137

@@ -1210,7 +1209,7 @@ range_reverse(PyObject *seq, PyObject *Py_UNUSED(ignored))
12101209
it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
12111210
if (it == NULL)
12121211
return NULL;
1213-
it->index = it->start = it->step = NULL;
1212+
it->start = it->step = NULL;
12141213

12151214
/* start + (len - 1) * step */
12161215
it->len = range->length;
@@ -1235,8 +1234,6 @@ range_reverse(PyObject *seq, PyObject *Py_UNUSED(ignored))
12351234
if (!it->step)
12361235
goto create_failure;
12371236

1238-
it->index = _PyLong_GetZero();
1239-
Py_INCREF(it->index);
12401237
return (PyObject *)it;
12411238

12421239
create_failure:

0 commit comments

Comments
 (0)