Skip to content

Commit 60656ea

Browse files
committed
py: Define and use MP_OBJ_ITER_BUF_NSLOTS to get size of stack iter buf.
It improves readability of code and reduces the chance to make a mistake. This patch also fixes a bug with nan-boxing builds by rounding up the calculation of the new NSLOTS variable, giving the correct number of slots (being 4) even if mp_obj_t is larger than the native machine size.
1 parent 507119f commit 60656ea

5 files changed

Lines changed: 24 additions & 18 deletions

File tree

py/compile.c

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3005,6 +3005,7 @@ STATIC void compile_scope(compiler_t *comp, scope_t *scope, pass_kind_t pass) {
30053005
// There are 4 slots on the stack for the iterator, and the first one is
30063006
// NULL to indicate that the second one points to the iterator object.
30073007
if (scope->kind == SCOPE_GEN_EXPR) {
3008+
// TODO static assert that MP_OBJ_ITER_BUF_NSLOTS == 4
30083009
EMIT(load_null);
30093010
compile_load_id(comp, qstr_arg);
30103011
EMIT(load_null);

py/emitbc.c

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -735,7 +735,7 @@ void mp_emit_bc_unwind_jump(emit_t *emit, mp_uint_t label, mp_uint_t except_dept
735735
// need to pop the iterator if we are breaking out of a for loop
736736
emit_write_bytecode_byte(emit, MP_BC_POP_TOP);
737737
// also pop the iter_buf
738-
for (size_t i = 0; i < sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) - 1; ++i) {
738+
for (size_t i = 0; i < MP_OBJ_ITER_BUF_NSLOTS - 1; ++i) {
739739
emit_write_bytecode_byte(emit, MP_BC_POP_TOP);
740740
}
741741
}
@@ -778,7 +778,7 @@ void mp_emit_bc_end_finally(emit_t *emit) {
778778
}
779779

780780
void mp_emit_bc_get_iter(emit_t *emit, bool use_stack) {
781-
emit_bc_pre(emit, use_stack ? sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) - 1 : 0);
781+
emit_bc_pre(emit, use_stack ? MP_OBJ_ITER_BUF_NSLOTS - 1 : 0);
782782
emit_write_bytecode_byte(emit, use_stack ? MP_BC_GET_ITER_STACK : MP_BC_GET_ITER);
783783
}
784784

@@ -788,7 +788,7 @@ void mp_emit_bc_for_iter(emit_t *emit, mp_uint_t label) {
788788
}
789789

790790
void mp_emit_bc_for_iter_end(emit_t *emit) {
791-
emit_bc_pre(emit, -(sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t)));
791+
emit_bc_pre(emit, -MP_OBJ_ITER_BUF_NSLOTS);
792792
}
793793

794794
void mp_emit_bc_pop_block(emit_t *emit) {

py/emitnative.c

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1761,7 +1761,7 @@ STATIC void emit_native_get_iter(emit_t *emit, bool use_stack) {
17611761
emit_pre_pop_reg(emit, &vtype, REG_ARG_1);
17621762
assert(vtype == VTYPE_PYOBJ);
17631763
if (use_stack) {
1764-
emit_get_stack_pointer_to_reg_for_push(emit, REG_ARG_2, sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t));
1764+
emit_get_stack_pointer_to_reg_for_push(emit, REG_ARG_2, MP_OBJ_ITER_BUF_NSLOTS);
17651765
emit_call(emit, MP_F_NATIVE_GETITER);
17661766
} else {
17671767
// mp_getiter will allocate the iter_buf on the heap
@@ -1773,8 +1773,8 @@ STATIC void emit_native_get_iter(emit_t *emit, bool use_stack) {
17731773

17741774
STATIC void emit_native_for_iter(emit_t *emit, mp_uint_t label) {
17751775
emit_native_pre(emit);
1776-
emit_get_stack_pointer_to_reg_for_pop(emit, REG_ARG_1, sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t));
1777-
adjust_stack(emit, 4);
1776+
emit_get_stack_pointer_to_reg_for_pop(emit, REG_ARG_1, MP_OBJ_ITER_BUF_NSLOTS);
1777+
adjust_stack(emit, MP_OBJ_ITER_BUF_NSLOTS);
17781778
emit_call(emit, MP_F_NATIVE_ITERNEXT);
17791779
ASM_MOV_IMM_TO_REG(emit->as, (mp_uint_t)MP_OBJ_STOP_ITERATION, REG_TEMP1);
17801780
ASM_JUMP_IF_REG_EQ(emit->as, REG_RET, REG_TEMP1, label);
@@ -1784,7 +1784,7 @@ STATIC void emit_native_for_iter(emit_t *emit, mp_uint_t label) {
17841784
STATIC void emit_native_for_iter_end(emit_t *emit) {
17851785
// adjust stack counter (we get here from for_iter ending, which popped the value for us)
17861786
emit_native_pre(emit);
1787-
adjust_stack(emit, -(sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t)));
1787+
adjust_stack(emit, -MP_OBJ_ITER_BUF_NSLOTS);
17881788
emit_post(emit);
17891789
}
17901790

py/obj.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,10 @@ typedef struct _mp_obj_iter_buf_t {
422422
mp_obj_t buf[3];
423423
} mp_obj_iter_buf_t;
424424

425+
// The number of slots that an mp_obj_iter_buf_t needs on the Python value stack.
426+
// It's rounded up in case mp_obj_base_t is smaller than mp_obj_t (eg for OBJ_REPR_D).
427+
#define MP_OBJ_ITER_BUF_NSLOTS ((sizeof(mp_obj_iter_buf_t) + sizeof(mp_obj_t) - 1) / sizeof(mp_obj_t))
428+
425429
typedef void (*mp_print_fun_t)(const mp_print_t *print, mp_obj_t o, mp_print_kind_t kind);
426430
typedef mp_obj_t (*mp_make_new_fun_t)(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args);
427431
typedef mp_obj_t (*mp_call_fun_t)(mp_obj_t fun, size_t n_args, size_t n_kw, const mp_obj_t *args);

py/vm.c

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -684,7 +684,7 @@ unwind_jump:;
684684
if (unum != 0) {
685685
// pop iter and iter_buf
686686
sp--;
687-
sp -= sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t);
687+
sp -= MP_OBJ_ITER_BUF_NSLOTS;
688688
}
689689
DISPATCH_WITH_PEND_EXC_CHECK();
690690
}
@@ -729,19 +729,20 @@ unwind_jump:;
729729
SET_TOP(mp_getiter(TOP(), NULL));
730730
DISPATCH();
731731

732-
// An iterator for a for-loop takes 4 slots on the stack. They are either
733-
// used to store the iterator object itself, or the first slot is NULL and
732+
// An iterator for a for-loop takes MP_OBJ_ITER_BUF_NSLOTS slots on
733+
// the Python value stack. These slots are either used to store the
734+
// iterator object itself, or the first slot is MP_OBJ_NULL and
734735
// the second slot holds a reference to the iterator object.
735736
ENTRY(MP_BC_GET_ITER_STACK): {
736737
MARK_EXC_IP_SELECTIVE();
737738
mp_obj_t obj = TOP();
738739
mp_obj_iter_buf_t *iter_buf = (mp_obj_iter_buf_t*)sp;
739-
sp += sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) - 1;
740+
sp += MP_OBJ_ITER_BUF_NSLOTS - 1;
740741
obj = mp_getiter(obj, iter_buf);
741742
if (obj != MP_OBJ_FROM_PTR(iter_buf)) {
742743
// Iterator didn't use the stack so indicate that with MP_OBJ_NULL.
743-
sp[-3] = MP_OBJ_NULL;
744-
sp[-2] = obj;
744+
sp[-MP_OBJ_ITER_BUF_NSLOTS + 1] = MP_OBJ_NULL;
745+
sp[-MP_OBJ_ITER_BUF_NSLOTS + 2] = obj;
745746
}
746747
DISPATCH();
747748
}
@@ -751,14 +752,14 @@ unwind_jump:;
751752
DECODE_ULABEL; // the jump offset if iteration finishes; for labels are always forward
752753
code_state->sp = sp;
753754
mp_obj_t obj;
754-
if (sp[-3] == MP_OBJ_NULL) {
755-
obj = sp[-2];
755+
if (sp[-MP_OBJ_ITER_BUF_NSLOTS + 1] == MP_OBJ_NULL) {
756+
obj = sp[-MP_OBJ_ITER_BUF_NSLOTS + 2];
756757
} else {
757-
obj = MP_OBJ_FROM_PTR(&sp[-3]);
758+
obj = MP_OBJ_FROM_PTR(&sp[-MP_OBJ_ITER_BUF_NSLOTS + 1]);
758759
}
759760
mp_obj_t value = mp_iternext_allow_raise(obj);
760761
if (value == MP_OBJ_STOP_ITERATION) {
761-
sp -= 4; // pop the exhausted iterator
762+
sp -= MP_OBJ_ITER_BUF_NSLOTS; // pop the exhausted iterator
762763
ip += ulab; // jump to after for-block
763764
} else {
764765
PUSH(value); // push the next iteration value
@@ -1334,7 +1335,7 @@ unwind_jump:;
13341335
const byte *ip = code_state->ip + 1;
13351336
DECODE_ULABEL; // the jump offset if iteration finishes; for labels are always forward
13361337
code_state->ip = ip + ulab; // jump to after for-block
1337-
code_state->sp -= 4; // pop the exhausted iterator
1338+
code_state->sp -= MP_OBJ_ITER_BUF_NSLOTS; // pop the exhausted iterator
13381339
goto outer_dispatch_loop; // continue with dispatch loop
13391340
} else if (*code_state->ip == MP_BC_YIELD_FROM) {
13401341
// StopIteration inside yield from call means return a value of

0 commit comments

Comments
 (0)