commit 4d2bab14445b7261e4598f2aad5032d08e95b3ef
author Damien George <damien.p.george@gmail.com> Fri Feb 10 15:39:33 2017 +1100
committer Damien George <damien.p.george@gmail.com> Thu Feb 16 18:38:07 2017 +1100
tree 81d565cfa893425a70ad1d0bd7759d670347f367
parent 861b00178370d507516273ea95dd48acd31f36ff

py/compile: Optimise list/dict/set comprehensions to use stack iter.

py/compile.c

diff --git a/py/compile.c b/py/compile.c
index 3967350..58f6631 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -2372,7 +2372,9 @@
     close_over_variables_etc(comp, this_scope, 0, 0);
 
     compile_node(comp, pns_comp_for->nodes[1]); // source of the iterator
-    EMIT_ARG(get_iter, false);
+    if (kind == SCOPE_GEN_EXPR) {
+        EMIT_ARG(get_iter, false);
+    }
     EMIT_ARG(call_function, 1, 0, 0);
 }
 
@@ -3072,10 +3074,15 @@
 
         // There are 4 slots on the stack for the iterator, and the first one is
         // NULL to indicate that the second one points to the iterator object.
-        EMIT(load_null);
-        compile_load_id(comp, qstr_arg);
-        EMIT(load_null);
-        EMIT(load_null);
+        if (scope->kind == SCOPE_GEN_EXPR) {
+            EMIT(load_null);
+            compile_load_id(comp, qstr_arg);
+            EMIT(load_null);
+            EMIT(load_null);
+        } else {
+            compile_load_id(comp, qstr_arg);
+            EMIT_ARG(get_iter, true);
+        }
 
         compile_scope_comp_iter(comp, pns_comp_for, pns->nodes[0], 0);