Skip to content

Commit 088740e

Browse files
dpgeorgedpgeorge
committed
py: Optimise storage of iterator so it takes only 4 slots on Py stack.
1 parent 6e769da commit 088740e

6 files changed

Lines changed: 65 additions & 29 deletions

File tree

‎py/compile.c‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2887,7 +2887,7 @@ STATIC void compile_scope_comp_iter(compiler_t *comp, mp_parse_node_struct_t *pn
28872887
EMIT(yield_value);
28882888
EMIT(pop_top);
28892889
} else {
2890-
EMIT_ARG(store_comp, comp->scope_cur->kind, 5 * for_depth + 6);
2890+
EMIT_ARG(store_comp, comp->scope_cur->kind, 4 * for_depth + 5);
28912891
}
28922892
} else if (MP_PARSE_NODE_IS_STRUCT_KIND(pn_iter, PN_comp_if)) {
28932893
// if condition
@@ -3070,13 +3070,13 @@ STATIC void compile_scope(compiler_t *comp, scope_t *scope, pass_kind_t pass) {
30703070
#endif
30713071
}
30723072

3073-
// dummy 4 objects
3074-
EMIT(load_null);
3073+
// There are 4 slots on the stack for the iterator, and the first one is
3074+
// NULL to indicate that the second one points to the iterator object.
30753075
EMIT(load_null);
3076+
compile_load_id(comp, qstr_arg);
30763077
EMIT(load_null);
30773078
EMIT(load_null);
30783079

3079-
compile_load_id(comp, qstr_arg);
30803080
compile_scope_comp_iter(comp, pns_comp_for, pns->nodes[0], 0);
30813081

30823082
if (scope->kind == SCOPE_GEN_EXPR) {

‎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); ++i) {
738+
for (size_t i = 0; i < sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) - 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) : 0);
781+
emit_bc_pre(emit, use_stack ? sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) - 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, bool use_stack) {
791-
emit_bc_pre(emit, use_stack ? -1 - sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) : -1);
791+
emit_bc_pre(emit, -(use_stack ? sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) : 1));
792792
}
793793

794794
void mp_emit_bc_pop_block(emit_t *emit) {

‎py/emitnative.c‎

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,8 @@ STATIC byte mp_f_n_args[MP_F_NUMBER_OF] = {
105105
[MP_F_NATIVE_CALL_FUNCTION_N_KW] = 3,
106106
[MP_F_CALL_METHOD_N_KW] = 3,
107107
[MP_F_CALL_METHOD_N_KW_VAR] = 3,
108-
[MP_F_GETITER] = 1,
109-
[MP_F_ITERNEXT] = 1,
108+
[MP_F_NATIVE_GETITER] = 2,
109+
[MP_F_NATIVE_ITERNEXT] = 1,
110110
[MP_F_NLR_PUSH] = 1,
111111
[MP_F_NLR_POP] = 0,
112112
[MP_F_NATIVE_RAISE] = 1,
@@ -1808,20 +1808,20 @@ STATIC void emit_native_get_iter(emit_t *emit, bool use_stack) {
18081808
assert(vtype == VTYPE_PYOBJ);
18091809
if (use_stack) {
18101810
emit_get_stack_pointer_to_reg_for_push(emit, REG_ARG_2, sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t));
1811+
emit_call(emit, MP_F_NATIVE_GETITER);
18111812
} else {
18121813
// mp_getiter will allocate the iter_buf on the heap
18131814
ASM_MOV_IMM_TO_REG(emit->as, 0, REG_ARG_2);
1815+
emit_call(emit, MP_F_NATIVE_GETITER);
1816+
emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET);
18141817
}
1815-
emit_call(emit, MP_F_GETITER);
1816-
emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET);
18171818
}
18181819

18191820
STATIC void emit_native_for_iter(emit_t *emit, mp_uint_t label) {
18201821
emit_native_pre(emit);
1821-
vtype_kind_t vtype;
1822-
emit_access_stack(emit, 1, &vtype, REG_ARG_1);
1823-
assert(vtype == VTYPE_PYOBJ);
1824-
emit_call(emit, MP_F_ITERNEXT);
1822+
emit_get_stack_pointer_to_reg_for_pop(emit, REG_ARG_1, sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t));
1823+
adjust_stack(emit, 4);
1824+
emit_call(emit, MP_F_NATIVE_ITERNEXT);
18251825
ASM_MOV_IMM_TO_REG(emit->as, (mp_uint_t)MP_OBJ_STOP_ITERATION, REG_TEMP1);
18261826
ASM_JUMP_IF_REG_EQ(emit->as, REG_RET, REG_TEMP1, label);
18271827
emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET);
@@ -1830,10 +1830,7 @@ STATIC void emit_native_for_iter(emit_t *emit, mp_uint_t label) {
18301830
STATIC void emit_native_for_iter_end(emit_t *emit, bool use_stack) {
18311831
// adjust stack counter (we get here from for_iter ending, which popped the value for us)
18321832
emit_native_pre(emit);
1833-
adjust_stack(emit, -1);
1834-
if (use_stack) {
1835-
adjust_stack(emit, -(sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t)));
1836-
}
1833+
adjust_stack(emit, -(use_stack ? sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) : 1));
18371834
emit_post(emit);
18381835
}
18391836

‎py/nativeglue.c‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,32 @@ void mp_native_raise(mp_obj_t o) {
9898
}
9999
}
100100

101+
// wrapper that handles iterator buffer
102+
STATIC mp_obj_t mp_native_getiter(mp_obj_t obj, mp_obj_iter_buf_t *iter) {
103+
if (iter == NULL) {
104+
return mp_getiter(obj, NULL);
105+
} else {
106+
obj = mp_getiter(obj, iter);
107+
if (obj != MP_OBJ_FROM_PTR(iter)) {
108+
// Iterator didn't use the stack so indicate that with MP_OBJ_NULL.
109+
iter->base.type = MP_OBJ_NULL;
110+
iter->buf[0] = obj;
111+
}
112+
return NULL;
113+
}
114+
}
115+
116+
// wrapper that handles iterator buffer
117+
STATIC mp_obj_t mp_native_iternext(mp_obj_iter_buf_t *iter) {
118+
mp_obj_t obj;
119+
if (iter->base.type == MP_OBJ_NULL) {
120+
obj = iter->buf[0];
121+
} else {
122+
obj = MP_OBJ_FROM_PTR(iter);
123+
}
124+
return mp_iternext(obj);
125+
}
126+
101127
// these must correspond to the respective enum in runtime0.h
102128
void *const mp_fun_table[MP_F_NUMBER_OF] = {
103129
mp_convert_obj_to_native,
@@ -127,8 +153,8 @@ void *const mp_fun_table[MP_F_NUMBER_OF] = {
127153
mp_native_call_function_n_kw,
128154
mp_call_method_n_kw,
129155
mp_call_method_n_kw_var,
130-
mp_getiter,
131-
mp_iternext,
156+
mp_native_getiter,
157+
mp_native_iternext,
132158
nlr_push,
133159
nlr_pop,
134160
mp_native_raise,

‎py/runtime0.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,8 @@ typedef enum {
127127
MP_F_NATIVE_CALL_FUNCTION_N_KW,
128128
MP_F_CALL_METHOD_N_KW,
129129
MP_F_CALL_METHOD_N_KW_VAR,
130-
MP_F_GETITER,
131-
MP_F_ITERNEXT,
130+
MP_F_NATIVE_GETITER,
131+
MP_F_NATIVE_ITERNEXT,
132132
MP_F_NLR_PUSH,
133133
MP_F_NLR_POP,
134134
MP_F_NATIVE_RAISE,

‎py/vm.c‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -728,23 +728,36 @@ unwind_jump:;
728728
SET_TOP(mp_getiter(TOP(), NULL));
729729
DISPATCH();
730730

731+
// An iterator for a for-loop takes 4 slots on the stack. They are either
732+
// used to store the iterator object itself, or the first slot is NULL and
733+
// the second slot holds a reference to the iterator object.
731734
ENTRY(MP_BC_GET_ITER_STACK): {
732735
MARK_EXC_IP_SELECTIVE();
733736
mp_obj_t obj = TOP();
734737
mp_obj_iter_buf_t *iter_buf = (mp_obj_iter_buf_t*)sp;
735-
sp += sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t);
736-
SET_TOP(mp_getiter(obj, iter_buf));
738+
sp += sizeof(mp_obj_iter_buf_t) / sizeof(mp_obj_t) - 1;
739+
obj = mp_getiter(obj, iter_buf);
740+
if (obj != MP_OBJ_FROM_PTR(iter_buf)) {
741+
// Iterator didn't use the stack so indicate that with MP_OBJ_NULL.
742+
sp[-3] = MP_OBJ_NULL;
743+
sp[-2] = obj;
744+
}
737745
DISPATCH();
738746
}
739747

740748
ENTRY(MP_BC_FOR_ITER): {
741749
MARK_EXC_IP_SELECTIVE();
742750
DECODE_ULABEL; // the jump offset if iteration finishes; for labels are always forward
743751
code_state->sp = sp;
744-
assert(TOP());
745-
mp_obj_t value = mp_iternext_allow_raise(TOP());
752+
mp_obj_t obj;
753+
if (sp[-3] == MP_OBJ_NULL) {
754+
obj = sp[-2];
755+
} else {
756+
obj = MP_OBJ_FROM_PTR(&sp[-3]);
757+
}
758+
mp_obj_t value = mp_iternext_allow_raise(obj);
746759
if (value == MP_OBJ_STOP_ITERATION) {
747-
sp -= 5; // pop the exhausted iterator
760+
sp -= 4; // pop the exhausted iterator
748761
ip += ulab; // jump to after for-block
749762
} else {
750763
PUSH(value); // push the next iteration value
@@ -1294,7 +1307,7 @@ unwind_jump:;
12941307
const byte *ip = code_state->ip + 1;
12951308
DECODE_ULABEL; // the jump offset if iteration finishes; for labels are always forward
12961309
code_state->ip = ip + ulab; // jump to after for-block
1297-
code_state->sp -= 5; // pop the exhausted iterator
1310+
code_state->sp -= 4; // pop the exhausted iterator
12981311
goto outer_dispatch_loop; // continue with dispatch loop
12991312
} else if (*code_state->ip == MP_BC_YIELD_FROM) {
13001313
// StopIteration inside yield from call means return a value of

0 commit comments

Comments
 (0)