Skip to content

Commit e429410

Browse files
committed
Make FOR_ITER_RANGE mutate the local
1 parent dc80fda commit e429410

File tree

4 files changed

+50
-30
lines changed

4 files changed

+50
-30
lines changed

Include/internal/pycore_range.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@ extern "C" {
1010

1111
typedef struct {
1212
PyObject_HEAD
13-
long index;
14-
long start;
15-
long step;
16-
long len;
13+
sdigit index;
14+
sdigit start;
15+
sdigit step;
16+
sdigit len;
1717
} _PyRangeIterObject;
1818

1919
#ifdef __cplusplus

Objects/rangeobject.c

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -770,7 +770,7 @@ rangeiter_next(_PyRangeIterObject *r)
770770
/* cast to unsigned to avoid possible signed overflow
771771
in intermediate calculations. */
772772
return PyLong_FromLong((long)(r->start +
773-
(unsigned long)(r->index++) * r->step));
773+
(digit)(r->index++) * r->step));
774774
return NULL;
775775
}
776776

@@ -911,9 +911,9 @@ fast_range_iter(long start, long stop, long step, long len)
911911
_PyRangeIterObject *it = PyObject_New(_PyRangeIterObject, &PyRangeIter_Type);
912912
if (it == NULL)
913913
return NULL;
914-
it->start = start;
915-
it->step = step;
916-
it->len = len;
914+
it->start = Py_SAFE_DOWNCAST(start, long, sdigit);
915+
it->step = Py_SAFE_DOWNCAST(step, long, sdigit);
916+
it->len = Py_SAFE_DOWNCAST(len, long, sdigit);
917917
it->index = 0;
918918
return (PyObject *)it;
919919
}
@@ -1097,20 +1097,13 @@ range_iter(PyObject *seq)
10971097
goto long_range;
10981098
}
10991099
ulen = get_len_of_range(lstart, lstop, lstep);
1100-
if (ulen > (unsigned long)LONG_MAX) {
1100+
if (ulen > PyLong_MASK ||
1101+
lstart > PyLong_MASK || lstart < -(long)PyLong_MASK ||
1102+
lstop > PyLong_MASK || lstop < -(long)PyLong_MASK ||
1103+
lstep > PyLong_MASK || lstep < -(long)PyLong_MASK)
1104+
{
11011105
goto long_range;
11021106
}
1103-
/* check for potential overflow of lstart + ulen * lstep */
1104-
if (ulen) {
1105-
if (lstep > 0) {
1106-
if (lstop > LONG_MAX - (lstep - 1))
1107-
goto long_range;
1108-
}
1109-
else {
1110-
if (lstop < LONG_MIN + (-1 - lstep))
1111-
goto long_range;
1112-
}
1113-
}
11141107
return fast_range_iter(lstart, lstop, lstep, (long)ulen);
11151108

11161109
long_range:

Python/ceval.c

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4399,17 +4399,42 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
43994399
_PyRangeIterObject *r = (_PyRangeIterObject *)TOP();
44004400
DEOPT_IF(Py_TYPE(r) != &PyRangeIter_Type, FOR_ITER);
44014401
STAT_INC(FOR_ITER, hit);
4402-
if (r->index < r->len) {
4403-
PyObject *res = PyLong_FromLong(
4404-
(long)(r->start + (unsigned long)(r->index++) * r->step));
4405-
if (res == NULL) {
4406-
goto error;
4402+
_Py_CODEUNIT next = next_instr[INLINE_CACHE_ENTRIES_FOR_ITER];
4403+
assert(_PyOpcode_Deopt[_Py_OPCODE(next)] == STORE_FAST);
4404+
PyObject **local_ptr = &GETLOCAL(_Py_OPARG(next));
4405+
PyObject *local = *local_ptr;
4406+
if (r->index >= r->len) {
4407+
goto iterator_exhausted_no_error;
4408+
}
4409+
JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + 1);
4410+
sdigit value = r->start + (digit)(r->index++) * r->step;
4411+
if (value < _PY_NSMALLPOSINTS && value >= -_PY_NSMALLNEGINTS) {
4412+
*local_ptr = Py_NewRef(&_PyLong_SMALL_INTS[_PY_NSMALLNEGINTS+value]);
4413+
Py_XDECREF(local);
4414+
NOTRACE_DISPATCH();
4415+
}
4416+
if (local && PyLong_CheckExact(local) && Py_REFCNT(local) == 1) {
4417+
if (value > 0) {
4418+
assert(value <= PyLong_MASK);
4419+
((PyLongObject *)local)->ob_digit[0] = value;
4420+
Py_SET_SIZE(local, 1);
4421+
}
4422+
else {
4423+
assert(value >= -(sdigit)PyLong_MASK);
4424+
((PyLongObject *)local)->ob_digit[0] = -(sdigit)value;
4425+
Py_SET_SIZE(local, -1);
44074426
}
4408-
PUSH(res);
4409-
JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER);
44104427
NOTRACE_DISPATCH();
44114428
}
4412-
goto iterator_exhausted_no_error;
4429+
PyObject *res = PyLong_FromLong(value);
4430+
if (res == NULL) {
4431+
// undo the JUMPBY
4432+
next_instr -= INLINE_CACHE_ENTRIES_FOR_ITER + 1;
4433+
goto error;
4434+
}
4435+
*local_ptr = res;
4436+
Py_XDECREF(local);
4437+
NOTRACE_DISPATCH();
44134438
}
44144439

44154440
TARGET(BEFORE_ASYNC_WITH) {

Python/specialize.c

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2064,16 +2064,18 @@ _Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr)
20642064
assert(_PyOpcode_Caches[FOR_ITER] == INLINE_CACHE_ENTRIES_FOR_ITER);
20652065
_PyForIterCache *cache = (_PyForIterCache *)(instr + 1);
20662066
PyTypeObject *tp = Py_TYPE(iter);
2067+
_Py_CODEUNIT next = instr[1+INLINE_CACHE_ENTRIES_FOR_ITER];
2068+
int next_op = _PyOpcode_Deopt[_Py_OPCODE(next)];
20672069
if (tp == &PyListIter_Type) {
20682070
_Py_SET_OPCODE(*instr, FOR_ITER_LIST);
20692071
goto success;
20702072
}
2071-
else if (tp == &PyRangeIter_Type) {
2073+
else if (tp == &PyRangeIter_Type && next_op == STORE_FAST) {
20722074
_Py_SET_OPCODE(*instr, FOR_ITER_RANGE);
20732075
goto success;
20742076
}
20752077
else {
2076-
SPECIALIZATION_FAIL(JUMP_BACKWARD,
2078+
SPECIALIZATION_FAIL(FOR_ITER,
20772079
_PySpecialization_ClassifyIterator(iter));
20782080
goto failure;
20792081
}

0 commit comments

Comments
 (0)