do not create new generator scopes while translating the scopes if we are in the iterator expression of the outermost generator comprehension

Author: cosminbasca

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/parser/ScopeInfo.java

@@ -184,7 +184,7 @@ public boolean isFreeVar(String identifier) {
         return this.freeVars.contains(identifier);
     }
 
-    private FrameSlot[] getFrameSlots(Collection<String> identifiers, ScopeInfo scope) {
+    private static FrameSlot[] getFrameSlots(Collection<String> identifiers, ScopeInfo scope) {
         assert scope != null : "getting frame slots: scope cannot be null!";
         FrameSlot[] slots = new FrameSlot[identifiers.size()];
         int i = 0;

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/parser/ScopeTranslator.java

@@ -47,7 +47,7 @@ public class ScopeTranslator<T> extends Python3BaseVisitor<T> {
     private final PythonCore core;
     private final boolean interactive;
     private final boolean trackCells;
-    private int comprehensionOrTestDepth = 0;
+    private ScopeInfo currentGeneratorScope = null;
 
     public ScopeTranslator(PythonCore core, TranslationEnvironment environment, boolean interactive, boolean trackCells) {
         this.core = core;
@@ -360,33 +360,37 @@ public T visitComp_for(Python3Parser.Comp_forContext ctx) {
 
     @Override
     public T visitOr_test(Python3Parser.Or_testContext ctx) {
-        ScopeInfo generatorScope = null;
+        boolean pushedCurrentGeneratorScope = false;
         if (ctx.getParent() instanceof Python3Parser.Comp_forContext) {
-            if (comprehensionOrTestDepth == 0 && environment.getCurrentScopeLoopCount() == 1) {
+            if (currentGeneratorScope == null && environment.getCurrentScopeLoopCount() == 1) {
                 // the generator iterator needs to be early evaluated in the parent scope
-                generatorScope = environment.pushCurentScope();
+                currentGeneratorScope = environment.pushCurentScope();
+                pushedCurrentGeneratorScope = true;
             }
-            comprehensionOrTestDepth++;
         }
         try {
             return super.visitOr_test(ctx);
         } finally {
             if (ctx.getParent() instanceof Python3Parser.Comp_forContext) {
-                comprehensionOrTestDepth--;
-                if (comprehensionOrTestDepth == 0 && generatorScope != null && generatorScope.getLoopCount() == 1) {
+                if (pushedCurrentGeneratorScope && currentGeneratorScope.getLoopCount() == 1) {
                     // restore the current scope
                     environment.popCurrentScope();
+                    currentGeneratorScope = null;
                 }
             }
         }
     }
 
     private T visitGenerator(ParserRuleContext ctx, Function<ParserRuleContext, T> block) {
-        environment.beginScope(ctx, ScopeKind.Generator);
+        if (currentGeneratorScope == null) {
+            environment.beginScope(ctx, ScopeKind.Generator);
+        }
         try {
             return block.apply(ctx);
         } finally {
-            environment.endScope(ctx);
+            if (currentGeneratorScope == null) {
+                environment.endScope(ctx);
+            }
         }
     }