Fix: allocate yieldFromGenerator local eagerly

Author: steve-s

graalpython/com.oracle.graal.python.pegparser.test/testData/parser/goldenFiles/YieldStatementTests/customIter01.scope

@@ -27,7 +27,7 @@ ScopeEnvironment
                         StopIteration: [Use, GlobalImplicit]
                         self: [DefParam, Local]
       Scope gen Function
-            Flags: [IsNested, HasFreeVars, IsGenerator]
+            Flags: [IsNested, HasFreeVars, IsGenerator, IsGeneratorWithYieldFrom]
             Symbols:
                   MyIter: [Use, Free]
-                  ret: [DefLocal, DefNonLocal, Free]
+                  ret: [DefLocal, DefNonLocal, Free]

graalpython/com.oracle.graal.python.pegparser.test/testData/parser/goldenFiles/YieldStatementTests/yield12.scope

@@ -3,4 +3,4 @@ ScopeEnvironment
     Symbols:
       f: [DefLocal, Local]
     Scope f Function
-        Flags: [IsGenerator]
+        Flags: [IsGenerator, IsGeneratorWithYieldFrom]

graalpython/com.oracle.graal.python.pegparser.test/testData/parser/goldenFiles/YieldStatementTests/yield13.scope

@@ -3,6 +3,6 @@ ScopeEnvironment
     Symbols:
       f: [DefLocal, Local]
     Scope f Function
-        Flags: [IsGenerator]
+        Flags: [IsGenerator, IsGeneratorWithYieldFrom]
         Symbols:
-            f: [Use, GlobalImplicit]
+            f: [Use, GlobalImplicit]

graalpython/com.oracle.graal.python.pegparser/src/com/oracle/graal/python/pegparser/scope/Scope.java

@@ -119,7 +119,8 @@ enum ScopeFlags {
         NeedsClassClosure,
         NeedsClassDict,
         IsVisitingIterTarget,
-        CanSeeClassScope
+        CanSeeClassScope,
+        IsGeneratorWithYieldFrom,
     }
 
     EnumSet<ScopeFlags> flags = EnumSet.noneOf(ScopeFlags.class);
@@ -243,6 +244,10 @@ public boolean isGenerator() {
         return flags.contains(ScopeFlags.IsGenerator);
     }
 
+    public boolean isGeneratorWithYieldFrom() {
+        return flags.contains(ScopeFlags.IsGeneratorWithYieldFrom);
+    }
+
     public boolean isCoroutine() {
         return flags.contains(ScopeFlags.IsCoroutine);
     }

graalpython/com.oracle.graal.python.pegparser/src/com/oracle/graal/python/pegparser/scope/ScopeEnvironment.java

@@ -920,6 +920,7 @@ public Void visit(ExprTy.YieldFrom node) {
                 node.value.accept(this);
             }
             currentScope.flags.add(ScopeFlags.IsGenerator);
+            currentScope.flags.add(ScopeFlags.IsGeneratorWithYieldFrom);
             if (currentScope.flags.contains(ScopeFlags.IsComprehension)) {
                 throw raiseIfComprehensionBlock(node);
             }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/compiler/bytecode_dsl/RootNodeCompiler.java

@@ -238,7 +238,6 @@ public final class RootNodeCompiler implements BaseBytecodeDSLVisitor<BytecodeDS
 
     // Mutable (must be reset)
     private SourceRange currentLocation;
-    boolean savedYieldFromGenerator;
     BytecodeLocal yieldFromGenerator;
     boolean inExceptStar;
 
@@ -452,7 +451,7 @@ flags, orderedTruffleStringArray(names),
                         null,
                         nodes);
         rootNode.setMetadata(codeUnit, ctx.errorCallback);
-        if (codeUnit.isGeneratorOrCoroutine() && savedYieldFromGenerator) {
+        if (codeUnit.isCoroutine() || codeUnit.isAsyncGenerator() || scope.isGeneratorWithYieldFrom()) {
             rootNode.yieldFromGeneratorIndex = yieldFromGenerator.getLocalIndex();
         }
 
@@ -595,7 +594,6 @@ public void reset() {
         currentLocation = null;
         currentSaveExceptionLocal = null;
         prevSaveExceptionLocal = null;
-        this.savedYieldFromGenerator = false;
         this.inExceptStar = false;
     }
 
@@ -1608,6 +1606,10 @@ public void setUpFrame(ArgumentsTy args, Builder b) {
         if (scope.isCoroutine() || scope.isGenerator()) {
             generatorExceptionStateLocal = b.createLocal();
         }
+
+        if (scope.isGeneratorWithYieldFrom() || scope.isCoroutine()) {
+            yieldFromGenerator = b.createLocal();
+        }
     }
 
     private void copyArguments(ArgumentsTy args, Builder b) {
@@ -2728,10 +2730,7 @@ public void emitYieldFrom(Runnable generatorOrCoroutineProducer, BytecodeLocal r
             generatorOrCoroutineProducer.run();
             b.endStoreLocal();
 
-            if (!savedYieldFromGenerator) {
-                savedYieldFromGenerator = true;
-                yieldFromGenerator = b.createLocal();
-            }
+            assert yieldFromGenerator != null;
             b.beginStoreLocal(yieldFromGenerator);
             b.emitLoadLocal(generator);
             b.endStoreLocal();