Skip to content

Commit 7877642

Browse files
serhiy-storchakaserhiy-storchaka
authored
gh-89189: More compact range iterator (GH-27986)
1 parent 9628136 commit 7877642

File tree

7 files changed

+88
-50
lines changed

7 files changed

+88
-50
lines changed

‎Include/internal/pycore_range.h

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ extern "C" {
1010

1111
typedef struct {
1212
PyObject_HEAD
13-
long index;
1413
long start;
1514
long step;
1615
long len;

‎Lib/test/test_range.py

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -407,11 +407,7 @@ def test_iterator_pickling_overflowing_index(self):
407407
for proto in range(pickle.HIGHEST_PROTOCOL + 1):
408408
with self.subTest(proto=proto):
409409
it = iter(range(2**32 + 2))
410-
_, _, idx = it.__reduce__()
411-
self.assertEqual(idx, 0)
412-
it.__setstate__(2**32 + 1) # undocumented way to set r->index
413-
_, _, idx = it.__reduce__()
414-
self.assertEqual(idx, 2**32 + 1)
410+
it.__setstate__(2**32 + 1) # undocumented way to advance an iterator
415411
d = pickle.dumps(it, proto)
416412
it = pickle.loads(d)
417413
self.assertEqual(next(it), 2**32 + 1)
@@ -442,6 +438,38 @@ def test_large_exhausted_iterator_pickling(self):
442438
self.assertEqual(list(i), [])
443439
self.assertEqual(list(i2), [])
444440

441+
def test_iterator_unpickle_compat(self):
442+
testcases = [
443+
b'c__builtin__\niter\n(c__builtin__\nxrange\n(I10\nI20\nI2\ntRtRI2\nb.',
444+
b'c__builtin__\niter\n(c__builtin__\nxrange\n(K\nK\x14K\x02tRtRK\x02b.',
445+
b'\x80\x02c__builtin__\niter\nc__builtin__\nxrange\nK\nK\x14K\x02\x87R\x85RK\x02b.',
446+
b'\x80\x03cbuiltins\niter\ncbuiltins\nrange\nK\nK\x14K\x02\x87R\x85RK\x02b.',
447+
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.',
448+
449+
b'c__builtin__\niter\n(c__builtin__\nxrange\n(L-36893488147419103232L\nI20\nI2\ntRtRL18446744073709551623L\nb.',
450+
b'c__builtin__\niter\n(c__builtin__\nxrange\n(L-36893488147419103232L\nK\x14K\x02tRtRL18446744073709551623L\nb.',
451+
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.',
452+
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.',
453+
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.',
454+
]
455+
for t in testcases:
456+
it = pickle.loads(t)
457+
self.assertEqual(list(it), [14, 16, 18])
458+
459+
def test_iterator_setstate(self):
460+
it = iter(range(10, 20, 2))
461+
it.__setstate__(2)
462+
self.assertEqual(list(it), [14, 16, 18])
463+
it = reversed(range(10, 20, 2))
464+
it.__setstate__(3)
465+
self.assertEqual(list(it), [12, 10])
466+
it = iter(range(-2**65, 20, 2))
467+
it.__setstate__(2**64 + 7)
468+
self.assertEqual(list(it), [14, 16, 18])
469+
it = reversed(range(10, 2**65, 2))
470+
it.__setstate__(2**64 - 7)
471+
self.assertEqual(list(it), [12, 10])
472+
445473
def test_odd_bug(self):
446474
# This used to raise a "SystemError: NULL result without error"
447475
# because the range validation step was eating the exception

‎Lib/test/test_sys.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1484,7 +1484,8 @@ def delx(self): del self.__x
14841484
# PyCapsule
14851485
# XXX
14861486
# rangeiterator
1487-
check(iter(range(1)), size('4l'))
1487+
check(iter(range(1)), size('3l'))
1488+
check(iter(range(2**65)), size('3P'))
14881489
# reverse
14891490
check(reversed(''), size('nP'))
14901491
# range

‎Misc/NEWS.d/next/Core and Builtins/2021-08-29-15-55-19.bpo-45026.z7nTA3.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Optimize the :class:`range` object iterator. It is now smaller, faster
2+
iteration of ranges containing large numbers. Smaller pickles, faster
3+
unpickling.

‎Objects/rangeobject.c

Lines changed: 42 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -756,18 +756,19 @@ PyTypeObject PyRange_Type = {
756756
static PyObject *
757757
rangeiter_next(_PyRangeIterObject *r)
758758
{
759-
if (r->index < r->len)
760-
/* cast to unsigned to avoid possible signed overflow
761-
in intermediate calculations. */
762-
return PyLong_FromLong((long)(r->start +
763-
(unsigned long)(r->index++) * r->step));
759+
if (r->len > 0) {
760+
long result = r->start;
761+
r->start = result + r->step;
762+
r->len--;
763+
return PyLong_FromLong(result);
764+
}
764765
return NULL;
765766
}
766767

767768
static PyObject *
768769
rangeiter_len(_PyRangeIterObject *r, PyObject *Py_UNUSED(ignored))
769770
{
770-
return PyLong_FromLong(r->len - r->index);
771+
return PyLong_FromLong(r->len);
771772
}
772773

773774
PyDoc_STRVAR(length_hint_doc,
@@ -794,8 +795,8 @@ rangeiter_reduce(_PyRangeIterObject *r, PyObject *Py_UNUSED(ignored))
794795
if (range == NULL)
795796
goto err;
796797
/* return the result */
797-
return Py_BuildValue(
798-
"N(N)l", _PyEval_GetBuiltin(&_Py_ID(iter)), range, r->index);
798+
return Py_BuildValue("N(N)O", _PyEval_GetBuiltin(&_Py_ID(iter)),
799+
range, Py_None);
799800
err:
800801
Py_XDECREF(start);
801802
Py_XDECREF(stop);
@@ -814,7 +815,8 @@ rangeiter_setstate(_PyRangeIterObject *r, PyObject *state)
814815
index = 0;
815816
else if (index > r->len)
816817
index = r->len; /* exhausted iterator */
817-
r->index = index;
818+
r->start += index * r->step;
819+
r->len -= index;
818820
Py_RETURN_NONE;
819821
}
820822

@@ -904,13 +906,11 @@ fast_range_iter(long start, long stop, long step, long len)
904906
it->start = start;
905907
it->step = step;
906908
it->len = len;
907-
it->index = 0;
908909
return (PyObject *)it;
909910
}
910911

911912
typedef struct {
912913
PyObject_HEAD
913-
PyObject *index;
914914
PyObject *start;
915915
PyObject *step;
916916
PyObject *len;
@@ -919,7 +919,8 @@ typedef struct {
919919
static PyObject *
920920
longrangeiter_len(longrangeiterobject *r, PyObject *no_args)
921921
{
922-
return PyNumber_Subtract(r->len, r->index);
922+
Py_INCREF(r->len);
923+
return r->len;
923924
}
924925

925926
static PyObject *
@@ -946,8 +947,8 @@ longrangeiter_reduce(longrangeiterobject *r, PyObject *Py_UNUSED(ignored))
946947
}
947948

948949
/* return the result */
949-
return Py_BuildValue(
950-
"N(N)O", _PyEval_GetBuiltin(&_Py_ID(iter)), range, r->index);
950+
return Py_BuildValue("N(N)O", _PyEval_GetBuiltin(&_Py_ID(iter)),
951+
range, Py_None);
951952
}
952953

953954
static PyObject *
@@ -970,7 +971,22 @@ longrangeiter_setstate(longrangeiterobject *r, PyObject *state)
970971
if (cmp > 0)
971972
state = r->len;
972973
}
973-
Py_XSETREF(r->index, Py_NewRef(state));
974+
PyObject *product = PyNumber_Multiply(state, r->step);
975+
if (product == NULL)
976+
return NULL;
977+
PyObject *new_start = PyNumber_Add(r->start, product);
978+
Py_DECREF(product);
979+
if (new_start == NULL)
980+
return NULL;
981+
PyObject *new_len = PyNumber_Subtract(r->len, state);
982+
if (new_len == NULL) {
983+
Py_DECREF(new_start);
984+
return NULL;
985+
}
986+
PyObject *tmp = r->start;
987+
r->start = new_start;
988+
Py_SETREF(r->len, new_len);
989+
Py_DECREF(tmp);
974990
Py_RETURN_NONE;
975991
}
976992

@@ -987,7 +1003,6 @@ static PyMethodDef longrangeiter_methods[] = {
9871003
static void
9881004
longrangeiter_dealloc(longrangeiterobject *r)
9891005
{
990-
Py_XDECREF(r->index);
9911006
Py_XDECREF(r->start);
9921007
Py_XDECREF(r->step);
9931008
Py_XDECREF(r->len);
@@ -997,29 +1012,21 @@ longrangeiter_dealloc(longrangeiterobject *r)
9971012
static PyObject *
9981013
longrangeiter_next(longrangeiterobject *r)
9991014
{
1000-
PyObject *product, *new_index, *result;
1001-
if (PyObject_RichCompareBool(r->index, r->len, Py_LT) != 1)
1015+
if (PyObject_RichCompareBool(r->len, _PyLong_GetZero(), Py_GT) != 1)
10021016
return NULL;
10031017

1004-
new_index = PyNumber_Add(r->index, _PyLong_GetOne());
1005-
if (!new_index)
1018+
PyObject *new_start = PyNumber_Add(r->start, r->step);
1019+
if (new_start == NULL) {
10061020
return NULL;
1007-
1008-
product = PyNumber_Multiply(r->index, r->step);
1009-
if (!product) {
1010-
Py_DECREF(new_index);
1011-
return NULL;
1012-
}
1013-
1014-
result = PyNumber_Add(r->start, product);
1015-
Py_DECREF(product);
1016-
if (result) {
1017-
Py_SETREF(r->index, new_index);
10181021
}
1019-
else {
1020-
Py_DECREF(new_index);
1022+
PyObject *new_len = PyNumber_Subtract(r->len, _PyLong_GetOne());
1023+
if (new_len == NULL) {
1024+
Py_DECREF(new_start);
1025+
return NULL;
10211026
}
1022-
1027+
PyObject *result = r->start;
1028+
r->start = new_start;
1029+
Py_SETREF(r->len, new_len);
10231030
return result;
10241031
}
10251032

@@ -1108,7 +1115,6 @@ range_iter(PyObject *seq)
11081115
it->start = Py_NewRef(r->start);
11091116
it->step = Py_NewRef(r->step);
11101117
it->len = Py_NewRef(r->length);
1111-
it->index = Py_NewRef(_PyLong_GetZero());
11121118
return (PyObject *)it;
11131119
}
11141120

@@ -1186,7 +1192,7 @@ range_reverse(PyObject *seq, PyObject *Py_UNUSED(ignored))
11861192
it = PyObject_New(longrangeiterobject, &PyLongRangeIter_Type);
11871193
if (it == NULL)
11881194
return NULL;
1189-
it->index = it->start = it->step = NULL;
1195+
it->start = it->step = NULL;
11901196

11911197
/* start + (len - 1) * step */
11921198
it->len = Py_NewRef(range->length);
@@ -1210,7 +1216,6 @@ range_reverse(PyObject *seq, PyObject *Py_UNUSED(ignored))
12101216
if (!it->step)
12111217
goto create_failure;
12121218

1213-
it->index = Py_NewRef(_PyLong_GetZero());
12141219
return (PyObject *)it;
12151220

12161221
create_failure:

‎Python/bytecodes.c

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2620,14 +2620,15 @@ dummy_func(
26202620
STAT_INC(FOR_ITER, hit);
26212621
_Py_CODEUNIT next = next_instr[INLINE_CACHE_ENTRIES_FOR_ITER];
26222622
assert(_PyOpcode_Deopt[_Py_OPCODE(next)] == STORE_FAST);
2623-
if (r->index >= r->len) {
2623+
if (r->len <= 0) {
26242624
STACK_SHRINK(1);
26252625
Py_DECREF(r);
26262626
JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg + 1);
26272627
}
26282628
else {
2629-
long value = (long)(r->start +
2630-
(unsigned long)(r->index++) * r->step);
2629+
long value = r->start;
2630+
r->start = value + r->step;
2631+
r->len--;
26312632
if (_PyLong_AssignValue(&GETLOCAL(_Py_OPARG(next)), value) < 0) {
26322633
goto error;
26332634
}

‎Python/generated_cases.c.h

Lines changed: 4 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)