Skip to content

Commit ad97f2a

Browse files
dpgeorgedpgeorge
committed
Merge pull request adafruit#136 from pfalcon/for-range-downto
compile_for_stmt_optimised_range(): Properly handle negative & unknown s...
2 parents e9b4b7a + 899c69f commit ad97f2a

2 files changed

Lines changed: 31 additions & 6 deletions

File tree

‎py/compile.c‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1460,10 +1460,14 @@ void compile_for_stmt_optimised_range(compiler_t *comp, mp_parse_node_t pn_var,
14601460

14611461
EMIT(label_assign, continue_label);
14621462

1463-
// compile: if var < end: goto top
1463+
// compile: if var <cond> end: goto top
14641464
compile_node(comp, pn_var);
14651465
compile_node(comp, pn_end);
1466-
EMIT(compare_op, RT_COMPARE_OP_LESS);
1466+
if (MP_PARSE_NODE_LEAF_ARG(pn_step) >= 0) {
1467+
EMIT(compare_op, RT_COMPARE_OP_LESS);
1468+
} else {
1469+
EMIT(compare_op, RT_COMPARE_OP_MORE);
1470+
}
14671471
EMIT(pop_jump_if_true, top_label);
14681472

14691473
// break/continue apply to outer loop (if any) in the else block
@@ -1482,14 +1486,19 @@ void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
14821486
// for viper it will be much, much faster
14831487
if (/*comp->scope_cur->emit_options == EMIT_OPT_VIPER &&*/ MP_PARSE_NODE_IS_ID(pns->nodes[0]) && MP_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[1], PN_power)) {
14841488
mp_parse_node_struct_t *pns_it = (mp_parse_node_struct_t*)pns->nodes[1];
1485-
if (MP_PARSE_NODE_IS_ID(pns_it->nodes[0]) && MP_PARSE_NODE_LEAF_ARG(pns_it->nodes[0]) == MP_QSTR_range && MP_PARSE_NODE_IS_STRUCT_KIND(pns_it->nodes[1], PN_trailer_paren) && MP_PARSE_NODE_IS_NULL(pns_it->nodes[2])) {
1489+
if (MP_PARSE_NODE_IS_ID(pns_it->nodes[0])
1490+
&& MP_PARSE_NODE_LEAF_ARG(pns_it->nodes[0]) == MP_QSTR_range
1491+
&& MP_PARSE_NODE_IS_STRUCT_KIND(pns_it->nodes[1], PN_trailer_paren)
1492+
&& MP_PARSE_NODE_IS_NULL(pns_it->nodes[2])) {
14861493
mp_parse_node_t pn_range_args = ((mp_parse_node_struct_t*)pns_it->nodes[1])->nodes[0];
14871494
mp_parse_node_t *args;
14881495
int n_args = list_get(&pn_range_args, PN_arglist, &args);
1496+
mp_parse_node_t pn_range_start;
1497+
mp_parse_node_t pn_range_end;
1498+
mp_parse_node_t pn_range_step;
1499+
bool optimize = false;
14891500
if (1 <= n_args && n_args <= 3) {
1490-
mp_parse_node_t pn_range_start;
1491-
mp_parse_node_t pn_range_end;
1492-
mp_parse_node_t pn_range_step;
1501+
optimize = true;
14931502
if (n_args == 1) {
14941503
pn_range_start = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, 0);
14951504
pn_range_end = args[0];
@@ -1502,7 +1511,13 @@ void compile_for_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
15021511
pn_range_start = args[0];
15031512
pn_range_end = args[1];
15041513
pn_range_step = args[2];
1514+
// We need to know sign of step. This is possible only if it's constant
1515+
if (!MP_PARSE_NODE_IS_SMALL_INT(pn_range_step)) {
1516+
optimize = false;
1517+
}
15051518
}
1519+
}
1520+
if (optimize) {
15061521
compile_for_stmt_optimised_range(comp, pns->nodes[0], pn_range_start, pn_range_end, pn_range_step, pns->nodes[2], pns->nodes[3]);
15071522
return;
15081523
}

‎tests/basics/tests/for1.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,3 +7,13 @@ def f():
77
print(x, y, z)
88

99
f()
10+
11+
# range with negative step
12+
for i in range(3, -1, -1):
13+
print(i)
14+
15+
a = -1
16+
# range with non-constant step - we optimize constant steps, so this
17+
# will be executed differently
18+
for i in range(3, -1, a):
19+
print(i)

0 commit comments

Comments
 (0)