py: restrict further when for-range optimisation is done.

dpgeorge · dpgeorge · commit d79338969b14 · 2013-11-28T19:12:18.000Z
diff --git a/py/compile.c b/py/compile.c
@@ -48,6 +48,7 @@ typedef struct _compiler_t {
     qstr qstr_native;
     qstr qstr_viper;
     qstr qstr_asm_thumb;
+    qstr qstr_range;
 
     bool is_repl;
     pass_kind_t pass;
@@ -1409,9 +1410,9 @@ void compile_for_stmt(compiler_t *comp, py_parse_node_struct_t *pns) {
     // this bit optimises: for <x> in range(...), turning it into an explicitly incremented variable
     // this is actually slower, but uses no heap memory
     // for viper it will be much, much faster
-    if (/*comp->scope_cur->emit_options == EMIT_OPT_VIPER &&*/ PY_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[1], PN_power)) {
+    if (/*comp->scope_cur->emit_options == EMIT_OPT_VIPER &&*/ PY_PARSE_NODE_IS_ID(pns->nodes[0]) && PY_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[1], PN_power)) {
         py_parse_node_struct_t *pns_it = (py_parse_node_struct_t*)pns->nodes[1];
-        if (PY_PARSE_NODE_IS_ID(pns_it->nodes[0]) && PY_PARSE_NODE_IS_STRUCT_KIND(pns_it->nodes[1], PN_trailer_paren) && PY_PARSE_NODE_IS_NULL(pns_it->nodes[2])) {
+        if (PY_PARSE_NODE_IS_ID(pns_it->nodes[0]) && PY_PARSE_NODE_LEAF_ARG(pns_it->nodes[0]) == comp->qstr_range && PY_PARSE_NODE_IS_STRUCT_KIND(pns_it->nodes[1], PN_trailer_paren) && PY_PARSE_NODE_IS_NULL(pns_it->nodes[2])) {
             py_parse_node_t pn_range_args = ((py_parse_node_struct_t*)pns_it->nodes[1])->nodes[0];
             py_parse_node_t *args;
             int n_args = list_get(&pn_range_args, PN_arglist, &args);
@@ -2853,6 +2854,7 @@ bool py_compile(py_parse_node_t pn, bool is_repl) {
     comp->qstr_native = qstr_from_str_static("native");
     comp->qstr_viper = qstr_from_str_static("viper");
     comp->qstr_asm_thumb = qstr_from_str_static("asm_thumb");
+    comp->qstr_range = qstr_from_str_static("range");
 
     comp->is_repl = is_repl;
     comp->had_error = false;
