Skip to content

Commit aa5b0bd

Browse files
markshannonmarkshannon
committed
Allow range starts other than zero
1 parent 76f1814 commit aa5b0bd

File tree

10 files changed

+100
-33
lines changed

10 files changed

+100
-33
lines changed

‎Include/internal/pycore_range.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ typedef struct {
1717

1818
// Does this range have start == 0, step == 1 and step in compact int range?
1919
int _PyRange_IsSimpleCompact(PyObject *range);
20+
Py_ssize_t _PyRange_GetStartIfCompact(PyObject *range);
2021
Py_ssize_t _PyRange_GetStopIfCompact(PyObject *range);
2122

2223
#ifdef __cplusplus

‎Include/internal/pycore_stackref.h

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,14 @@ PyStackRef_IsTaggedInt(_PyStackRef ref)
101101
return (ref.index & 1) == 1;
102102
}
103103

104+
static inline bool
105+
PyStackRef_TaggedIntLessThan(_PyStackRef a, _PyStackRef b)
106+
{
107+
assert(PyStackRef_IsTaggedInt(a));
108+
assert(PyStackRef_IsTaggedInt(b));
109+
return ((intptr_t)a.bits) < ((intptr_t)b.bits);
110+
}
111+
104112
static inline PyObject *
105113
_PyStackRef_AsPyObjectBorrow(_PyStackRef ref, const char *filename, int linenumber)
106114
{
@@ -274,6 +282,14 @@ PyStackRef_IncrementTaggedIntNoOverflow(_PyStackRef ref)
274282
return (_PyStackRef){ .bits = ref.bits + 4 };
275283
}
276284

285+
static inline bool
286+
PyStackRef_TaggedIntLessThan(_PyStackRef a, _PyStackRef b)
287+
{
288+
assert(PyStackRef_IsTaggedInt(a));
289+
assert(PyStackRef_IsTaggedInt(b));
290+
return ((intptr_t)a.bits) < ((intptr_t)b.bits);
291+
}
292+
277293
#define PyStackRef_IsDeferredOrTaggedInt(ref) (((ref).bits & Py_TAG_REFCNT) != 0)
278294

279295
extern _PyStackRef PyStackRef_BoxInt(_PyStackRef i);

‎Lib/test/test_opcache.py

Lines changed: 38 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1832,15 +1832,6 @@ def for_iter_tuple():
18321832
self.assert_specialized(for_iter_tuple, "FOR_ITER_TUPLE")
18331833
self.assert_no_opcode(for_iter_tuple, "FOR_ITER")
18341834

1835-
r = range(10)
1836-
def for_iter_range():
1837-
for i in r:
1838-
self.assertIn(i, r)
1839-
1840-
for_iter_range()
1841-
self.assert_specialized(for_iter_range, "FOR_ITER_RANGE")
1842-
self.assert_no_opcode(for_iter_range, "FOR_ITER")
1843-
18441835
def for_iter_generator():
18451836
for i in (i for i in range(10)):
18461837
i + 1
@@ -1849,6 +1840,44 @@ def for_iter_generator():
18491840
self.assert_specialized(for_iter_generator, "FOR_ITER_GEN")
18501841
self.assert_no_opcode(for_iter_generator, "FOR_ITER")
18511842

1843+
@cpython_only
1844+
@requires_specialization_ft
1845+
def test_for_iter_range(self):
1846+
r = range(10)
1847+
def for_iter1():
1848+
for i in r:
1849+
self.assertIn(i, r)
1850+
1851+
for_iter1()
1852+
self.assert_specialized(for_iter1, "FOR_ITER_RANGE")
1853+
self.assert_no_opcode(for_iter1, "FOR_ITER")
1854+
1855+
r = range(-10, 0)
1856+
def for_iter2():
1857+
for i in r:
1858+
self.assertIn(i, r)
1859+
1860+
for_iter2()
1861+
self.assert_specialized(for_iter2, "FOR_ITER_RANGE")
1862+
self.assert_no_opcode(for_iter2, "FOR_ITER")
1863+
1864+
r = range(100_000, 100_010)
1865+
def for_iter3():
1866+
for i in r:
1867+
self.assertIn(i, r)
1868+
1869+
for_iter3()
1870+
self.assert_specialized(for_iter3, "FOR_ITER_RANGE")
1871+
self.assert_no_opcode(for_iter3, "FOR_ITER")
1872+
1873+
r = range((1 << 65), (1 << 65)+10)
1874+
def for_iter4():
1875+
for i in r:
1876+
self.assertIn(i, r)
1877+
1878+
for_iter4()
1879+
self.assert_specialized(for_iter4, "FOR_ITER")
1880+
18521881

18531882
if __name__ == "__main__":
18541883
unittest.main()

‎Objects/rangeobject.c

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -160,20 +160,29 @@ int
160160
_PyRange_IsSimpleCompact(PyObject *range) {
161161
assert(PyRange_Check(range));
162162
rangeobject *r = (rangeobject*)range;
163-
if (r->start == _PyLong_GetZero() && r->step == _PyLong_GetOne() &&
164-
_PyLong_IsNonNegativeCompact((PyLongObject *)r->stop)
163+
if (_PyLong_IsCompact((PyLongObject *)r->start) &&
164+
_PyLong_IsCompact((PyLongObject *)r->stop) &&
165+
r->step == _PyLong_GetOne()
165166
) {
166167
return 1;
167168
}
168169
return 0;
169170
}
170171

172+
Py_ssize_t
173+
_PyRange_GetStartIfCompact(PyObject *range) {
174+
assert(PyRange_Check(range));
175+
rangeobject *r = (rangeobject*)range;
176+
assert(_PyLong_IsCompact((PyLongObject *)r->start));
177+
return _PyLong_CompactValue((PyLongObject *)r->start);
178+
}
179+
171180
Py_ssize_t
172181
_PyRange_GetStopIfCompact(PyObject *range) {
173182
assert(PyRange_Check(range));
174183
rangeobject *r = (rangeobject*)range;
175-
assert(_PyLong_IsNonNegativeCompact((PyLongObject *)r->stop));
176-
return _PyLong_GetNonNegativeCompactValue((PyLongObject *)r->stop);
184+
assert(_PyLong_IsCompact((PyLongObject *)r->stop));
185+
return _PyLong_CompactValue((PyLongObject *)r->stop);
177186
}
178187

179188
PyDoc_STRVAR(range_doc,

‎Python/bytecodes.c

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3080,10 +3080,11 @@ dummy_func(
30803080
else {
30813081
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iterable);
30823082
if (tp == &PyRange_Type && _PyRange_IsSimpleCompact(iter_o)) {
3083+
Py_ssize_t start = _PyRange_GetStartIfCompact(iter_o);
30833084
Py_ssize_t stop = _PyRange_GetStopIfCompact(iter_o);
30843085
PyStackRef_CLOSE(iterable);
30853086
iter = PyStackRef_TagInt(stop);
3086-
index_or_null = PyStackRef_TagInt(0);
3087+
index_or_null = PyStackRef_TagInt(start);
30873088
}
30883089
else {
30893090
iter_o = PyObject_GetIter(iter_o);
@@ -3171,8 +3172,7 @@ dummy_func(
31713172
/* before: [iter]; after: [iter, iter()] *or* [] (and jump over END_FOR.) */
31723173
if (PyStackRef_IsTaggedInt(null_or_index)) {
31733174
if (PyStackRef_IsTaggedInt(iter)) {
3174-
if (PyStackRef_Is(iter, null_or_index)) {
3175-
null_or_index = PyStackRef_TagInt(-1);
3175+
if (!PyStackRef_TaggedIntLessThan(null_or_index, iter)) {
31763176
JUMPBY(oparg + 1);
31773177
DISPATCH();
31783178

@@ -3243,14 +3243,11 @@ dummy_func(
32433243

32443244

32453245
inst(INSTRUMENTED_FOR_ITER, (unused/1, iter, null_or_index -- iter, null_or_index, next)) {
3246-
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
32473246
if (PyStackRef_IsTaggedInt(null_or_index)) {
32483247
if (PyStackRef_IsTaggedInt(iter)) {
3249-
if (PyStackRef_Is(iter, null_or_index)) {
3250-
null_or_index = PyStackRef_TagInt(-1);
3248+
if (!PyStackRef_TaggedIntLessThan(null_or_index, iter)) {
32513249
JUMPBY(oparg + 1);
32523250
DISPATCH();
3253-
32543251
}
32553252
next = PyStackRef_BoxInt(null_or_index);
32563253
if (PyStackRef_IsNull(next)) {
@@ -3259,6 +3256,7 @@ dummy_func(
32593256
null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
32603257
}
32613258
else {
3259+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
32623260
next = _PyForIter_NextWithIndex(iter_o, null_or_index);
32633261
if (PyStackRef_IsNull(next)) {
32643262
JUMPBY(oparg + 1);
@@ -3268,6 +3266,7 @@ dummy_func(
32683266
INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
32693267
}
32703268
else {
3269+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
32713270
PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
32723271
if (next_o != NULL) {
32733272
next = PyStackRef_FromPyObjectSteal(next_o);
@@ -3423,7 +3422,7 @@ dummy_func(
34233422
}
34243423

34253424
replaced op(_ITER_JUMP_RANGE, (iter, null_or_index -- iter, null_or_index)) {
3426-
if (PyStackRef_Is(iter, null_or_index)) {
3425+
if (!PyStackRef_TaggedIntLessThan(null_or_index, iter)) {
34273426
// Jump over END_FOR instruction.
34283427
JUMPBY(oparg + 1);
34293428
DISPATCH();
@@ -3432,7 +3431,7 @@ dummy_func(
34323431

34333432
// Only used by Tier 2
34343433
op(_GUARD_NOT_EXHAUSTED_RANGE, (iter, null_or_index -- iter, null_or_index)) {
3435-
EXIT_IF(PyStackRef_Is(iter, null_or_index));
3434+
EXIT_IF(!PyStackRef_TaggedIntLessThan(null_or_index, iter));
34363435
}
34373436

34383437
op(_ITER_NEXT_RANGE, (iter, null_or_index -- iter, null_or_index, next)) {

‎Python/executor_cases.c.h

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

‎Python/generated_cases.c.h

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

‎Python/specialize.c

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2912,6 +2912,7 @@ _Py_Specialize_ForIter(_PyStackRef iter, _PyStackRef null_or_index, _Py_CODEUNIT
29122912

29132913
if (PyStackRef_IsNull(null_or_index)) {
29142914
#ifdef Py_GIL_DISABLED
2915+
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
29152916
// Only specialize for uniquely referenced iterators, so that we know
29162917
// they're only referenced by this one thread. This is more limiting
29172918
// than we need (even `it = iter(mylist); for item in it:` won't get

‎Python/stackrefs.c

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,14 @@ _PyStackRef PyStackRef_TagInt(intptr_t i)
202202
return (_PyStackRef){ .index = (i << 1) + 1 };
203203
}
204204

205+
bool
206+
PyStackRef_TaggedIntLessThan(_PyStackRef a, _PyStackRef b)
207+
{
208+
assert(PyStackRef_IsTaggedInt(a));
209+
assert(PyStackRef_IsTaggedInt(b));
210+
return a.bits < b.bits;
211+
}
212+
205213
intptr_t
206214
PyStackRef_UntagInt(_PyStackRef i)
207215
{

‎Tools/cases_generator/analyzer.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -685,6 +685,7 @@ def has_error_without_pop(op: parser.CodeDef) -> bool:
685685
"_PyRange_GetStopIfCompact",
686686
"PyStackRef_BoxInt",
687687
"PyStackRef_TYPE",
688+
"PyStackRef_TaggedIntLessThan",
688689
)
689690

690691

0 commit comments

Comments
 (0)